"""
WebSocket Handler - WebSocket 协议处理器

职责：
- 处理 WebSocket 连接生命周期
- 解析客户端消息，路由到对应的 Session 方法
- 将 Session 事件转发到客户端
- 实现 WebSocket 协议（心跳、错误处理等）
"""

import asyncio
import json
from typing import Dict, Any, Optional, List, Callable
from fastapi import WebSocket, WebSocketDisconnect

from core.logger import log_event
from core.config import AppConfig, save_settings, load_settings
from core.interaction import get_interaction_manager
from tools.base import BaseTool
from .constants import MessageType, EventType, AgentType
from .session import Session, ConnectionSessionManager


class WebSocketHandler:
    """
    WebSocket 协议处理器

    职责：
    1. 管理 WebSocket 连接生命周期
    2. 解析客户端消息并路由
    3. 转发 Session 事件到客户端
    """

    def __init__(
        self,
        session_manager: ConnectionSessionManager,
        model_name: str,
        config: AppConfig,
        mcp_tools_getter: Optional[Callable[[], List[BaseTool]]] = None,
    ):
        """
        初始化 WebSocket 处理器

        Args:
            session_manager: 会话管理器
            model_name: 默认模型名称
            config: 应用配置
            mcp_tools_getter: MCP 工具获取函数
        """
        self.session_manager = session_manager
        self.model_name = model_name
        self.config = config
        self.mcp_tools_getter = mcp_tools_getter

        # WebSocket 连接映射: session_id -> WebSocket
        self._connections: Dict[str, WebSocket] = {}

    async def handle_connection(self, websocket: WebSocket, session_id: str, agent_type: str = "general"):
        """
        处理 WebSocket 连接

        Args:
            websocket: WebSocket 连接
            session_id: 会话 ID
            agent_type: 客户端请求的 Agent 类型
        """
        await websocket.accept()
        self._connections[session_id] = websocket

        log_event(
            "ws_connected",
            session_id=session_id,
            agent_type=agent_type,
        )

        try:
            # 创建或获取会话
            session, is_new, resolved_agent_type = await self._get_or_create_session(session_id, websocket, agent_type)

            # 注册交互回调
            self._register_interaction_callback(session_id, websocket)

            # 只有新建会话才发送 CONNECTED 事件，重连不发送
            if is_new:
                import os
                await self._send_event(websocket, {
                    "type": EventType.CONNECTED,
                    "data": {
                        "session_id": session_id,
                        "agent_type": resolved_agent_type.value,
                        "working_path": os.getcwd(),
                    },
                })

            # 消息处理循环
            await self._message_loop(websocket, session_id, session)

        except WebSocketDisconnect:
            log_event(
                "ws_client_disconnected",
                session_id=session_id,
            )

        except Exception as e:
            import traceback
            log_event(
                "ws_error",
                session_id=session_id,
                error=str(e),
                traceback=traceback.format_exc(),
            )
            # 尝试发送错误信息给客户端
            try:
                await self._send_error(websocket, "SERVER_ERROR", str(e))
            except:
                pass

        finally:
            await self._cleanup_connection(session_id)

    async def _get_or_create_session(
        self,
        session_id: str,
        websocket: WebSocket,
        agent_type: str = "general",
    ) -> tuple[Session, bool, AgentType]:
        """获取或创建会话

        Returns:
            (session, is_new, resolved_agent_type): session 实例、是否是新建会话、解析后的 agent_type
        """
        session = self.session_manager.get_session(session_id)

        # 将客户端传来的 agent_type 字符串转为枚举
        try:
            resolved_agent_type = AgentType(agent_type)
        except ValueError:
            resolved_agent_type = AgentType.GENERAL

        # 创建事件回调（每次连接都需要更新）
        async def event_callback(event: Dict[str, Any]):
            await self._send_event(websocket, event)

        if not session:
            # 新建会话
            session = self.session_manager.create_session(
                session_id=session_id,
                model_name=self.model_name,
                config=self.config,
                mcp_tools_getter=self.mcp_tools_getter,
            )

            # 初始化会话
            await session.initialize(
                agent_type=resolved_agent_type,
                event_callback=event_callback,
            )
            return session, True, resolved_agent_type
        else:
            # 会话已存在（重连），更新事件回调和连接状态
            session._event_callback = event_callback
            session._is_connected = True
            log_event(
                "session_reconnected",
                session_id=session_id,
            )
            return session, False, resolved_agent_type

    def _register_interaction_callback(
        self,
        session_id: str,
        websocket: WebSocket,
    ):
        """注册交互回调"""
        interaction_manager = get_interaction_manager()
        root_session_id = session_id.split('.')[0]

        async def send_interaction(message: Dict[str, Any]):
            await websocket.send_json(message)

        interaction_manager.register_session(root_session_id, send_interaction)

        log_event(
            "interaction_registered",
            root_session_id=root_session_id,
        )

    async def _message_loop(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
    ):
        """消息处理循环"""
        while True:
            # 接收消息
            message_text = await websocket.receive_text()

            try:
                message = json.loads(message_text)
            except json.JSONDecodeError as e:
                log_event(
                    "ws_invalid_json",
                    session_id=session_id,
                    error=str(e),
                )
                await self._send_error(websocket, "INVALID_JSON", "无效的 JSON 格式")
                continue

            msg_type = message.get("type", "")
            data = message.get("data", {})

            log_event(
                "ws_received",
                session_id=session_id,
                type=msg_type,
            )

            # 路由消息
            try:
                await self._route_message(websocket, session_id, session, msg_type, data)
            except Exception as e:
                log_event(
                    "ws_route_message_error",
                    session_id=session_id,
                    msg_type=msg_type,
                    error=str(e),
                )
                await self._send_error(websocket, "HANDLER_ERROR", f"处理消息失败: {str(e)}")

    async def _route_message(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
        msg_type: str,
        data: Dict[str, Any],
    ):
        """路由消息到对应处理器"""
        if msg_type == MessageType.USER_MESSAGE:
            await self._handle_user_message(session, data)

        elif msg_type == MessageType.INTERRUPT:
            await self._handle_interrupt(session)

        elif msg_type == MessageType.SWITCH_AGENT:
            await self._handle_switch_agent(websocket, session_id, session, data)

        elif msg_type == MessageType.SWITCH_MODEL:
            await self._handle_switch_model(websocket, session_id, session, data)

        elif msg_type == MessageType.PERMISSION_RESPONSE:
            await self._handle_permission_response(session, data)

        elif msg_type == MessageType.INTERACTION_RESPONSE:
            await self._handle_interaction_response(session_id, data)

        elif msg_type == MessageType.QUERY_TODO:
            await self._handle_query_todo(websocket, session)

        elif msg_type == MessageType.QUERY_TOOLS:
            await self._handle_query_tools(websocket, session)

        elif msg_type == MessageType.QUERY_SKILLS:
            await self._handle_query_skills(websocket)

        elif msg_type == MessageType.COMPACT:
            await self._handle_compact(websocket, session_id, session)

        elif msg_type == MessageType.INIT:
            await self._handle_init(websocket, session_id, session, data)

        elif msg_type == "ping":
            await self._handle_ping(websocket)

        elif msg_type == "command":
            await self._handle_command(websocket, session_id, session, data)

        else:
            log_event(
                "ws_unknown_message_type",
                session_id=session_id,
                type=msg_type,
            )

    # ==================== 消息处理器 ====================

    async def _handle_user_message(self, session: Session, data: Dict[str, Any]):
        """处理用户消息"""
        content = data.get("content", "")
        if content.strip():
            # 在单独的任务中运行 Agent，避免阻塞消息循环
            # 这样权限响应等消息可以在 Agent 运行期间被接收
            async def run_agent():
                try:
                    await session.execute(content)
                except Exception as e:
                    log_event(
                        "agent_execute_error",
                        session_id=session.session_id,
                        error=str(e),
                    )
                    # 确保客户端收到错误事件，否则会卡死
                    if session._event_callback:
                        await session._event_callback({
                            "type": EventType.TASK_ERROR,
                            "data": {"message": f"Agent 执行异常: {str(e)}"},
                        })

            # 创建任务但不等待完成，让消息循环继续运行
            task = asyncio.create_task(run_agent())

            # 保存任务引用，防止被垃圾回收
            if not hasattr(self, '_agent_tasks'):
                self._agent_tasks = {}
            self._agent_tasks[session.session_id] = task

    async def _handle_interrupt(self, session: Session):
        """处理中断请求"""
        session.interrupt()
        log_event(
            "interrupt_requested",
            session_id=session.session_id,
        )

    async def _handle_switch_agent(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
        data: Dict[str, Any],
    ):
        """处理 Agent 切换"""
        agent_type_str = data.get("agent_type", "general")

        try:
            # 转换为 AgentType 枚举
            agent_type = AgentType(agent_type_str)
            old_type = await session.switch_agent(agent_type)

            # 发送切换成功事件
            await self._send_event(websocket, {
                "type": EventType.AGENT_SWITCHED,
                "data": {
                    "old_type": old_type.value,
                    "new_type": agent_type.value,
                },
            })

        except ValueError:
            await self._send_error(
                websocket,
                "INVALID_AGENT_TYPE",
                f"无效的 Agent 类型: {agent_type_str}",
            )

        except Exception as e:
            log_event(
                "switch_agent_error",
                session_id=session_id,
                error=str(e),
            )
            await self._send_error(websocket, "SWITCH_AGENT_ERROR", str(e))

    async def _handle_switch_model(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
        data: Dict[str, Any],
    ):
        """处理模型切换"""
        model_name = data.get("model_name", "")

        try:
            # Session 内部创建 client 并切换，返回旧/新模型信息
            model_info = await session.switch_model(model_name)

            # 更新 handler 的默认 model_name（新会话也用新模型）
            self.model_name = model_name

            # 保存到 setting.json
            settings = load_settings()
            settings["default_model"] = model_name
            save_settings(settings)

            # 发送切换成功事件
            await self._send_event(websocket, {
                "type": EventType.MODEL_SWITCHED,
                "data": {
                    "old_model": model_info["old_model"],
                    "new_model": model_info["new_model"],
                    "model_name": model_name,
                },
            })

        except ValueError as e:
            await self._send_error(
                websocket,
                "INVALID_MODEL",
                str(e),
            )

        except Exception as e:
            log_event(
                "switch_model_error",
                session_id=session_id,
                error=str(e),
            )
            await self._send_error(websocket, "SWITCH_MODEL_ERROR", str(e))

    async def _handle_permission_response(
        self,
        session: Session,
        data: Dict[str, Any],
    ):
        """处理权限响应"""
        request_id = data.get("request_id")
        approved = data.get("approved", False)
        remember = data.get("remember", False)
        remember_pattern = data.get("remember_pattern")

        if request_id:
            await session.handle_permission_response(
                request_id=request_id,
                approved=approved,
                remember=remember,
                remember_pattern=remember_pattern,
            )

    async def _handle_interaction_response(
        self,
        session_id: str,
        data: Dict[str, Any],
    ):
        """处理交互响应"""
        interaction_manager = get_interaction_manager()
        handled = interaction_manager.handle_response(data)

        log_event(
            "interaction_response_handled",
            session_id=session_id,
            request_id=data.get("request_id"),
            handled=handled,
        )

    async def _handle_query_todo(self, websocket: WebSocket, session: Session):
        """处理查询 TODO 请求"""
        # 从 session 获取 todo 信息
        todo_items = []

        # 可以从 agent 或其他地方获取 todo
        if session.agent:
            agent = session.agent
            # 尝试获取 todo 信息
            if hasattr(agent, '_todos'):
                todo_items = agent._todos.get(session.session_id, [])

        await websocket.send_json({
            "type": EventType.QUERY_TODO_RESPONSE,
            "data": {"todos": todo_items},
        })

    async def _handle_query_tools(self, websocket: WebSocket, session: Session):
        """处理查询工具列表请求"""
        tools = []
        categorized_tools = {}

        if session.tool_registry:
            # 获取当前 Agent 允许的工具列表
            allowed_tools = []
            if session.agent:
                allowed_tools = session.agent.get_tools(session.tool_registry)
            else:
                # 如果没有 agent，回退到所有工具
                allowed_tools = session.tool_registry.get_all_tools()

            # 允许的工具名称集合
            allowed_tool_names = {t.name for t in allowed_tools}

            # 过滤工具列表
            for tool in allowed_tools:
                tools.append({
                    "name": tool.name,
                    "description": tool.description,
                })

            # 获取分类工具列表（仅包含当前 Agent 允许的工具）
            categorized_tools = session.tool_registry.get_tools_by_category(
                allowed_tool_names=allowed_tool_names
            )

        await websocket.send_json({
            "type": EventType.QUERY_TOOLS_RESPONSE,
            "data": {"tools": tools, "categorized_tools": categorized_tools},
        })

    async def _handle_query_skills(self, websocket: WebSocket):
        """处理查询 skills 列表请求"""
        from core.skill import get_skill_manager

        manager = get_skill_manager()
        skills = [
            {"name": s.name, "description": s.description, "path": s.location}
            for s in manager.get_all_skills()
        ]

        await websocket.send_json({
            "type": EventType.QUERY_SKILLS_RESPONSE,
            "data": {"skills": skills},
        })

    async def _handle_compact(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
    ):
        """处理压缩上下文请求"""
        # 先检查上下文占比
        context_info = session.get_context_info(session_id)
        percent = context_info.get("percent", 0)

        if percent < 40:
            await websocket.send_json({
                "type": EventType.COMPACT_RESULT,
                "data": {
                    "success": True,
                    "skipped": True,
                    "message": f"上下文占比 {percent:.1f}% 未达到 40%，无需压缩",
                    "percent": percent,
                },
            })
            return

        # 执行压缩（强制执行，忽略阈值检查）
        try:
            result = await session.compact_session(force=True)
            await websocket.send_json({
                "type": EventType.COMPACT_RESULT,
                "data": {
                    "success": True,
                    "skipped": False,
                    "message": f"压缩完成（压缩前占比 {percent:.1f}%）",
                    "compressed_count": result.get("compressed_count", 0),
                },
            })
        except Exception as e:
            await websocket.send_json({
                "type": EventType.COMPACT_RESULT,
                "data": {
                    "success": False,
                    "message": f"压缩失败: {str(e)}",
                },
            })

    async def _handle_ping(self, websocket: WebSocket):
        """处理心跳"""
        await websocket.send_json({"type": "pong", "data": {}})

    async def _handle_init(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
        data: Dict[str, Any],
    ):
        """处理初始化项目请求"""
        from pathlib import Path

        force = data.get("force", False)
        agent_md = Path("AGENT.md")

        # 检查文件是否存在
        if agent_md.exists() and not force:
            await websocket.send_json({
                "type": EventType.TEXT_CHUNK,
                "data": {"content": "⚠️  AGENT.md 已存在。使用 /init --force 强制覆盖。\n"},
            })
            return

        # 读取提示词文件
        prompt_path = Path(__file__).parent.parent / "assets" / "prompts" / "commands" / "init.txt"
        if not prompt_path.exists():
            await websocket.send_json({
                "type": EventType.TEXT_CHUNK,
                "data": {"content": "❌ 提示词文件 assets/prompts/commands/init.txt 不存在\n"},
            })
            return

        prompt = prompt_path.read_text(encoding="utf-8")

        # 如果有 force，添加提示
        if force:
            prompt += "\n\n**注意：AGENT.md 已存在，请覆盖它。**"

        # 执行 Agent
        async def run_agent():
            try:
                await session.execute(prompt)
            except Exception as e:
                log_event(
                    "agent_execute_error",
                    session_id=session_id,
                    error=str(e),
                )
                if session._event_callback:
                    await session._event_callback({
                        "type": EventType.TASK_ERROR,
                        "data": {"message": f"Agent 执行异常: {str(e)}"},
                    })

        task = asyncio.create_task(run_agent())
        if not hasattr(self, '_agent_tasks'):
            self._agent_tasks = {}
        self._agent_tasks[session_id] = task

    async def _handle_command(
        self,
        websocket: WebSocket,
        session_id: str,
        session: Session,
        data: Dict[str, Any],
    ):
        """处理命令"""
        command = data.get("command", "")
        args = data.get("args", [])

        result = await self._execute_command(session_id, session, command, args)

        await websocket.send_json({
            "type": "command_result",
            "data": result,
        })

    async def _execute_command(
        self,
        session_id: str,
        session: Session,
        command: str,
        args: list,
    ) -> Dict[str, Any]:
        """执行命令"""
        if command == "status":
            return {
                "success": True,
                "message": (
                    f"Session: {session_id}\n"
                    f"Agent: {session.agent_type.value}\n"
                    f"Running: {session.is_running}\n"
                    f"Connected: {session.is_connected}"
                ),
            }

        elif command == "clear":
            session.clear_history()
            return {"success": True, "message": "历史已清除"}

        elif command in ["exit", "quit"]:
            return {"success": True, "message": "再见！", "should_exit": True}

        else:
            return {"success": False, "message": f"未知命令: {command}"}

    # ==================== 辅助方法 ====================

    async def _send_event(self, websocket: WebSocket, event: Dict[str, Any]):
        """发送事件到客户端"""
        try:
            # 确保 type 字段是字符串（处理枚举值）
            if "type" in event and hasattr(event["type"], "value"):
                event = {**event, "type": event["type"].value}

            # Debug: 记录事件详细信息
            import sys
            if event.get("type") == "connected":
                data = event.get("data", {})
                print(f"DEBUG: connected event data keys: {list(data.keys())}", file=sys.stderr)
                print(f"DEBUG: connected event data: {data}", file=sys.stderr)

            log_event("ws_sending_event", event_type=event.get("type", "unknown"), data=event.get("data"))
            await websocket.send_json(event)
        except Exception as e:
            log_event(
                "ws_send_error",
                error=str(e),
            )

    async def _send_error(
        self,
        websocket: WebSocket,
        error_code: str,
        message: str,
    ):
        """发送错误事件"""
        await self._send_event(websocket, {
            "type": EventType.ERROR,
            "data": {
                "error_code": error_code,
                "message": message,
            },
        })

    async def _cleanup_connection(self, session_id: str):
        """清理 WebSocket 连接（保留 session 供重连）"""
        # 移除连接映射
        if session_id in self._connections:
            del self._connections[session_id]

        # 注销交互回调（WebSocket 引用已失效）
        interaction_manager = get_interaction_manager()
        root_session_id = session_id.split('.')[0]
        interaction_manager.unregister_session(root_session_id)

        # 标记 session 为未连接，但不销毁（保留上下文供重连）
        session = self.session_manager.get_session(session_id)
        if session:
            session._is_connected = False

        log_event(
            "ws_disconnected",
            session_id=session_id,
        )

    def get_active_connections(self) -> int:
        """获取活跃连接数"""
        return len(self._connections)
