"""
Mind CLI Bridge — Cloud 端 WebSocket Tunnel 管理器。

管理所有 Mind CLI 用户的 Tunnel 连接：
- 接收 CLI 的 WebSocket 握手 + JWT 认证
- 能力审计（审批通过的工具白名单）
- 工具调用派发（Cloud LLM → Tunnel → CLI → 结果）
"""

import asyncio
import json
import logging
import time
import uuid
from typing import Any

from aiohttp import web

logger = logging.getLogger("mindcli_bridge")

# ── 初版白名单（所有内置工具默认通过） ────────────────────
DEFAULT_APPROVED_TOOLS = [
    "terminal",
    "file_read",
    "file_write",
    "file_ops",
    "grep",
    "code_execution",
]

# 工具调用超时
TOOL_CALL_TIMEOUT = 120  # 秒
# 心跳检测：60s 无消息判定离线
HEARTBEAT_TIMEOUT = 60


class MindCLIBridge:
    """
    管理所有 CLI 用户的 Tunnel 连接。

    每个 userId 最多维持一个活跃 WebSocket 连接。
    """

    def __init__(self):
        # userId → WebSocketResponse
        self._connections: dict[str, web.WebSocketResponse] = {}
        # userId → 能力报告
        self._capabilities: dict[str, dict] = {}
        # userId → 最后心跳时间
        self._last_heartbeat: dict[str, float] = {}
        # 待处理的 tool_call 请求
        self._pending: dict[str, asyncio.Future] = {}

    async def handle_tunnel_connect(self, request: web.Request) -> web.WebSocketResponse:
        """
        WebSocket 握手入口。

        路由：GET /mindos-next/ws/cli-tunnel
        认证：Header Authorization: Bearer <JWT>
        """
        ws = web.WebSocketResponse(heartbeat=30, autoping=True)
        await ws.prepare(request)

        # ── JWT 认证 ──
        token = _extract_token(request)
        if not token:
            await ws.send_json({"type": "error", "message": "Missing Authorization header"})
            await ws.close()
            return ws

        # 使用 MindPass 认证策略
        user = await self._verify_token(token, request)
        if not user:
            await ws.send_json({"type": "error", "message": "Invalid or expired token"})
            await ws.close()
            return ws

        user_id = user["userId"]

        # 关闭旧连接
        old_ws = self._connections.get(user_id)
        if old_ws and not old_ws.closed:
            logger.info("[Bridge] 关闭 %s 的旧连接", user_id)
            await old_ws.close()

        self._connections[user_id] = ws
        self._last_heartbeat[user_id] = time.time()
        logger.info("[Bridge] CLI 连接: userId=%s", user_id)

        # ── 发送连接确认 ──
        await ws.send_json({
            "type": "connected",
            "userId": user_id,
        })

        # ── 等待能力报告 ──
        try:
            msg = await asyncio.wait_for(ws.receive_json(), timeout=10)
            if msg.get("type") == "capability_report":
                self._capabilities[user_id] = msg
                logger.info(
                    "[Bridge] 能力报告: userId=%s, tools=%d, mcp=%d",
                    user_id,
                    len(msg.get("tools", [])),
                    len(msg.get("mcp_servers", [])),
                )

                # 审批工具
                approved = self._audit_capabilities(msg)
                await ws.send_json({
                    "type": "approved_tools",
                    "tools": approved,
                })
        except asyncio.TimeoutError:
            logger.warning("[Bridge] %s 未在 10s 内发送能力报告", user_id)

        # ── 消息循环 ──
        try:
            async for msg in ws:
                if msg.type == web.WSMsgType.TEXT:
                    data = json.loads(msg.data)
                    await self._handle_message(user_id, data)
                elif msg.type in (web.WSMsgType.ERROR, web.WSMsgType.CLOSE):
                    break
        except Exception as e:
            logger.warning("[Bridge] %s 连接异常: %s", user_id, e)
        finally:
            self._cleanup(user_id)

        return ws

    def is_connected(self, user_id: str) -> bool:
        """检查用户的 CLI 是否在线。"""
        ws = self._connections.get(user_id)
        return ws is not None and not ws.closed

    def get_capabilities(self, user_id: str) -> dict | None:
        """获取用户 CLI 的能力报告。"""
        return self._capabilities.get(user_id)

    async def dispatch_tool_call(
        self, user_id: str, tool_name: str, args: dict
    ) -> dict:
        """
        派发工具调用：Cloud LLM → Tunnel → CLI → 结果。

        Args:
            user_id: 用户 ID
            tool_name: 工具名（如 "terminal"）
            args: 工具参数

        Returns:
            {"output": "...", "exit_code": 0} 或 {"error": "..."}
        """
        ws = self._connections.get(user_id)
        if not ws or ws.closed:
            return {"error": "CLI not connected"}

        call_id = str(uuid.uuid4())[:8]

        # 创建 Future 等待 CLI 响应
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._pending[call_id] = future

        try:
            # 发送 JSON-RPC tool_call
            await ws.send_json({
                "jsonrpc": "2.0",
                "method": "tool_call",
                "params": {"tool": tool_name, "args": args},
                "id": call_id,
            })
            logger.info("[Bridge] → %s: %s (id=%s)", user_id, tool_name, call_id)

            # 等待结果
            result = await asyncio.wait_for(future, timeout=TOOL_CALL_TIMEOUT)
            logger.info("[Bridge] ← %s: %s 完成 (id=%s)", user_id, tool_name, call_id)
            return result

        except asyncio.TimeoutError:
            logger.warning("[Bridge] %s: %s 超时 (id=%s)", user_id, tool_name, call_id)
            return {"error": f"Tool call timed out after {TOOL_CALL_TIMEOUT}s"}
        finally:
            self._pending.pop(call_id, None)

    # ── 内部方法 ──────────────────────────────────────

    async def _verify_token(self, token: str, request: web.Request) -> dict | None:
        """JWT 验证（复用 MindPass 认证策略）。"""
        try:
            from platforms.sse_base.auth_strategies import MindPassAuth
            auth = MindPassAuth()
            return await auth.verify(token)
        except Exception as e:
            logger.error("[Bridge] Token 验证失败: %s", e)
            return None

    def _audit_capabilities(self, report: dict) -> list[str]:
        """
        审计 CLI 上报的能力，返回审批通过的工具列表。

        初版：白名单 = DEFAULT_APPROVED_TOOLS 与上报工具的交集。
        后续：从 DB 读取 per-user 配置。
        """
        reported_tools = {t["name"] for t in report.get("tools", [])}
        approved = [t for t in DEFAULT_APPROVED_TOOLS if t in reported_tools]
        return approved

    async def _handle_message(self, user_id: str, data: dict) -> None:
        """处理 CLI 发来的消息。"""
        self._last_heartbeat[user_id] = time.time()

        # JSON-RPC 响应（工具调用结果）
        if data.get("jsonrpc") == "2.0" and "id" in data:
            call_id = data["id"]
            future = self._pending.get(call_id)
            if future and not future.done():
                result = data.get("result", {"error": "Empty result"})
                future.set_result(result)
            return

        # 心跳 pong
        if data.get("type") == "pong":
            return

        logger.debug("[Bridge] %s 未知消息: %s", user_id, data.get("type"))

    def _cleanup(self, user_id: str) -> None:
        """连接关闭时清理。"""
        self._connections.pop(user_id, None)
        self._capabilities.pop(user_id, None)
        self._last_heartbeat.pop(user_id, None)

        # 取消所有待处理的 Future
        to_cancel = [k for k, v in self._pending.items() if not v.done()]
        for k in to_cancel:
            self._pending[k].set_exception(ConnectionError("CLI disconnected"))
            del self._pending[k]

        logger.info("[Bridge] CLI 断开: userId=%s", user_id)


def _extract_token(request: web.Request) -> str | None:
    """从 Authorization header 或 query 参数提取 token。"""
    auth_header = request.headers.get("Authorization", "")
    if auth_header.startswith("Bearer "):
        return auth_header[7:]
    # 备选：query 参数（WebSocket 某些客户端不支持 header）
    return request.query.get("token")
