"""
Session — 单个 CMDC Agent Session 的操作接口。

支持同步和异步两种访问方式。
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any

import httpx

from cmdc._http import raise_for_status
from cmdc.approval import ApprovalPolicy, ApprovalPolicyLike, make_policy
from cmdc.sse import AsyncSSEIterator, SSEIterator
from cmdc.tools import LocalTool
from cmdc.types import (
    Event,
    Message,
    PromptResponse,
    SessionInfo,
    Stats,
    TokenUsage,
)
from cmdc.ws import AsyncWSConnection, WSConnection

if TYPE_CHECKING:
    pass


class Session:
    """单个 Agent Session 的操作接口（同步）。

    通过 ``CMDC.create_session()`` 或 ``CMDC.session()`` context manager 获取。

    用法::

        with client.session(model="deepseek:deepseek-chat") as session:
            for event in session.prompt("你好"):
                if event.type == "message_delta":
                    print(event.data["delta"], end="")
    """

    def __init__(
        self,
        session_id: str,
        base_url: str,
        http_client: httpx.Client,
        headers: dict[str, str],
    ) -> None:
        self.session_id = session_id
        self._base_url = base_url.rstrip("/")
        self._client = http_client
        self._headers = headers

    # ------------------------------------------------------------------
    # 基础 URL 构建
    # ------------------------------------------------------------------

    def _url(self, path: str) -> str:
        return f"{self._base_url}/v1/sessions/{self.session_id}{path}"

    def _events_url(self) -> str:
        return self._url("/events")

    def _ws_url(self) -> str:
        url = self._url("/ws")
        # 将 http/https 换成 ws/wss
        return url.replace("https://", "wss://").replace("http://", "ws://")

    # ------------------------------------------------------------------
    # SSE 流式迭代（同步）
    # ------------------------------------------------------------------

    def prompt(
        self,
        text: str,
        *,
        approval_policy: ApprovalPolicyLike | None = None,
    ) -> Iterator[Event]:
        """向 Agent 发送 prompt，返回 SSE 事件迭代器（同步）。

        发送 prompt 后立即开始监听 SSE 流，自动处理审批策略。

        Args:
            text: 用户消息文本。
            approval_policy: 审批策略（可选）。支持字符串
                ``"auto_approve"`` / ``"auto_reject"`` / ``"prompt"``，
                ``ApprovalPolicy`` 实例，或 ``(Event) -> bool`` callable。

        Yields:
            :class:`~cmdc.types.Event` — 逐个 Agent 事件。

        用法::

            for event in session.prompt("写一个排序"):
                if event.type == "message_delta":
                    print(event.delta, end="")
        """
        policy: ApprovalPolicy | None = None
        if approval_policy is not None:
            policy = make_policy(approval_policy)

        # 发送 prompt（异步，202 立即返回）
        resp = self._client.post(
            self._url("/prompt"),
            headers=self._headers,
            json={"text": text},
        )
        raise_for_status(resp)

        # 建立 SSE 连接，迭代事件
        with SSEIterator(self._client, self._events_url(), self._headers) as it:
            for event in it:
                # 自动处理 approval_required
                if policy is not None and event.type == "approval_required":
                    approved = policy.decide(event)
                    approval_id = event.approval_id
                    if approval_id:
                        if approved:
                            self.approve(approval_id)
                        else:
                            self.reject(approval_id)
                yield event

    def prompt_sync(self, text: str) -> str:
        """发送 prompt，阻塞等待，返回完整 assistant 回复文本。

        不支持审批策略（如需 HITL，请用 :meth:`prompt`）。

        Returns:
            最终 assistant 消息的完整文本。
        """
        parts: list[str] = []
        for event in self.prompt(text):
            if event.type == "message_delta" and event.delta:
                parts.append(event.delta)
        return "".join(parts)

    # ------------------------------------------------------------------
    # 发送 prompt 不监听流（仅返回 request_id）
    # ------------------------------------------------------------------

    def send_prompt(self, text: str) -> PromptResponse:
        """投递 prompt 并立即返回，不监听 SSE 流。

        适合已有独立 SSE 订阅的场景。
        """
        resp = self._client.post(
            self._url("/prompt"),
            headers=self._headers,
            json={"text": text},
        )
        raise_for_status(resp)
        body = resp.json()
        return PromptResponse(
            request_id=body["requestId"],
            session_id=body["sessionId"],
            queued=body.get("queued", False),
        )

    # ------------------------------------------------------------------
    # 控制 API
    # ------------------------------------------------------------------

    def approve(self, approval_id: str) -> None:
        """审批通过。"""
        resp = self._client.post(
            self._url("/approve"),
            headers=self._headers,
            json={"approvalId": approval_id},
        )
        raise_for_status(resp)

    def reject(self, approval_id: str) -> None:
        """审批拒绝。"""
        resp = self._client.post(
            self._url("/reject"),
            headers=self._headers,
            json={"approvalId": approval_id},
        )
        raise_for_status(resp)

    def respond(self, ref: str, response: str) -> None:
        """回答 Agent 的提问（ask_user 事件）。"""
        resp = self._client.post(
            self._url("/respond"),
            headers=self._headers,
            json={"ref": ref, "response": response},
        )
        raise_for_status(resp)

    def stop(self) -> None:
        """停止 Agent，删除 Session。"""
        resp = self._client.delete(
            self._url(""),
            headers=self._headers,
        )
        raise_for_status(resp)

    # ------------------------------------------------------------------
    # 查询
    # ------------------------------------------------------------------

    def get_info(self) -> SessionInfo:
        """查询 Session 当前状态。"""
        resp = self._client.get(self._url(""), headers=self._headers)
        raise_for_status(resp)
        body = resp.json()
        return SessionInfo(
            session_id=body["sessionId"],
            state=body["state"],
            turns=body["turns"],
            tool_calls=body["toolCalls"],
            total_tokens=body["totalTokens"],
            uptime_ms=body["uptimeMs"],
        )

    def get_stats(self) -> Stats:
        """获取 Session 用量统计。"""
        resp = self._client.get(self._url("/stats"), headers=self._headers)
        raise_for_status(resp)
        body = resp.json()
        usage = body.get("tokenUsage", {})
        return Stats(
            session_id=body["sessionId"],
            turns=body["turns"],
            total_tokens=body["totalTokens"],
            cost_usd=body.get("costUsd", 0.0),
            prompt_count=body.get("promptCount", 0),
            token_usage=TokenUsage(
                prompt_tokens=usage.get("promptTokens", 0),
                completion_tokens=usage.get("completionTokens", 0),
                total_tokens=usage.get("totalTokens", 0),
            ),
        )

    def get_messages(self) -> list[Message]:
        """获取对话历史消息列表。"""
        resp = self._client.get(self._url("/messages"), headers=self._headers)
        raise_for_status(resp)
        body = resp.json()
        return [
            Message(
                role=m["role"],
                content=m.get("content", ""),
                tool_calls=m.get("toolCalls", []),
            )
            for m in body.get("messages", [])
        ]

    # ------------------------------------------------------------------
    # 本地工具注册
    # ------------------------------------------------------------------

    def register_tool(self, tool: LocalTool, callback_url: str) -> None:
        """向 Gateway 注册本地回调工具。

        Args:
            tool: :class:`~cmdc.tools.LocalTool` 工具描述符。
            callback_url: Gateway 回调时 POST 的本地 HTTP 端点 URL。
        """
        payload = tool.to_registration_payload(callback_url)
        resp = self._client.post(
            self._url("/tools"),
            headers=self._headers,
            json=payload,
        )
        raise_for_status(resp)

    # ------------------------------------------------------------------
    # WebSocket 连接
    # ------------------------------------------------------------------

    def connect_ws(self) -> WSConnection:
        """返回同步 WebSocket 连接对象。

        需要手动作为 context manager 使用::

            with session.connect_ws() as ws:
                ws.send_prompt("你好")
                for event in ws.iter_events():
                    ...
        """
        api_key = self._headers.get("X-API-Key", "")
        url = f"{self._ws_url()}?api_key={api_key}"
        return WSConnection(url)

    # ------------------------------------------------------------------
    # Context Manager（with client.session(...)）
    # ------------------------------------------------------------------

    def __enter__(self) -> "Session":
        return self

    def __exit__(self, *args: Any) -> None:
        try:
            self.stop()
        except Exception:
            pass


# ===========================================================================
# 异步版本
# ===========================================================================


class AsyncSession:
    """单个 Agent Session 的操作接口（异步）。

    通过 ``CMDC.create_session_async()`` 或 ``CMDC.asession()`` 获取。

    用法::

        async with await client.asession(model="deepseek:deepseek-chat") as session:
            async for event in session.aprompt("你好"):
                if event.type == "message_delta":
                    print(event.delta, end="")
    """

    def __init__(
        self,
        session_id: str,
        base_url: str,
        http_client: httpx.AsyncClient,
        headers: dict[str, str],
    ) -> None:
        self.session_id = session_id
        self._base_url = base_url.rstrip("/")
        self._client = http_client
        self._headers = headers

    def _url(self, path: str) -> str:
        return f"{self._base_url}/v1/sessions/{self.session_id}{path}"

    def _ws_url(self) -> str:
        url = self._url("/ws")
        return url.replace("https://", "wss://").replace("http://", "ws://")

    async def aprompt(
        self,
        text: str,
        *,
        approval_policy: ApprovalPolicyLike | None = None,
    ) -> AsyncIterator[Event]:
        """异步 SSE 流式迭代。

        用法::

            async for event in session.aprompt("你好"):
                if event.type == "message_delta":
                    print(event.delta, end="")
        """
        policy: ApprovalPolicy | None = None
        if approval_policy is not None:
            policy = make_policy(approval_policy)

        resp = await self._client.post(
            self._url("/prompt"),
            headers=self._headers,
            json={"text": text},
        )
        raise_for_status(resp)

        return self._stream_events(policy)

    async def _stream_events(
        self, policy: ApprovalPolicy | None
    ) -> AsyncIterator[Event]:
        async with AsyncSSEIterator(
            self._client, self._url("/events"), self._headers
        ) as it:
            async for event in it:
                if policy is not None and event.type == "approval_required":
                    approved = policy.decide(event)
                    approval_id = event.approval_id
                    if approval_id:
                        if approved:
                            await self.approve(approval_id)
                        else:
                            await self.reject(approval_id)
                yield event

    async def approve(self, approval_id: str) -> None:
        resp = await self._client.post(
            self._url("/approve"),
            headers=self._headers,
            json={"approvalId": approval_id},
        )
        raise_for_status(resp)

    async def reject(self, approval_id: str) -> None:
        resp = await self._client.post(
            self._url("/reject"),
            headers=self._headers,
            json={"approvalId": approval_id},
        )
        raise_for_status(resp)

    async def respond(self, ref: str, response: str) -> None:
        resp = await self._client.post(
            self._url("/respond"),
            headers=self._headers,
            json={"ref": ref, "response": response},
        )
        raise_for_status(resp)

    async def stop(self) -> None:
        resp = await self._client.delete(
            self._url(""),
            headers=self._headers,
        )
        raise_for_status(resp)

    def connect_ws(self) -> AsyncWSConnection:
        api_key = self._headers.get("X-API-Key", "")
        url = f"{self._ws_url()}?api_key={api_key}"
        return AsyncWSConnection(url)

    async def __aenter__(self) -> "AsyncSession":
        return self

    async def __aexit__(self, *args: Any) -> None:
        try:
            await self.stop()
        except Exception:
            pass
