"""
Gshield SDK - GshieldStreamClient 主类
"""
import json
import asyncio
import uuid
import logging
import threading
from typing import Optional, Dict

import websockets
import websockets.client

from .models import SessionConfig, ClassificationResult
from .session import StreamSession
from .exceptions import (
    GshieldConnectionError,
    SessionError,
    ProtocolError,
)

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


class GshieldStreamClient:
    """
    Gshield 流式文本安全检测 WebSocket 客户端。

    用法（异步）::

        async with GshieldStreamClient("ws://host:port/ws/classify") as client:
            session = await client.create_session(classify_every_n_chars=100)
            session.on_result(lambda e: print(e.result.main_label))
            for chunk in llm_stream():
                await session.send_delta(chunk)
            final = await session.finish()

    Args:
        url: WebSocket 服务地址，例如 ``ws://localhost:6006/ws/classify``
        extra_headers: 可选的额外请求头（用于认证等）
        ping_interval: WebSocket 心跳间隔（秒），默认 20
        ping_timeout: WebSocket 心跳超时（秒），默认 10
        max_reconnect_attempts: 最大重连次数，默认 3
    """

    def __init__(
        self,
        url: str,
        extra_headers: Optional[dict] = None,
        ping_interval: float = 20,
        ping_timeout: float = 10,
        max_reconnect_attempts: int = 3,
    ):
        self.url = url
        self.extra_headers = extra_headers
        self.ping_interval = ping_interval
        self.ping_timeout = ping_timeout
        self.max_reconnect_attempts = max_reconnect_attempts

        self._ws: Optional[websockets.client.WebSocketClientProtocol] = None
        self._sessions: Dict[str, StreamSession] = {}
        self._listener_task: Optional[asyncio.Task] = None
        self._connected = False

    # ============================================================
    # 连接管理
    # ============================================================

    async def connect(self):
        """建立 WebSocket 连接"""
        try:
            self._ws = await websockets.connect(
                self.url,
                extra_headers=self.extra_headers,
                ping_interval=self.ping_interval,
                ping_timeout=self.ping_timeout,
            )
            self._connected = True
            self._listener_task = asyncio.create_task(self._listen())
            logger.info(f"已连接到 {self.url}")
        except Exception as e:
            raise GshieldConnectionError(f"连接失败: {e}") from e

    async def disconnect(self):
        """断开 WebSocket 连接"""
        self._connected = False
        if self._listener_task:
            self._listener_task.cancel()
            try:
                await self._listener_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()
            self._ws = None

        # 强制停止所有会话
        for session in self._sessions.values():
            session._force_stop()
        self._sessions.clear()
        logger.info("连接已断开")

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.disconnect()

    # ============================================================
    # 会话管理
    # ============================================================

    async def create_session(
        self,
        session_id: Optional[str] = None,
        classify_every_n_chars: int = 100,
        classify_interval_secs: float = 3.0,
    ) -> StreamSession:
        """
        创建一个新的流式检测会话。

        Args:
            session_id: 可选的自定义会话 ID，不提供则由服务端生成
            classify_every_n_chars: 每累积 N 个新字符触发一次检测 (0=禁用)
            classify_interval_secs: 每隔 N 秒触发一次检测 (0=禁用)

        Returns:
            StreamSession 会话对象
        """
        if not self._connected or not self._ws:
            raise GshieldConnectionError("未连接到服务器")

        msg = {
            "type": "session.start",
            "config": {
                "classify_every_n_chars": classify_every_n_chars,
                "classify_interval_secs": classify_interval_secs,
            }
        }
        if session_id:
            msg["session_id"] = session_id

        await self._ws.send(json.dumps(msg, ensure_ascii=False))

        # 等待 session.created 响应（通过 listener 分发）
        # 创建一个临时 future 来等待
        pending_id = session_id or "__pending__"
        created_future = asyncio.get_running_loop().create_future()
        self._pending_creates = getattr(self, '_pending_creates', {})
        self._pending_creates[pending_id] = created_future

        try:
            resp = await asyncio.wait_for(created_future, timeout=10.0)
        except asyncio.TimeoutError:
            self._pending_creates.pop(pending_id, None)
            raise SessionError("创建会话超时")

        sid = resp["session_id"]
        config_data = resp.get("config", {})
        config = SessionConfig(
            classify_every_n_chars=config_data.get("classify_every_n_chars", classify_every_n_chars),
            classify_interval_secs=config_data.get("classify_interval_secs", classify_interval_secs),
        )

        session = StreamSession(session_id=sid, ws=self._ws, config=config)
        self._sessions[sid] = session
        logger.info(f"会话已创建: {sid}")
        return session

    # ============================================================
    # 内部 Listener
    # ============================================================

    async def _listen(self):
        """后台监听 WebSocket 消息并分发到对应 session"""
        try:
            async for raw in self._ws:
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    logger.warning(f"收到非 JSON 消息: {raw[:100]}")
                    continue

                msg_type = msg.get("type")
                session_id = msg.get("session_id")

                # 处理 session.created（需要唤醒 create_session）
                if msg_type == "session.created":
                    pending = getattr(self, '_pending_creates', {})
                    # 先尝试用 session_id 匹配
                    fut = pending.pop(session_id, None)
                    # 再尝试 __pending__（无指定 id 的情况）
                    if fut is None:
                        fut = pending.pop("__pending__", None)
                    if fut and not fut.done():
                        fut.set_result(msg)
                    continue

                # 处理 error（可能在 session 创建阶段）
                if msg_type == "error" and session_id:
                    # 检查是否有 pending create
                    pending = getattr(self, '_pending_creates', {})
                    fut = pending.pop(session_id, None) or pending.pop("__pending__", None)
                    if fut and not fut.done():
                        fut.set_exception(SessionError(
                            f"[{msg.get('code')}] {msg.get('message')}"
                        ))
                        continue

                # 分发到对应 session
                if session_id and session_id in self._sessions:
                    await self._sessions[session_id]._handle_message(msg)
                elif msg_type == "error":
                    logger.error(f"全局错误: [{msg.get('code')}] {msg.get('message')}")

        except websockets.exceptions.ConnectionClosed:
            logger.warning("WebSocket 连接已关闭")
            self._connected = False
            for session in self._sessions.values():
                session._force_stop()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"Listener 异常: {e}")
            self._connected = False


# ============================================================
# 同步封装
# ============================================================

class SyncGshieldStreamClient:
    """
    GshieldStreamClient 的同步封装，供非 async 环境使用。

    用法::

        with SyncGshieldStreamClient("ws://host:port/ws/classify") as client:
            session = client.create_session()
            for chunk in llm_stream():
                client.send_delta(session, chunk)
            final = client.finish(session)
            print(final)

    内部启动一个独立的事件循环线程。
    """

    def __init__(self, url: str, **kwargs):
        self._url = url
        self._kwargs = kwargs
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._thread: Optional[threading.Thread] = None
        self._client: Optional[GshieldStreamClient] = None

    def connect(self):
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._thread.start()

        self._client = GshieldStreamClient(self._url, **self._kwargs)
        future = asyncio.run_coroutine_threadsafe(self._client.connect(), self._loop)
        future.result(timeout=15)

    def disconnect(self):
        if self._client:
            future = asyncio.run_coroutine_threadsafe(self._client.disconnect(), self._loop)
            future.result(timeout=10)
        if self._loop:
            self._loop.call_soon_threadsafe(self._loop.stop)
        if self._thread:
            self._thread.join(timeout=5)

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, *args):
        self.disconnect()

    def create_session(
        self,
        session_id: Optional[str] = None,
        classify_every_n_chars: int = 100,
        classify_interval_secs: float = 3.0,
    ) -> StreamSession:
        future = asyncio.run_coroutine_threadsafe(
            self._client.create_session(
                session_id=session_id,
                classify_every_n_chars=classify_every_n_chars,
                classify_interval_secs=classify_interval_secs,
            ),
            self._loop,
        )
        return future.result(timeout=15)

    def send_delta(self, session: StreamSession, delta: str) -> int:
        future = asyncio.run_coroutine_threadsafe(
            session.send_delta(delta),
            self._loop,
        )
        return future.result(timeout=10)

    def finish(self, session: StreamSession) -> Optional[ClassificationResult]:
        future = asyncio.run_coroutine_threadsafe(
            session.finish(),
            self._loop,
        )
        return future.result(timeout=30)
