"""
Enhanced Orchestrator Session Management

Session lifecycle methods extracted from EnhancedTaskOrchestrator.
"""
import structlog
from typing import Optional, Dict, Any

from sessions.session_manager import SessionManager, AgentSession, SessionState
from sessions.heartbeat import HeartbeatMonitor, HeartbeatClient
from brain.context_manager import ContextManager

logger = structlog.get_logger(__name__)


async def create_session(
    session_manager: SessionManager,
    context_manager: ContextManager,
    heartbeat: HeartbeatMonitor,
    voice_sessions: Dict[str, AgentSession],
    heartbeat_clients: Dict[str, HeartbeatClient],
    voice_session_id: str,
    user_id: str = "",
    metadata: Optional[Dict[str, Any]] = None,
    system_prompt: str = "",
) -> AgentSession:
    """Creates a new agent session for a voice session."""
    session = await session_manager.create_session(
        agent_type="voice-orchestrator",
        user_id=user_id,
        context={"voice_session_id": voice_session_id},
        metadata=metadata
    )

    context_manager.create_context(
        session_id=session.session_id,
        system_prompt=system_prompt,
        max_messages=50
    )

    heartbeat_client = HeartbeatClient(
        session_id=session.session_id,
        monitor=heartbeat,
        interval_seconds=10
    )
    await heartbeat_client.start()
    heartbeat.register(session.session_id, "voice-orchestrator")

    voice_sessions[voice_session_id] = session
    heartbeat_clients[session.session_id] = heartbeat_client

    logger.info(
        "Created agent session",
        session_id=session.session_id[:8],
        voice_session_id=voice_session_id
    )
    return session


async def end_session(
    session_manager: SessionManager,
    heartbeat: HeartbeatMonitor,
    voice_sessions: Dict[str, AgentSession],
    heartbeat_clients: Dict[str, HeartbeatClient],
    voice_session_id: str,
) -> None:
    """Ends an agent session."""
    session = voice_sessions.get(voice_session_id)
    if not session:
        return

    if session.session_id in heartbeat_clients:
        await heartbeat_clients[session.session_id].stop()
        del heartbeat_clients[session.session_id]

    heartbeat.unregister(session.session_id)
    session.complete()
    await session_manager.update_session(session)
    del voice_sessions[voice_session_id]

    logger.info(
        "Ended agent session",
        session_id=session.session_id[:8],
        duration_seconds=session.get_duration().total_seconds()
    )


async def recover_session(
    session_manager: SessionManager,
    heartbeat: HeartbeatMonitor,
    voice_sessions: Dict[str, AgentSession],
    heartbeat_clients: Dict[str, HeartbeatClient],
    tasks: Dict[str, Any],
    process_task_fn,
    voice_session_id: str,
    session_id: str,
) -> Optional[AgentSession]:
    """Recovers a session from checkpoint."""
    session = await session_manager.get_session(session_id)

    if not session:
        logger.warning("Session not found for recovery", session_id=session_id)
        return None

    if session.state != SessionState.ACTIVE:
        logger.warning(
            "Session not active for recovery",
            session_id=session_id, state=session.state.value
        )
        return None

    session.resume()

    heartbeat_client = HeartbeatClient(
        session_id=session.session_id,
        monitor=heartbeat,
        interval_seconds=10
    )
    await heartbeat_client.start()
    heartbeat.register(session.session_id, "voice-orchestrator")

    voice_sessions[voice_session_id] = session
    heartbeat_clients[session.session_id] = heartbeat_client

    # Recover pending tasks from checkpoints
    from models.task import TaskState
    for checkpoint in reversed(session.checkpoints):
        if checkpoint.name == "task_queued":
            task_id = checkpoint.data.get("task_id")
            if task_id and task_id in tasks:
                task = tasks[task_id]
                if task.state == TaskState.QUEUED:
                    await process_task_fn(task)
                    logger.info("Recovered pending task", task_id=task_id[:8])

    logger.info(
        "Recovered session",
        session_id=session.session_id[:8],
        checkpoints=len(session.checkpoints)
    )
    return session