"""ACP session/* method handlers.

session/prompt dispatches to the AI Agent bridge (MobileAgent v3.5).
Session management (new/load/list) maintains in-memory session state.
"""

from __future__ import annotations

import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Awaitable

logger = logging.getLogger(__name__)

StreamCallback = Callable[[dict[str, Any]], Awaitable[None]]

_MAX_SESSIONS = 50
_SESSION_TTL = 3600.0  # 1 hour


@dataclass
class Session:
    id: str
    history: list[dict[str, Any]] = field(default_factory=list)
    created_at: float = field(default_factory=time.time)
    last_active: float = field(default_factory=time.time)


class SessionHandler:
    """Handles ACP ``session/*`` JSON-RPC methods."""

    def __init__(
        self,
        on_prompt: StreamCallback | None = None,
        max_sessions: int = _MAX_SESSIONS,
        session_ttl: float = _SESSION_TTL,
    ) -> None:
        self._on_prompt = on_prompt
        self._sessions: dict[str, Session] = {}
        self._active_session_id: str | None = None
        self._max_sessions = max_sessions
        self._session_ttl = session_ttl

    @property
    def active_session(self) -> Session | None:
        if self._active_session_id:
            return self._sessions.get(self._active_session_id)
        return None

    async def handle_new(self, msg: dict[str, Any]) -> dict[str, Any]:
        self._evict_expired()
        session_id = str(uuid.uuid4())
        self._sessions[session_id] = Session(id=session_id)
        self._active_session_id = session_id
        return {"session_id": session_id}

    async def handle_load(self, msg: dict[str, Any]) -> dict[str, Any]:
        params = msg.get("params", {})
        session_id = params.get("session_id", "")
        if session_id not in self._sessions:
            raise ValueError(f"Session not found: {session_id}")
        self._sessions[session_id].last_active = time.time()
        self._active_session_id = session_id
        return {"session_id": session_id}

    async def handle_list(self, msg: dict[str, Any]) -> dict[str, Any]:
        self._evict_expired()
        return {
            "sessions": [
                {"id": s.id, "message_count": len(s.history)}
                for s in self._sessions.values()
            ]
        }

    async def handle_prompt(self, msg: dict[str, Any]) -> dict[str, Any]:
        params = msg.get("params", {})
        prompt = params.get("prompt", "")
        if not prompt:
            raise ValueError("prompt is required")

        if not self._active_session_id or self._active_session_id not in self._sessions:
            await self.handle_new({})

        session = self.active_session
        assert session is not None

        session.last_active = time.time()
        session.history.append({"role": "user", "content": prompt})

        if self._on_prompt:
            await self._on_prompt({"prompt": prompt, "session_id": session.id})

        return {
            "status": "accepted",
            "session_id": session.id,
            "prompt": prompt,
        }

    def _evict_expired(self) -> None:
        """Remove sessions older than TTL or exceeding max count."""
        now = time.time()
        expired = [
            sid for sid, s in self._sessions.items()
            if (now - s.last_active) > self._session_ttl
        ]
        for sid in expired:
            del self._sessions[sid]
            if self._active_session_id == sid:
                self._active_session_id = None

        if len(self._sessions) > self._max_sessions:
            by_age = sorted(self._sessions.items(), key=lambda kv: kv[1].last_active)
            for sid, _ in by_age[: len(self._sessions) - self._max_sessions]:
                del self._sessions[sid]
                if self._active_session_id == sid:
                    self._active_session_id = None
