"""Pool of live Agent SDK sessions across every Claude agent (§3.2).

One :class:`Session` is one ``ClaudeSDKClient`` (one claude subprocess).
The pool owns the two decisions the adapters used to make on their own:
when a session is closed for idleness (TTL by conversation kind) and
which one goes when memory runs out (measured RSS of the subprocess tree
against the cgroup limit, ``max_live`` where there is no limit). Eviction
only ever picks ``idle && !running_turn && !pending_question`` sessions
that are neither pinned (the master) nor ``dirty`` (mirror gap not yet
repaired); forks and jobs go first.
"""

from __future__ import annotations

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

import psutil

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Iterator, Mapping

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

_log = logging.getLogger("beaver_gateway.core.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


class SessionClient(Protocol):
    async def connect(self) -> None: ...
    async def query(self, prompt: str) -> None: ...
    def receive_response(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 (``core/policy``); dies with the process."""

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

    @property
    def evictable(self) -> bool:
        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

    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

    @staticmethod
    def rss_of(session: Session) -> int | None:
        pid = session.pid
        if pid is None:
            return None
        try:
            process = psutil.Process(pid)
            return process.memory_info().rss + sum(
                c.memory_info().rss for c in process.children(recursive=True)
            )
        except psutil.Error:
            return None

    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]:
        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)
        try:
            await session.client.disconnect()
        except Exception:  # noqa: BLE001
            _log.exception("disconnect failed for session %s", session.session_id)

    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]]:
        now = time.monotonic()
        return [
            {
                "key": s.key,
                "agent": s.agent,
                "kind": s.kind,
                "session_id": s.session_id,
                "pid": s.pid,
                "rss": self.rss_of(s),
                "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_question": s.pending_question,
                "pinned": s.pinned,
                "dirty": s.dirty,
            }
            for s in self._sessions.values()
        ]


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