"""Typed root-session controls backed by the Prime Agent host bridge."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal

from rlm import host_request

AgentSessionDelivery = Literal["steer", "follow_up"]
AgentSessionCapability = Literal["send", "recent_turns", "stop", "restart"]


def _field(payload: dict[str, Any], snake: str, camel: str) -> Any:
    value = payload.get(snake)
    return payload.get(camel) if value is None else value


def _required_string(payload: dict[str, Any], snake: str, camel: str, label: str) -> str:
    value = _field(payload, snake, camel)
    if not isinstance(value, str) or not value.strip():
        raise RuntimeError(f"{label} returned an invalid {snake}")
    return value


def _optional_string(payload: dict[str, Any], snake: str, camel: str) -> str | None:
    value = _field(payload, snake, camel)
    if value is None:
        return None
    if not isinstance(value, str):
        raise RuntimeError(f"session entry has an invalid {snake}")
    return value


@dataclass(frozen=True)
class TurnSummary:
    user: str
    assistant: str
    tool_calls: int
    started_at: int | None = None
    ended_at: int | None = None
    message_count: int = 0
    truncated: bool = False


@dataclass(frozen=True)
class AgentSessionRecentTurns:
    turns: list[TurnSummary]
    count: int
    max_chars: int
    truncated: bool


@dataclass(frozen=True)
class AgentSessionSendReceipt:
    operation_id: str
    status: Literal["accepted", "queued"]
    delivery: AgentSessionDelivery
    target_active_session_id: str
    target_session_id: str
    target_session_name: str | None = None


@dataclass(frozen=True)
class AgentSessionOperationReceipt:
    operation_id: str
    status: Literal["accepted"]


@dataclass
class AgentSessionHandle:
    active_session_id: str
    session_id: str
    session_name: str | None = None
    cwd: str | None = None
    runtime_kind: str | None = "top-level"
    status: str | None = None
    is_current: bool = False
    is_streaming: bool = False
    is_compacting: bool = False
    attached_clients: int = 0
    message_count: int = 0
    queued_count: int = 0
    is_session_active: bool = False
    capabilities: tuple[AgentSessionCapability, ...] = field(default=())

    @property
    def id(self) -> str:
        return self.active_session_id

    @property
    def name(self) -> str:
        return self.session_name or self.active_session_id

    @property
    def running(self) -> bool:
        return self.is_session_active or self.is_streaming or self.status in {"running", "busy", "model", "tool"}

    async def refresh(self) -> AgentSessionHandle:
        refreshed = await sessions.get(self.active_session_id)
        self._merge(refreshed)
        return self

    async def info(self) -> dict[str, Any]:
        payload = await _request("get", {"target": self.active_session_id})
        refreshed = _session_from_payload(_required_mapping(payload, "session", "agent.sessions.get"))
        self._merge(refreshed)
        return _session_to_dict(refreshed)

    async def observe(self, count: int = 8, max_chars: int = 800) -> list[TurnSummary]:
        return (await self.recent_turns(count=count, max_chars=max_chars)).turns

    async def recent_turns(self, count: int = 3, max_chars: int = 6_000) -> AgentSessionRecentTurns:
        self._require_capability("recent_turns")
        if not isinstance(count, int) or isinstance(count, bool):
            raise TypeError("count must be an int")
        if not isinstance(max_chars, int) or isinstance(max_chars, bool):
            raise TypeError("max_chars must be an int")
        payload = await _request(
            "recent_turns",
            {"target": self.active_session_id, "count": count, "max_chars": max_chars},
        )
        agent_payload = payload.get("agent")
        if isinstance(agent_payload, dict):
            self._merge(_session_from_payload(agent_payload))
        raw_turns = payload.get("turns")
        if not isinstance(raw_turns, list):
            raise RuntimeError("agent.sessions.recent_turns returned an invalid turns list")
        result_count = payload.get("count", count)
        result_max_chars = payload.get("maxChars", payload.get("max_chars", max_chars))
        truncated = payload.get("truncated", False)
        if (
            isinstance(result_count, bool)
            or not isinstance(result_count, int)
            or isinstance(result_max_chars, bool)
            or not isinstance(result_max_chars, int)
            or not isinstance(truncated, bool)
        ):
            raise RuntimeError("agent.sessions.recent_turns returned invalid metadata")
        return AgentSessionRecentTurns(
            turns=[_turn_from_payload(turn) for turn in raw_turns],
            count=result_count,
            max_chars=result_max_chars,
            truncated=truncated,
        )

    async def send(self, message: str, delivery: AgentSessionDelivery = "follow_up") -> AgentSessionSendReceipt:
        self._require_capability("send")
        if not isinstance(message, str) or not message.strip():
            raise TypeError("message must be a non-empty str")
        if delivery not in {"steer", "follow_up"}:
            raise ValueError("delivery must be 'steer' or 'follow_up'")
        payload = await _request(
            "send",
            {"target": self.active_session_id, "message": message, "delivery": delivery},
        )
        return _send_receipt_from_payload(payload)

    async def stop(self) -> AgentSessionOperationReceipt:
        self._require_capability("stop")
        receipt = _operation_from_payload(await _request("stop", {"target": self.active_session_id}))
        self.status = "stopped"
        self.is_session_active = False
        return receipt

    async def restart(self) -> AgentSessionOperationReceipt:
        self._require_capability("restart")
        return _operation_from_payload(await _request("restart", {"target": self.active_session_id}))

    def _require_capability(self, capability: AgentSessionCapability) -> None:
        if capability not in self.capabilities:
            raise RuntimeError(f"agent session {self.name!r} does not support {capability}")

    def _merge(self, other: AgentSessionHandle) -> None:
        self.active_session_id = other.active_session_id
        self.session_id = other.session_id
        self.session_name = other.session_name
        self.cwd = other.cwd
        self.runtime_kind = other.runtime_kind
        self.status = other.status
        self.is_current = other.is_current
        self.is_streaming = other.is_streaming
        self.is_compacting = other.is_compacting
        self.attached_clients = other.attached_clients
        self.message_count = other.message_count
        self.queued_count = other.queued_count
        self.is_session_active = other.is_session_active
        self.capabilities = other.capabilities


class SessionNamespace:
    async def list(self) -> list[AgentSessionHandle]:
        payload = await _request("list")
        entries = payload.get("sessions")
        if not isinstance(entries, list):
            raise RuntimeError("agent.sessions.list returned an invalid sessions list")
        return [_session_from_payload(entry) for entry in entries]

    async def get(
        self,
        target: str | None = None,
        *,
        active_session_id: str | None = None,
        session_id: str | None = None,
        name: str | None = None,
    ) -> AgentSessionHandle:
        selectors = [value for value in (target, active_session_id, session_id, name) if value is not None]
        if len(selectors) != 1 or not isinstance(selectors[0], str) or not selectors[0].strip():
            raise TypeError("sessions.get requires exactly one non-empty target selector")
        payload = await _request("get", {"target": selectors[0].strip()})
        return _session_from_payload(_required_mapping(payload, "session", "agent.sessions.get"))


sessions = SessionNamespace()


async def _request(operation: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
    result = await host_request(f"agent.sessions.{operation}", payload)
    if not isinstance(result, dict):
        raise RuntimeError(f"agent.sessions.{operation} returned an invalid result")
    return result


def _session_from_payload(payload: Any) -> AgentSessionHandle:
    if not isinstance(payload, dict):
        raise RuntimeError("agent.sessions returned an invalid session entry")
    active_session_id = _required_string(payload, "active_session_id", "activeSessionId", "agent.sessions")
    session_id = _required_string(payload, "session_id", "sessionId", "agent.sessions")
    session_name = _optional_string(payload, "session_name", "sessionName")
    cwd = _optional_string(payload, "cwd", "cwd")
    runtime_kind = _optional_string(payload, "runtime_kind", "runtimeKind")
    status = _optional_string(payload, "status", "status")
    capabilities_raw = payload.get("capabilities", [])
    if not isinstance(capabilities_raw, list) or not all(isinstance(value, str) for value in capabilities_raw):
        raise RuntimeError("agent.sessions entry has invalid capabilities")
    capabilities = tuple(value for value in capabilities_raw if value in {"send", "recent_turns", "stop", "restart"})
    return AgentSessionHandle(
        active_session_id=active_session_id,
        session_id=session_id,
        session_name=session_name,
        cwd=cwd,
        runtime_kind=runtime_kind,
        status=status,
        is_current=_bool_field(payload, "is_current", "isCurrent"),
        is_streaming=_bool_field(payload, "is_streaming", "isStreaming"),
        is_compacting=_bool_field(payload, "is_compacting", "isCompacting"),
        attached_clients=_int_field(payload, "attached_clients", "attachedClients"),
        message_count=_int_field(payload, "message_count", "messageCount"),
        queued_count=_int_field(payload, "queued_count", "queuedCount"),
        is_session_active=_bool_field(payload, "is_session_active", "isSessionActive"),
        capabilities=capabilities,
    )


def _session_to_dict(session: AgentSessionHandle) -> dict[str, Any]:
    return {
        "active_session_id": session.active_session_id,
        "session_id": session.session_id,
        "session_name": session.session_name,
        "cwd": session.cwd,
        "runtime_kind": session.runtime_kind,
        "status": session.status,
        "is_current": session.is_current,
        "is_streaming": session.is_streaming,
        "is_compacting": session.is_compacting,
        "attached_clients": session.attached_clients,
        "message_count": session.message_count,
        "queued_count": session.queued_count,
        "is_session_active": session.is_session_active,
        "capabilities": list(session.capabilities),
    }


def _turn_from_payload(payload: Any) -> TurnSummary:
    if not isinstance(payload, dict):
        raise RuntimeError("agent.sessions.recent_turns returned an invalid turn")
    user = payload.get("user")
    assistant = payload.get("assistant")
    tool_calls = payload.get("toolCalls", payload.get("tool_calls"))
    started_at = payload.get("startedAt", payload.get("started_at"))
    ended_at = payload.get("endedAt", payload.get("ended_at"))
    message_count = payload.get("messageCount", payload.get("message_count", 0))
    truncated = payload.get("truncated", False)
    if (
        not isinstance(user, str)
        or not isinstance(assistant, str)
        or isinstance(tool_calls, bool)
        or not isinstance(tool_calls, int)
        or tool_calls < 0
        or (started_at is not None and (isinstance(started_at, bool) or not isinstance(started_at, int)))
        or (ended_at is not None and (isinstance(ended_at, bool) or not isinstance(ended_at, int)))
        or isinstance(message_count, bool)
        or not isinstance(message_count, int)
        or message_count < 0
        or not isinstance(truncated, bool)
    ):
        raise RuntimeError("agent.sessions.recent_turns returned an invalid turn")
    return TurnSummary(user, assistant, tool_calls, started_at, ended_at, message_count, truncated)


def _send_receipt_from_payload(payload: dict[str, Any]) -> AgentSessionSendReceipt:
    operation_id = _required_string(payload, "operation_id", "operationId", "agent.sessions.send")
    status = payload.get("status")
    delivery = payload.get("delivery")
    target = payload.get("target")
    if status not in {"accepted", "queued"} or delivery not in {"steer", "follow_up"} or not isinstance(target, dict):
        raise RuntimeError("agent.sessions.send returned an invalid receipt")
    target_active = _required_string(target, "active_session_id", "activeSessionId", "agent.sessions.send")
    target_session = _required_string(target, "session_id", "sessionId", "agent.sessions.send")
    target_name = _optional_string(target, "session_name", "sessionName")
    return AgentSessionSendReceipt(operation_id, status, delivery, target_active, target_session, target_name)


def _operation_from_payload(payload: dict[str, Any]) -> AgentSessionOperationReceipt:
    operation_id = _required_string(payload, "operation_id", "operationId", "agent.sessions.operation")
    if payload.get("status") != "accepted":
        raise RuntimeError("agent.sessions operation returned an invalid status")
    return AgentSessionOperationReceipt(operation_id, "accepted")


def _required_mapping(payload: dict[str, Any], key: str, operation: str) -> dict[str, Any]:
    value = payload.get(key)
    if not isinstance(value, dict):
        raise RuntimeError(f"{operation} returned an invalid {key}")
    return value


def _bool_field(payload: dict[str, Any], snake: str, camel: str) -> bool:
    value = _field(payload, snake, camel)
    return value if isinstance(value, bool) else False


def _int_field(payload: dict[str, Any], snake: str, camel: str) -> int:
    value = _field(payload, snake, camel)
    return value if isinstance(value, int) and not isinstance(value, bool) else 0