"""
WebSocket 双向连接。

适合 HITL 场景：连接后同时接收事件推送 + 发送 approve/reject/respond/prompt。
"""

from __future__ import annotations

import asyncio
import json
import threading
from collections.abc import AsyncIterator, Iterator
from queue import Empty, Queue
from typing import Any

import websockets
from websockets.sync.client import connect as ws_connect
from websockets.asyncio.client import connect as aws_connect

from cmdc.types import Event


class WSConnection:
    """同步 WebSocket 连接。

    用法::

        with WSConnection(url) as ws:
            ws.send_prompt("你好")
            for event in ws.iter_events():
                if event.type == "approval_required":
                    ws.approve(event.approval_id)
                elif event.type == "agent_end":
                    break
    """

    def __init__(self, url: str) -> None:
        self._url = url
        self._ws: Any = None

    def __enter__(self) -> "WSConnection":
        self._ws = ws_connect(self._url)
        return self

    def __exit__(self, *args: Any) -> None:
        if self._ws is not None:
            self._ws.close()

    def _send(self, payload: dict[str, Any]) -> None:
        if self._ws is None:
            raise RuntimeError("WSConnection is not connected")
        self._ws.send(json.dumps(payload))

    def send_prompt(self, text: str) -> dict[str, Any]:
        """向 Agent 发送 prompt 消息。"""
        self._send({"action": "prompt", "text": text})
        raw = self._ws.recv()
        return json.loads(raw)

    def approve(self, approval_id: str) -> dict[str, Any]:
        """审批通过。"""
        self._send({"action": "approve", "approvalId": approval_id})
        raw = self._ws.recv()
        return json.loads(raw)

    def reject(self, approval_id: str) -> dict[str, Any]:
        """审批拒绝。"""
        self._send({"action": "reject", "approvalId": approval_id})
        raw = self._ws.recv()
        return json.loads(raw)

    def respond(self, ref: str, response: str) -> dict[str, Any]:
        """回答 Agent 提问。"""
        self._send({"action": "respond", "ref": ref, "response": response})
        raw = self._ws.recv()
        return json.loads(raw)

    def iter_events(self) -> Iterator[Event]:
        """迭代接收所有 Agent 事件，直到 agent_end / agent_abort。"""
        if self._ws is None:
            raise RuntimeError("WSConnection is not connected")
        while True:
            try:
                raw = self._ws.recv()
            except Exception:
                break
            msg = json.loads(raw)
            event_type = msg.get("event")
            if not event_type:
                continue
            event = Event(type=event_type, data=msg.get("data", {}))
            yield event
            if event_type in ("agent_end", "agent_abort"):
                break


class AsyncWSConnection:
    """异步 WebSocket 连接。

    用法::

        async with AsyncWSConnection(url) as ws:
            await ws.send_prompt("你好")
            async for event in ws.iter_events():
                if event.type == "approval_required":
                    await ws.approve(event.approval_id)
                elif event.type == "agent_end":
                    break
    """

    def __init__(self, url: str) -> None:
        self._url = url
        self._ws: Any = None

    async def __aenter__(self) -> "AsyncWSConnection":
        self._ws = await aws_connect(self._url)
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._ws is not None:
            await self._ws.close()

    async def _send(self, payload: dict[str, Any]) -> None:
        if self._ws is None:
            raise RuntimeError("AsyncWSConnection is not connected")
        await self._ws.send(json.dumps(payload))

    async def send_prompt(self, text: str) -> dict[str, Any]:
        """向 Agent 发送 prompt 消息。"""
        await self._send({"action": "prompt", "text": text})
        raw = await self._ws.recv()
        return json.loads(raw)

    async def approve(self, approval_id: str) -> dict[str, Any]:
        """审批通过。"""
        await self._send({"action": "approve", "approvalId": approval_id})
        raw = await self._ws.recv()
        return json.loads(raw)

    async def reject(self, approval_id: str) -> dict[str, Any]:
        """审批拒绝。"""
        await self._send({"action": "reject", "approvalId": approval_id})
        raw = await self._ws.recv()
        return json.loads(raw)

    async def respond(self, ref: str, response: str) -> dict[str, Any]:
        """回答 Agent 提问。"""
        await self._send({"action": "respond", "ref": ref, "response": response})
        raw = await self._ws.recv()
        return json.loads(raw)

    async def iter_events(self) -> AsyncIterator[Event]:
        """异步迭代接收所有 Agent 事件。"""
        if self._ws is None:
            raise RuntimeError("AsyncWSConnection is not connected")
        while True:
            try:
                raw = await self._ws.recv()
            except Exception:
                break
            msg = json.loads(raw)
            event_type = msg.get("event")
            if not event_type:
                continue
            event = Event(type=event_type, data=msg.get("data", {}))
            yield event
            if event_type in ("agent_end", "agent_abort"):
                break
