"""Pool of live Agent SDK sessions across every Claude agent.

One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess). The
pool decides when a session is closed for idleness (TTL by conversation
kind) and which one goes when memory runs out (RSS of the subprocess tree
against the cgroup limit, or ``max_live`` with no limit).
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol

import psutil

from beaver_gateway.backends.processes import Processes, subtree

if TYPE_CHECKING:
    from collections.abc import (
        AsyncIterable,
        AsyncIterator,
        Callable,
        Iterator,
        Mapping,
    )

__all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"]

_log = logging.getLogger("beaver_gateway.backends.sessions")

DEFAULT_TTL: Mapping[str, float | None] = {
    "master": None,
    "branch": 7200.0,
    "deep": 1800.0,
    "job": 0.0,
    "fork": 0.0,
}
"""Idle seconds before a session is closed; ``None`` = never (pinned kinds)."""

_EVICT_ORDER = {"fork": 0, "job": 0, "deep": 1, "branch": 2, "master": 3}
_RSS_HEADROOM = 0.8
STDERR_LINES = 200


class SessionClient(Protocol):
    async def connect(self) -> None: ...
    async def query(self, prompt: str | AsyncIterable[dict[str, Any]]) -> None: ...
    def receive_messages(self) -> AsyncIterator[Any]: ...
    async def interrupt(self) -> None: ...
    async def disconnect(self) -> None: ...


@dataclass
class Session:
    key: str
    agent: str
    kind: str
    client: SessionClient
    session_id: str | None
    resumed: bool
    pinned: bool = False
    dirty: bool = False
    running_turn: str | None = None
    pending_question: bool = False
    interrupt_requested: bool = False
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    last_used: float = field(default_factory=time.monotonic)
    created_at: float = field(default_factory=time.monotonic)
    turns: int = 0
    state: dict[str, Any] = field(default_factory=dict)
    """Scratch for policy rules (``agents/policy.py``); dies with the process."""

    reader: asyncio.Task[None] | None = None
    """Reads the client for the session's whole life, not only during a turn."""
    inbox: asyncio.Queue[Any] = field(default_factory=asyncio.Queue)
    """Where the reader puts messages while a turn is attached."""
    pending: list[Any] = field(default_factory=list)
    """A turn the CLI started on its own (a background task reported back),
    buffered until a consumer attaches."""
    attached: bool = False
    woken: bool = False
    note: str | None = None
    """Summary of the last task notification, the pending turn's prompt."""
    observer: Callable[[Any], None] | None = None
    """The last turn's observer; sees subagent traffic between turns."""
    stderr: deque[str] = field(default_factory=lambda: deque(maxlen=STDERR_LINES))

    @property
    def busy(self) -> bool:
        return self.lock.locked() or self.running_turn is not None

    @property
    def evictable(self) -> bool:
        """Never evict a pinned (master) session or a ``dirty`` one.

        ``dirty`` means a mirror gap not yet repaired.
        """
        return not (self.pinned or self.dirty or self.busy or self.pending_question)

    @property
    def pid(self) -> int | None:
        transport = getattr(self.client, "_transport", None)
        process = getattr(transport, "_process", None)
        pid = getattr(process, "pid", None)
        return pid if isinstance(pid, int) else None


class SessionPool:
    def __init__(
        self,
        *,
        rss_limit: int | None = None,
        max_live: int = 8,
        ttl: Mapping[str, float | None] = DEFAULT_TTL,
        reap_interval: float = 60.0,
    ) -> None:
        self._sessions: dict[str, Session] = {}
        self._rss_limit = rss_limit if rss_limit is not None else cgroup_limit()
        self._max_live = max_live
        self._ttl = dict(ttl)
        self._reap_interval = reap_interval
        self._processes = Processes()

    def __len__(self) -> int:
        return len(self._sessions)

    def __iter__(self) -> Iterator[Session]:
        return iter(list(self._sessions.values()))

    def __contains__(self, key: object) -> bool:
        return key in self._sessions

    @property
    def rss_limit(self) -> int | None:
        return self._rss_limit

    def get(self, key: str) -> Session | None:
        return self._sessions.get(key)

    def add(self, session: Session) -> Session:
        self._sessions[session.key] = session
        return session

    def pop(self, key: str) -> Session | None:
        return self._sessions.pop(key, None)

    def rekey(self, old: str, new: str) -> Session | None:
        session = self._sessions.pop(old, None)
        if session is None:
            return None
        stale = self._sessions.pop(new, None)
        session.key = new
        self._sessions[new] = session
        return stale if stale is not session else None

    def ttl_for(self, kind: str) -> float | None:
        return self._ttl.get(kind, self._ttl.get("deep"))

    def rss(self) -> int:
        try:
            children = psutil.Process(os.getpid()).children(recursive=True)
        except psutil.Error:
            return 0
        total = 0
        for child in children:
            with contextlib.suppress(psutil.Error):
                total += child.memory_info().rss
        return total

    def over_limit(self) -> bool:
        if self._rss_limit is not None:
            return self.rss() > self._rss_limit * _RSS_HEADROOM
        return len(self._sessions) >= self._max_live

    def victims(self) -> list[Session]:
        """Evictable sessions, forks/jobs first, then oldest idle."""
        candidates = [s for s in self._sessions.values() if s.evictable]
        candidates.sort(key=lambda s: (_EVICT_ORDER.get(s.kind, 1), s.last_used))
        return candidates

    async def make_room(self) -> int:
        closed = 0
        while self.over_limit():
            victims = self.victims()
            if not victims:
                _log.warning(
                    "session pool over limit (%d live, rss=%d) but nothing evictable",
                    len(self._sessions),
                    self.rss(),
                )
                break
            await self.close(victims[0].key)
            closed += 1
        return closed

    async def close(self, key: str) -> None:
        session = self._sessions.pop(key, None)
        if session is None:
            return
        _log.info("closing session %s (%s, %s)", session.session_id, session.kind, key)
        if session.reader is not None and session.reader is not asyncio.current_task():
            session.reader.cancel()
        if session.attached:
            session.inbox.put_nowait(RuntimeError(f"session {key} was closed"))
        try:
            await session.client.disconnect()
        except Exception:  # noqa: BLE001
            _log.exception("disconnect failed for session %s", session.session_id)

    async def interrupt(self, key: str) -> bool:
        session = self._sessions.get(key)
        if session is None or not session.busy:
            return False
        session.interrupt_requested = True
        await session.client.interrupt()
        return True

    async def close_all(self, *, agent: str | None = None) -> None:
        for session in list(self._sessions.values()):
            if agent is None or session.agent == agent:
                await self.close(session.key)

    async def reap_once(self) -> int:
        now = time.monotonic()
        closed = 0
        for session in list(self._sessions.values()):
            ttl = self.ttl_for(session.kind)
            if ttl is None or not session.evictable:
                continue
            if now - session.last_used > ttl:
                await self.close(session.key)
                closed += 1
        return closed

    async def reap_loop(self) -> None:
        while True:
            await asyncio.sleep(self._reap_interval)
            try:
                await self.reap_once()
                await self.make_room()
            except Exception:  # noqa: BLE001
                _log.exception("session reaper failed")

    def snapshot(self) -> list[dict[str, Any]]:
        return self.census()["sessions"]

    def census(self) -> dict[str, Any]:
        now = time.monotonic()
        rows = self._processes.tree_parents_first()
        owned: set[int] = set()
        sessions: list[dict[str, Any]] = []
        for s in self._sessions.values():
            tree = subtree(rows, s.pid)
            owned.update(row["pid"] for row in tree)
            sessions.append(
                {
                    "key": s.key,
                    "agent": s.agent,
                    "kind": s.kind,
                    "session_id": s.session_id,
                    "pid": s.pid,
                    "rss": sum(row["rss"] for row in tree) if tree else None,
                    "cpu": _cpu(tree),
                    "idle_seconds": round(now - s.last_used, 1),
                    "age_seconds": round(now - s.created_at, 1),
                    "turns": s.turns,
                    "busy": s.busy,
                    "running_turn": s.running_turn,
                    "pending": len(s.pending),
                    "pending_question": s.pending_question,
                    "pinned": s.pinned,
                    "dirty": s.dirty,
                    "processes": tree,
                }
            )
        return {
            "rss": sum(row["rss"] for row in rows),
            "rss_limit": self._rss_limit,
            "cpu": _cpu(rows),
            "sessions": sessions,
            "others": [row for row in rows if row["pid"] not in owned],
        }

    async def kill(self, pid: int) -> bool:
        if not self._processes.kill_own_descendant(pid):
            return False
        for session in list(self._sessions.values()):
            if session.pid == pid:
                await self.close(session.key)
        return True


def _cpu(rows: list[dict[str, Any]]) -> float | None:
    rates = [row["cpu"] for row in rows if row["cpu"] is not None]
    return round(sum(rates), 1) if rates else None


def cgroup_limit() -> int | None:
    for path in (
        "/sys/fs/cgroup/memory.max",
        "/sys/fs/cgroup/memory/memory.limit_in_bytes",
    ):
        try:
            raw = Path(path).read_text(encoding="ascii").strip()
        except OSError:
            continue
        if raw.isdigit() and int(raw) < 1 << 60:
            return int(raw)
    return None