"""
WebSocket 客户端

支持：
- 协议级 ping/pong（websockets 库内置）
- 应用层心跳（{"type": "ping"}）
- 断线自动重连（指数退避，最多 5 次）
"""

import asyncio
import json
from typing import Optional, Callable, Dict, Any
import aiohttp
import websockets
from websockets.client import WebSocketClientProtocol

from core.logger import log_event, log_error


# ==================== 连接配置 ====================

# 协议级 ping：websockets 库自动发送 WebSocket ping 帧
WS_PING_INTERVAL = 30      # 每 30 秒发一次 ping 帧
WS_PING_TIMEOUT = 60       # 60 秒没收到 pong 视为断开（足够长以支持用户交互）

# 应用层心跳：发送 {"type": "ping"} 消息（保留但简化）
HEARTBEAT_INTERVAL = 30     # 每 30 秒发一次应用层心跳

# 自动重连
RECONNECT_MAX_RETRIES = 5   # 最大重连次数
RECONNECT_BASE_DELAY = 1.0  # 首次重连延迟（秒）
RECONNECT_MAX_DELAY = 16.0  # 最大重连延迟（秒）


class WebSocketClient:
    """WebSocket 客户端"""

    def __init__(self, url: str, session_id: str, agent_type: str = "build"):
        """
        Args:
            url: WebSocket 服务器 URL (例如: ws://localhost:8000)
            session_id: 会话 ID
            agent_type: Agent 类型（build/general/plan/explore）
        """
        self.url = url
        self.session_id = session_id
        self.agent_type = agent_type
        self.ws: Optional[WebSocketClientProtocol] = None
        self.event_handlers: Dict[str, Callable] = {}
        self._running = False
        # 断开连接和错误回调
        self._on_disconnect_callback: Optional[Callable] = None
        self._on_error_callback: Optional[Callable] = None
        # 内部任务
        self._heartbeat_task: Optional[asyncio.Task] = None
        self._reconnecting = False

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

    async def connect(self):
        """连接到服务器"""
        # session_id 为空时使用 /ws 端点（让服务端生成）
        # session_id 有值时使用 /ws/{session_id} 端点（让服务端重新生成）
        if self.session_id:
            ws_url = f"{self.url}/ws/{self.session_id}?agent_type={self.agent_type}"
        else:
            ws_url = f"{self.url}/ws?agent_type={self.agent_type}"
        self.ws = await websockets.connect(
            ws_url,
            ping_interval=WS_PING_INTERVAL,
            ping_timeout=WS_PING_TIMEOUT,
        )
        self._running = True
        log_event("client_connected", session_id=self.session_id)

    async def disconnect(self, reason: str = ""):
        """断开连接"""
        self._running = False
        self._stop_heartbeat()
        log_event("client_disconnect_requested", session_id=self.session_id, reason=reason)
        if self.ws:
            await self.ws.close()
            log_event("client_disconnected", session_id=self.session_id, reason=reason)

    async def _reconnect(self) -> bool:
        """
        自动重连（指数退避）

        Returns:
            bool: 重连是否成功
        """
        if self._reconnecting:
            return False
        self._reconnecting = True

        try:
            delay = RECONNECT_BASE_DELAY
            for attempt in range(1, RECONNECT_MAX_RETRIES + 1):
                log_event(
                    "ws_reconnect_attempt",
                    session_id=self.session_id,
                    attempt=attempt,
                    max_retries=RECONNECT_MAX_RETRIES,
                    delay=delay,
                )
                await asyncio.sleep(delay)

                try:
                    ws_url = f"{self.url}/ws/{self.session_id}?agent_type={self.agent_type}"
                    self.ws = await websockets.connect(
                        ws_url,
                        ping_interval=WS_PING_INTERVAL,
                        ping_timeout=WS_PING_TIMEOUT,
                    )
                    self._running = True
                    log_event(
                        "ws_reconnect_success",
                        session_id=self.session_id,
                        attempt=attempt,
                    )
                    return True
                except Exception as e:
                    log_error(
                        f"ws_reconnect_failed (attempt {attempt}/{RECONNECT_MAX_RETRIES}): "
                        f"[{type(e).__name__}] {str(e) or repr(e)}",
                    )
                    delay = min(delay * 2, RECONNECT_MAX_DELAY)

            log_error(
                f"ws_reconnect_exhausted: 重连 {RECONNECT_MAX_RETRIES} 次均失败",
                session_id=self.session_id,
            )
            return False
        finally:
            self._reconnecting = False

    # ==================== 心跳 ====================

    def _start_heartbeat(self):
        """启动应用层心跳"""
        self._stop_heartbeat()
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())

    def _stop_heartbeat(self):
        """停止应用层心跳"""
        if self._heartbeat_task and not self._heartbeat_task.done():
            self._heartbeat_task.cancel()
            try:
                # 等待任务取消完成
                import asyncio
                asyncio.get_event_loop().run_until_complete(self._heartbeat_task)
            except Exception:
                pass
        self._heartbeat_task = None

    async def _heartbeat_loop(self):
        """应用层心跳循环，失败触发重连"""
        try:
            while self._running:
                await asyncio.sleep(HEARTBEAT_INTERVAL)
                if self._running and self.ws:
                    try:
                        await self.ping()
                    except Exception as e:
                        log_error(f"heartbeat_send_failed: [{type(e).__name__}] {str(e) or repr(e)}")
                        break  # 跳出循环，触发重连
        except asyncio.CancelledError:
            pass

    # ==================== 事件注册 ====================

    def on(self, event_type: str, handler: Callable):
        """注册事件处理器"""
        self.event_handlers[event_type] = handler

    def on_disconnect(self, callback: Callable):
        """注册断开连接回调"""
        self._on_disconnect_callback = callback

    def on_error(self, callback: Callable):
        """注册错误回调"""
        self._on_error_callback = callback

    # ==================== 消息发送 ====================

    async def send_message(self, content: str):
        """发送用户消息"""
        message = {
            "type": "user_message",
            "data": {"content": content}
        }
        await self._send(message)

    async def switch_agent(self, agent_type: str):
        """切换 Agent 类型"""
        message = {
            "type": "switch_agent",
            "data": {
                "session_id": self.session_id,
                "agent_type": agent_type
            }
        }
        await self._send(message)

    async def switch_model(self, model_name: str):
        """切换模型"""
        message = {
            "type": "switch_model",
            "data": {
                "session_id": self.session_id,
                "model_name": model_name
            }
        }
        await self._send(message)

    async def execute_command(self, command: str, args: list = None):
        """执行命令"""
        message = {
            "type": "command",
            "data": {
                "session_id": self.session_id,
                "command": command,
                "args": args or []
            }
        }
        await self._send(message)

    async def interrupt(self):
        """中断当前执行"""
        message = {
            "type": "interrupt",
            "data": {"session_id": self.session_id}
        }
        await self._send(message)

    async def query_todo(self):
        """查询 TODO 列表"""
        message = {
            "type": "query_todo",
            "data": {}
        }
        await self._send(message)

    async def query_tools(self):
        """查询工具列表"""
        message = {
            "type": "query_tools",
            "data": {}
        }
        await self._send(message)

    async def query_skills(self):
        """查询 skills 列表"""
        message = {
            "type": "query_skills",
            "data": {}
        }
        await self._send(message)

    async def init_project(self, force: bool = False):
        """初始化项目"""
        message = {
            "type": "init",
            "data": {"force": force}
        }
        await self._send(message)

    async def context_compact(self):
        """请求压缩上下文"""
        message = {
            "type": "compact",
            "data": {}
        }
        await self._send(message)

    async def ping(self):
        """发送应用层心跳"""
        message = {"type": "ping"}
        await self._send(message)

    async def get_model_list(self) -> list:
        """获取模型列表（HTTP 接口）"""
        http_url = self.url.replace("ws://", "http://").replace("wss://", "https://")
        url = f"{http_url}/models"
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as resp:
                data = await resp.json()
                return data.get("models", [])

    async def get_context_info(self) -> dict:
        """获取上下文信息（HTTP 接口）"""
        http_url = self.url.replace("ws://", "http://").replace("wss://", "https://")
        url = f"{http_url}/context/{self.session_id}"
        async with aiohttp.ClientSession() as session:
            async with session.get(url) as resp:
                return await resp.json()

    async def send_permission_response(
        self,
        request_id: str,
        approved: bool,
        remember: bool = False,
        remember_pattern: str = None
    ):
        """发送权限响应（旧版兼容）"""
        log_event("ws_client_send_permission_response",
            request_id=request_id,
            approved=approved,
            remember=remember,
            remember_pattern=remember_pattern
        )

        message = {
            "type": "permission_response",
            "data": {
                "request_id": request_id,
                "approved": approved,
                "remember": remember,
                "remember_pattern": remember_pattern
            }
        }

        log_event("ws_client_sending_message", message=message)
        await self._send(message)
        log_event("ws_client_message_sent")

    async def send_interaction_response(
        self,
        request_id: str,
        approved: bool,
        value: Any = None,
        remember: bool = False,
        remember_scope: str = "session"
    ):
        """
        发送交互响应（统一交互通道）

        Args:
            request_id: 请求 ID
            approved: 是否批准/确认
            value: 用户输入的值
            remember: 是否记住选择
            remember_scope: 记忆范围 (session/pattern/always)
        """
        log_event("ws_client_send_interaction_response",
            request_id=request_id,
            approved=approved,
            value=value,
            remember=remember,
            remember_scope=remember_scope
        )

        message = {
            "type": "interaction_response",
            "data": {
                "request_id": request_id,
                "approved": approved,
                "value": value,
                "remember": remember,
                "remember_scope": remember_scope
            }
        }

        await self._send(message)
        log_event("ws_client_interaction_response_sent")

    async def _send(self, message: Dict[str, Any]):
        """发送消息到服务器"""
        if not self.ws:
            raise RuntimeError("Not connected")
        await self.ws.send(json.dumps(message))

    # ==================== 接收循环 ====================

    async def receive_loop(self):
        """接收消息循环（含自动重连）"""
        if not self.ws:
            raise RuntimeError("Not connected")

        # 启动应用层心跳
        self._start_heartbeat()

        while self._running:
            try:
                while self._running:
                    message_text = await self.ws.recv()
                    message = json.loads(message_text)
                    await self._handle_event(message)

            except websockets.exceptions.ConnectionClosed as e:
                log_event(
                    "connection_closed",
                    session_id=self.session_id,
                    code=e.code,
                    reason=str(e.reason),
                )
                # code=1000 表示正常关闭（用户主动退出）
                if e.code == 1000:
                    log_event("connection_closed_normally", session_id=self.session_id)
                    self._stop_heartbeat()
                    return

                # 非正常关闭，尝试自动重连
                self._stop_heartbeat()
                if await self._reconnect():
                    self._start_heartbeat()
                    continue
                # 重连失败，通知 CLI
                await self._notify_disconnect(
                    f"WebSocket 连接已断开 (code={e.code}, reason={e.reason or '未知'})，重连失败"
                )
                return

            except Exception as e:
                log_error(
                    f"ws_receive_error: [{type(e).__name__}] {str(e) or repr(e)}",
                    session_id=self.session_id,
                )
                # 尝试重连
                self._stop_heartbeat()
                if await self._reconnect():
                    self._start_heartbeat()
                    continue
                # 重连失败，通知 CLI
                await self._notify_error(str(e))
                return

        self._stop_heartbeat()

    # ==================== 内部回调 ====================

    async def _notify_disconnect(self, reason: str):
        """通知 CLI 连接断开"""
        if self._on_disconnect_callback:
            if asyncio.iscoroutinefunction(self._on_disconnect_callback):
                await self._on_disconnect_callback(reason)
            else:
                self._on_disconnect_callback(reason)

    async def _notify_error(self, error: str):
        """通知 CLI 发生错误"""
        if self._on_error_callback:
            if asyncio.iscoroutinefunction(self._on_error_callback):
                await self._on_error_callback(error)
            else:
                self._on_error_callback(error)

    async def _handle_event(self, event: Dict[str, Any]):
        """处理接收到的事件"""
        event_type = event.get("type")

        # pong 不需要上报
        if event_type == "pong":
            return

        # 调用注册的处理器
        if event_type in self.event_handlers:
            handler = self.event_handlers[event_type]
            try:
                if asyncio.iscoroutinefunction(handler):
                    await handler(event)
                else:
                    handler(event)
            except Exception as e:
                log_event(
                    "event_handler_error",
                    event_type=event_type,
                    error=str(e),
                )
        else:
            # 默认处理器
            log_event("unhandled_event", unhandled_type=event_type)
