"""
Gshield SDK - StreamSession 会话封装
"""
import json
import asyncio
import logging
from typing import Optional, Callable, Any, AsyncIterator

from .models import (
    ClassificationResult,
    StreamClassificationEvent,
    SessionConfig,
)
from .exceptions import (
    SessionError,
    ClassifyError,
    ProtocolError,
    GshieldConnectionError,
)

logger = logging.getLogger("gshield.session")


class StreamSession:
    """
    流式检测会话。

    在一个 WebSocket 连接上管理单次检测会话的全生命周期：
    - send_delta(): 发送文本增量
    - finish(): 结束会话并获取终检结果
    - on_result(): 注册检测结果回调
    - results(): 异步迭代器，逐个 yield 检测结果
    """

    def __init__(self, session_id: str, ws, config: SessionConfig):
        self.session_id = session_id
        self._ws = ws
        self.config = config
        self._active = True
        self._result_callbacks: list[Callable] = []
        self._result_queue: asyncio.Queue[Optional[StreamClassificationEvent]] = asyncio.Queue()
        self._listener_task: Optional[asyncio.Task] = None
        self._final_result: Optional[ClassificationResult] = None
        self._finished_event = asyncio.Event()

    # ---------- 公开 API ----------

    def on_result(self, callback: Callable[[StreamClassificationEvent], Any]):
        """
        注册检测结果回调函数。

        Args:
            callback: 接收 StreamClassificationEvent 参数的回调函数（同步或异步均可）
        """
        self._result_callbacks.append(callback)

    async def send_delta(self, delta: str) -> int:
        """
        发送文本增量。

        Args:
            delta: 新增的文本片段

        Returns:
            服务端确认后的当前累积文本长度
        """
        if not self._active:
            raise SessionError(f"会话 {self.session_id} 已结束")

        msg = {
            "type": "text.delta",
            "session_id": self.session_id,
            "delta": delta,
        }
        await self._ws.send(json.dumps(msg, ensure_ascii=False))

        # 注意：text.received 回复由 _listener 处理，这里不主动等待
        # 返回 -1 表示尚不确定长度（异步推送模式）
        return -1

    async def finish(self) -> Optional[ClassificationResult]:
        """
        结束会话，触发终检并返回最终分类结果。

        Returns:
            最终分类结果，如果缓冲区为空则返回 None
        """
        if not self._active:
            raise SessionError(f"会话 {self.session_id} 已结束")

        msg = {
            "type": "session.end",
            "session_id": self.session_id,
        }
        await self._ws.send(json.dumps(msg, ensure_ascii=False))
        self._active = False

        # 等待 session.ended 消息（由 listener 设置）
        await self._finished_event.wait()

        # 通知结果迭代器结束
        await self._result_queue.put(None)

        return self._final_result

    async def results(self) -> AsyncIterator[StreamClassificationEvent]:
        """
        异步迭代器，逐个 yield 检测结果（包括中间检测和终检）。

        用法::

            async for event in session.results():
                print(event.result.main_label)
        """
        while True:
            event = await self._result_queue.get()
            if event is None:
                break
            yield event

    @property
    def is_active(self) -> bool:
        return self._active

    # ---------- 内部：消息分发 ----------

    async def _handle_message(self, msg: dict):
        """
        处理与本 session 相关的服务端消息。
        由 GshieldStreamClient 的全局 listener 调用。
        """
        msg_type = msg.get("type")

        if msg_type == "text.received":
            # 轻量确认，忽略即可
            pass

        elif msg_type == "classification.result":
            result_data = msg.get("result", {})
            event = StreamClassificationEvent(
                session_id=self.session_id,
                seq=msg.get("seq", 0),
                text_length=msg.get("text_length", 0),
                result=ClassificationResult(**result_data),
                is_final=msg.get("is_final", False),
            )
            # 放入队列供异步迭代器消费
            await self._result_queue.put(event)

            # 触发回调
            for cb in self._result_callbacks:
                try:
                    ret = cb(event)
                    if asyncio.iscoroutine(ret):
                        await ret
                except Exception as e:
                    logger.warning(f"Session {self.session_id}: 回调异常: {e}")

        elif msg_type == "session.ended":
            final_data = msg.get("final_result")
            if final_data:
                self._final_result = ClassificationResult(**final_data)
            self._active = False
            self._finished_event.set()

        elif msg_type == "error":
            code = msg.get("code", "UNKNOWN")
            message = msg.get("message", "")
            logger.error(f"Session {self.session_id}: 服务端错误 [{code}] {message}")
            # 对于会话级错误，标记结束
            if code in ("SESSION_NOT_FOUND",):
                self._active = False
                self._finished_event.set()
                await self._result_queue.put(None)

    def _force_stop(self):
        """连接断开时强制停止"""
        self._active = False
        self._finished_event.set()
