"""Session service for managing agent sessions."""

import asyncio
import contextlib
import datetime
import json
import logging
import uuid
from collections.abc import AsyncIterator
from pathlib import Path

from infra_mvp.stream.schemas.session import SessionState
from infra_mvp.stream.services.event_converter import (
    output_event_to_json,
    user_interaction_to_input_event,
)
from yera.agents.decorator import AgentFunctionWrapper
from yera.events.runtime import PyRuntimeExecutor
from yera.events.stream import OutputEvent


class SessionService:
    """Service for managing agent sessions and generating SSE events."""

    def __init__(self, log_events: bool = False):
        """Initialise the Python session service.

        Args:
            log_events: Enable logging of events to file
        """
        self._sessions: dict[str, SessionState] = {}
        self._log_events = log_events
        self._events_log_file: Path | None = None

        if self._log_events:
            # Create timestamped events log file
            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            self._events_log_file = Path(f"yera-event-log-{timestamp}.jsonl")
            logging.info(
                f"Events logging enabled: events will be logged to {self._events_log_file}"
            )

    def generate_session_id(self) -> str:
        """Generate a unique session ID.

        Returns:
            UUID string for the session
        """
        return str(uuid.uuid4())

    def create_session(self, agent: AgentFunctionWrapper, args: tuple = ()) -> str:
        """Create a new session for the agent.

        Args:
            agent: The agent wrapper to execute.
            args: Positional arguments to pass to the agent's ``invoke`` method.

        Returns:
            Generated session_id
        """
        session_id = self.generate_session_id()
        executor = PyRuntimeExecutor(agent, args)

        session_state = SessionState(
            agent=agent,
            executor=executor,
            session_id=session_id,
        )

        self._sessions[session_id] = session_state
        return session_id

    def get_session(self, session_id: str) -> SessionState:
        """Retrieve an existing session.

        Args:
            session_id: Session identifier

        Returns:
            Python agent session state

        Raises:
            KeyError: If session not found
        """
        if session_id not in self._sessions:
            raise KeyError(f"Session '{session_id}' not found")
        return self._sessions[session_id]

    def add_interaction(self, session_id: str, interaction: dict) -> None:
        """Add an interaction to a session.

        Converts the UserInteraction to an InputEvent and pushes it to the
        executor's EventStream input queue. Also stores the interaction in
        session state for debugging.

        Args:
            session_id: Session identifier
            interaction: Interaction dictionary (UserInteraction format)

        Raises:
            KeyError: If session not found
            RuntimeError: If executor not started
        """
        session = self.get_session(session_id)
        executor = session.executor

        # Store interaction in session state for debugging
        session.interactions.append(interaction)

        # Push to executor's EventStream if executor is running
        if not executor.is_running():
            raise RuntimeError("Cannot add interaction: executor not started")

        # Form submission: expand to one InputEvent per sub-block (same order as agent's await_input calls)
        if interaction.get("blockType") == "form" and isinstance(
            interaction.get("value"), list
        ):
            for item in interaction["value"]:
                sub_interaction = {
                    "blockId": item["sub_block_id"],
                    "blockType": item["sub_block_type"],
                    "value": item["value"],
                }
                input_event = user_interaction_to_input_event(sub_interaction)
                executor.stream.push_input(input_event)
            return

        # Single interaction: one InputEvent
        input_event = user_interaction_to_input_event(interaction)
        executor.stream.push_input(input_event)

    async def generate_events(self, session_id: str) -> AsyncIterator[str]:
        """Generate SSE events from Python agent execution."""
        session = self.get_session(session_id)
        executor = session.executor
        executor.start()

        # Layouts that have had at least one await_user child; stream pauses on layout_end for these
        layouts_with_await_user: set[str] = set()

        while True:
            if not executor.is_running():
                async for sse in self._drain_remaining_events(session_id, executor):
                    yield sse
                return

            event = await asyncio.to_thread(executor.stream.pop_output, timeout=None)
            if event is None:
                continue

            yield self._format_sse_event(session_id, event)

            if event.block_type == "exit":
                self.cleanup_session(session_id)
                return

            if event.message_type == "await_user":
                if event.parent_block_id is not None:
                    layouts_with_await_user.add(event.parent_block_id)
                    # Do not return; keep streaming until layout_end for this layout
                else:
                    return

            if (
                event.block_type == "layout_end"
                and event.block_id in layouts_with_await_user
            ):
                return

    def _format_sse_event(self, session_id: str, event: OutputEvent) -> str:
        """Convert OutputEvent to SSE format and optionally log."""
        event_json = output_event_to_json(event)
        if self._log_events:
            self._log_event_to_file(session_id, event_json)
        return f"data: {json.dumps(event_json)}\n\n"

    async def _drain_remaining_events(
        self, session_id: str, executor: PyRuntimeExecutor
    ) -> AsyncIterator[str]:
        """Drain remaining events from a stopped executor."""
        if executor.stream is None:
            return

        while executor.stream.has_output():
            event = await asyncio.to_thread(executor.stream.pop_output, timeout=0.1)
            if event is None:
                break
            yield self._format_sse_event(session_id, event)

    def cleanup_session(self, session_id: str) -> None:
        """Clean up session-specific resources.

        Stops the executor process if running and removes the session from
        internal storage.

        Args:
            session_id: Session identifier to clean up

        Raises:
            KeyError: If session not found
        """
        if session_id not in self._sessions:
            raise KeyError(f"Session '{session_id}' not found")

        session = self._sessions[session_id]

        # Stop executor if it exists (terminates process)
        if session.executor is not None:
            session.executor.stop()

        # Remove session from storage
        del self._sessions[session_id]

    def cleanup_all(self) -> None:
        """Clean up all sessions.

        Iterates through all sessions and calls cleanup_session() for each.
        Called on server shutdown to ensure no orphaned processes remain.
        """
        # Create a list of session IDs to avoid modifying dict during iteration
        session_ids = list(self._sessions.keys())

        for session_id in session_ids:
            with contextlib.suppress(KeyError):
                self.cleanup_session(session_id)

    def _log_event_to_file(self, session_id: str, event_json: dict) -> None:
        """Log an event to the debug log file in JSONL format.

        Args:
            session_id: Session identifier
            event_json: Event data as JSON-serializable dict
        """
        if self._events_log_file is None:
            return

        log_entry = {
            "timestamp": datetime.datetime.now().isoformat(),
            "session_id": session_id,
            "event": event_json,
        }

        try:
            with open(self._events_log_file, "a", encoding="utf-8") as f:
                f.write(json.dumps(log_entry) + "\n")
        except Exception as e:
            logging.warning(f"Failed to write debug log: {e}")
