"""
小狸 LSP 客户端
==============
连接外部语言服务器（Language Server Protocol），把真实的诊断 / 悬停 /
定义跳转 / 引用 / 补全喂给模型——这是对标 opencode 的 P2 收尾能力。

- transport = stdio，使用 LSP 标准的 **Content-Length 分帧** JSON-RPC
  （与 MCP 的逐行 JSON 不同！LSP 消息头是 `Content-Length: N\\r\\n\\r\\n` + N 字节 JSON）。
- 同步 + 线程安全：与 mcp_client.py 一致，契合小狸「同步主循环 + ThreadPoolExecutor」。
- 按文件扩展名路由到对应的语言服务器（python→pylsp、rust→rust-analyzer …），
  配置驱动，不绑定具体 server；未安装/启动失败不影响主程序与其他 LSP。
- 单个 server 失败不阻断其他，也不阻断主程序启动。

后续可在此文件扩展 workspace 级多文件夹、incremental 文本同步等。
"""

import os
import re
import json
import time
import shutil
import sys
import threading
import subprocess
import logging
from typing import Dict, List, Any, Optional, Tuple

logger = logging.getLogger(__name__)


class LspClientError(RuntimeError):
    """LSP 客户端错误"""


# 扩展名 → LSP languageId
_EXT_TO_LANG = {
    ".py": "python", ".pyi": "python",
    ".rs": "rust",
    ".ts": "typescript", ".tsx": "typescript",
    ".js": "javascript", ".jsx": "javascript", ".mjs": "javascript", ".cjs": "javascript",
    ".go": "go",
    ".c": "c", ".h": "c",
    ".cpp": "cpp", ".cc": "cpp", ".cxx": "cpp", ".hpp": "cpp",
    ".java": "java",
    ".lua": "lua",
    ".sh": "shellscript", ".bash": "shellscript",
    ".json": "json", ".jsonc": "json",
    ".html": "html", ".htm": "html",
    ".css": "css", ".scss": "scss", ".less": "less",
    ".md": "markdown", ".markdown": "markdown",
    ".vue": "vue",
    ".rb": "ruby",
    ".php": "php",
    ".sql": "sql",
    ".yaml": "yaml", ".yml": "yaml",
    ".xml": "xml",
    ".toml": "toml",
}


def ext_to_language(path: str) -> Optional[str]:
    """从文件路径推断 LSP languageId"""
    ext = os.path.splitext(path)[1].lower()
    return _EXT_TO_LANG.get(ext)


def path_to_uri(path: str) -> str:
    """绝对路径 → file:// URI（Windows 下 file:///C:/...）"""
    abspath = os.path.abspath(path)
    # 统一为正斜杠
    norm = abspath.replace("\\", "/")
    if not norm.startswith("/"):
        norm = "/" + norm
    return "file://" + norm


def _parse_position(pos: Any) -> Tuple[int, int]:
    """兼容 {line,character} / [line,char] / (line,char) → (line, character)（0-based）"""
    if isinstance(pos, dict):
        return int(pos.get("line", 0)), int(pos.get("character", 0))
    if isinstance(pos, (list, tuple)) and len(pos) >= 2:
        return int(pos[0]), int(pos[1])
    return 0, 0


class LspClient:
    """单个语言服务器的 stdio 客户端（同步、线程安全、Content-Length 分帧）"""

    def __init__(self, name: str, command: str, args: Optional[List[str]] = None,
                 env: Optional[Dict[str, str]] = None, cwd: Optional[str] = None,
                 root_uri: Optional[str] = None, timeout: float = 30.0,
                 language_id: str = ""):
        self.name = name
        self.command = command
        self.args = list(args or [])
        self.env = env or {}
        self.cwd = cwd
        self.root_uri = root_uri
        self.timeout = timeout
        self.language_id = language_id

        self.proc = None
        self._lock = threading.RLock()
        self._next_id = 1
        self._pending: Dict[int, Dict[str, Any]] = {}
        self._reader_thread = None
        self._started = False
        self._opened: set = set()
        self._diagnostics: Dict[str, List[Dict]] = {}  # uri -> diagnostics
        self.capabilities: Dict[str, Any] = {}
        self._diag_lock = threading.Lock()

    # ── 生命周期 ──
    def start(self) -> None:
        if self._started:
            return
        full_env = dict(os.environ)
        full_env.update(self.env)
        cmd = [self.command] + self.args
        # 鲁棒性: command 为裸 "python"/"python3" 但 PATH 找不到时, 回退当前解释器
        if cmd and os.path.basename(cmd[0]) in ("python", "python3") and os.path.dirname(cmd[0]) == "":
            found = None
            for cand in (cmd[0], cmd[0] + ".exe"):
                if found is None and shutil.which(cand):
                    found = cand
            if not found:
                cmd = [sys.executable] + self.args
        try:
            # 二进制模式：避免 CJK 字符数与 Content-Length 字节数不一致
            self.proc = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.DEVNULL,
                cwd=self.cwd,
                env=full_env,
                bufsize=0,
            )
        except Exception as e:
            raise LspClientError(f"启动 LSP server '{self.name}' 失败: {e}")

        self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
        self._reader_thread.start()
        self._started = True

    def _reader_loop(self):
        """后台读线程：按 Content-Length 分帧解析 JSON-RPC，响应按 id 投递；
        通知（publishDiagnostics 等）走回调。"""
        try:
            while True:
                msg = self._read_message()
                if msg is None:
                    break
                self._dispatch(msg)
        except Exception:
            pass
        finally:
            with self._lock:
                pendings = list(self._pending.values())
                self._pending.clear()
            for pend in pendings:
                pend["error"] = {"code": -1, "message": "LSP server 连接断开"}
                pend["event"].set()

    def _read_message(self) -> Optional[Dict[str, Any]]:
        """从 stdout 读取一条 Content-Length 分帧消息，返回解析后的 dict 或 None(EOF)"""
        # 读头部，直到遇到空行（\r\n\r\n）
        header = b""
        while b"\r\n\r\n" not in header:
            chunk = self.proc.stdout.read(1)
            if not chunk:
                return None
            header += chunk
            if len(header) > 65536:  # 防御：头部异常大
                return None
        # 解析 Content-Length
        m = re.search(rb"Content-Length:\s*(\d+)", header)
        if not m:
            return None
        length = int(m.group(1))
        body = b""
        while len(body) < length:
            chunk = self.proc.stdout.read(length - len(body))
            if not chunk:
                return None
            body += chunk
        try:
            return json.loads(body.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

    def _dispatch(self, msg: Dict[str, Any]) -> None:
        msg_id = msg.get("id")
        method = msg.get("method")
        if msg_id is not None and "method" not in msg:
            # 响应
            with self._lock:
                pend = self._pending.pop(msg_id, None)
            if pend is not None:
                if "error" in msg:
                    pend["error"] = msg["error"]
                else:
                    pend["value"] = msg.get("result")
                pend["event"].set()
            return
        # 通知（无 id 或 method 在通知里）
        if method == "textDocument/publishDiagnostics":
            params = msg.get("params", {})
            uri = params.get("uri", "")
            diags = params.get("diagnostics", [])
            with self._diag_lock:
                self._diagnostics[uri] = diags
        # 其余通知（window/logMessage 等）忽略

    def _send(self, method: str, params: Any = None,
              is_notification: bool = False) -> Optional[Any]:
        if not self._started or self.proc is None:
            raise LspClientError(f"LSP client '{self.name}' 未启动")
        with self._lock:
            pend = None
            req_id = None
            if not is_notification:
                req_id = self._next_id
                self._next_id += 1
                pend = {"event": threading.Event(), "value": None, "error": None}
                self._pending[req_id] = pend
            req = {"jsonrpc": "2.0", "method": method}
            if not is_notification:
                req["id"] = req_id
            if params is not None:
                req["params"] = params
        try:
            data = json.dumps(req, ensure_ascii=False).encode("utf-8")
            header = f"Content-Length: {len(data)}\r\n\r\n".encode("ascii")
            self.proc.stdin.write(header + data)
            self.proc.stdin.flush()
        except (BrokenPipeError, ValueError, OSError) as e:
            raise LspClientError(f"向 LSP server '{self.name}' 写入失败: {e}")
        if is_notification:
            return None
        if not pend["event"].wait(self.timeout):
            with self._lock:
                self._pending.pop(req_id, None)
            raise LspClientError(f"LSP server '{self.name}' 响应超时: {method}")
        if pend["error"]:
            err = pend["error"]
            raise LspClientError(
                f"LSP server '{self.name}' 错误 [{err.get('code')}]: {err.get('message')}")
        return pend["value"]

    # ── LSP 协议方法 ──
    def initialize(self, root_uri: Optional[str] = None) -> Dict[str, Any]:
        if root_uri:
            self.root_uri = root_uri
        result = self._send("initialize", {
            "processId": os.getpid(),
            "rootUri": self.root_uri,
            "capabilities": {
                "textDocument": {
                    "synchronization": {"dynamicRegistration": True,
                                          "didSave": True},
                    "publishDiagnostics": {"relatedInformation": True},
                    "hover": {"contentFormat": ["markdown", "plaintext"]},
                    "definition": {"linkSupport": True},
                    "references": {},
                    "completion": {"completionItem": {"snippetSupport": True}},
                },
                "workspace": {"workspaceFolders": True,
                               "configuration": True},
            },
            "clientInfo": {"name": "xiaoli-cli", "version": "8.0"},
        })
        if isinstance(result, dict):
            self.capabilities = result.get("capabilities", {}) or {}
        # 通知服务器初始化完成
        self._send("initialized", {}, is_notification=True)
        return result or {}

    def _ensure_open(self, uri: str, text: str) -> None:
        if uri in self._opened:
            # 已打开则同步最新内容（full 同步）。didChange 是通知，无需等待响应。
            self._send("textDocument/didChange", {
                "textDocument": {"uri": uri},
                "contentChanges": [{"text": text}],
            }, is_notification=True)
            return
        lang = self.language_id or "plaintext"
        # didOpen 是通知，无需等待响应
        self._send("textDocument/didOpen", {
            "textDocument": {
                "uri": uri,
                "languageId": lang,
                "version": 1,
                "text": text,
            }
        }, is_notification=True)
        self._opened.add(uri)

    def _read_file(self, file_path: str) -> str:
        try:
            with open(file_path, "r", encoding="utf-8", errors="replace") as f:
                return f.read()
        except Exception as e:
            return f"# 无法读取文件 {file_path}: {e}"

    def diagnostics(self, file_path: str, wait: float = 2.0) -> List[Dict]:
        """返回某文件的诊断（先 didOpen 触发分析，等待并收集 publishDiagnostics）"""
        uri = path_to_uri(file_path)
        text = self._read_file(file_path)
        self._ensure_open(uri, text)
        # 等待服务器推送诊断（最多 wait 秒）
        deadline = time.time() + wait
        while time.time() < deadline:
            with self._diag_lock:
                if uri in self._diagnostics:
                    return list(self._diagnostics[uri])
            time.sleep(0.1)
        with self._diag_lock:
            return list(self._diagnostics.get(uri, []))

    def hover(self, file_path: str, line: int, character: int) -> Any:
        uri = path_to_uri(file_path)
        text = self._read_file(file_path)
        self._ensure_open(uri, text)
        return self._send("textDocument/hover", {
            "textDocument": {"uri": uri},
            "position": {"line": line, "character": character},
        })

    def definition(self, file_path: str, line: int, character: int) -> Any:
        uri = path_to_uri(file_path)
        text = self._read_file(file_path)
        self._ensure_open(uri, text)
        return self._send("textDocument/definition", {
            "textDocument": {"uri": uri},
            "position": {"line": line, "character": character},
        })

    def references(self, file_path: str, line: int, character: int) -> Any:
        uri = path_to_uri(file_path)
        text = self._read_file(file_path)
        self._ensure_open(uri, text)
        return self._send("textDocument/references", {
            "textDocument": {"uri": uri},
            "position": {"line": line, "character": character},
            "context": {"includeDeclaration": True},
        })

    def completion(self, file_path: str, line: int, character: int) -> Any:
        uri = path_to_uri(file_path)
        text = self._read_file(file_path)
        self._ensure_open(uri, text)
        return self._send("textDocument/completion", {
            "textDocument": {"uri": uri},
            "position": {"line": line, "character": character},
        })

    def has_diagnostics(self, file_path: str) -> bool:
        uri = path_to_uri(file_path)
        with self._diag_lock:
            return bool(self._diagnostics.get(uri))

    def shutdown(self) -> None:
        if self.proc is None:
            return
        try:
            self._send("shutdown", {})
        except Exception:
            pass
        try:
            self._send("exit", {}, is_notification=True)
        except Exception:
            pass
        try:
            self.proc.stdin.close()
        except Exception:
            pass
        try:
            self.proc.wait(timeout=3)
        except Exception:
            try:
                self.proc.kill()
            except Exception:
                pass
        self._started = False


class LspManager:
    """管理多个语言服务器，按文件扩展名路由诊断/智能感知请求"""

    def __init__(self):
        self.clients: Dict[str, LspClient] = {}      # language_id -> client
        self._ext_index: Dict[str, str] = {}         # ext -> language_id
        self._lock = threading.RLock()
        self._cleaned = False
        self._workspace_diag: Dict[str, List[Dict]] = {}  # file -> diagnostics（缓存）

    def connect_all(self, servers: Dict[str, Dict[str, Any]],
                    cwd: Optional[str] = None,
                    root_uri: Optional[str] = None) -> None:
        """连接所有已声明且未禁用的语言服务器。

        servers: {language_id 或扩展名: {"command", "args", "env", "cwd", "disabled", "languageId"}}
        键若是扩展名（以 '.' 开头）则直接登记该扩展名；否则视为 languageId，
        并把其常见扩展名都映射到它。
        """
        for key, spec in servers.items():
            if not isinstance(spec, dict):
                continue
            if spec.get("disabled"):
                logger.info(f"LSP server '{key}' 已禁用，跳过")
                continue
            command = spec.get("command")
            if not command:
                logger.warning(f"LSP server '{key}' 缺少 command，跳过")
                continue
            args = spec.get("args") or []
            env = spec.get("env") or {}
            server_cwd = spec.get("cwd") or cwd
            explicit_lang = spec.get("languageId") or ""
            if key.startswith("."):
                lang_id = explicit_lang or (ext_to_language(key) or key[1:])
                self._ext_index[key.lower()] = lang_id
            else:
                lang_id = key
                for ext, lid in _EXT_TO_LANG.items():
                    if lid == lang_id:
                        self._ext_index[ext] = lang_id
            if lang_id in self.clients:
                continue  # 同一语言只起一个 server
            client = LspClient(lang_id, command, args, env, server_cwd,
                               root_uri=root_uri, language_id=lang_id)
            try:
                client.start()
                client.initialize()
                with self._lock:
                    self.clients[lang_id] = client
                logger.info(f"已连接 LSP server '{lang_id}' ({command})")
            except Exception as e:
                try:
                    client.shutdown()
                except Exception:
                    pass
                logger.error(f"连接 LSP server '{lang_id}' 失败: {e}")

    def _resolve_client(self, file_path: str) -> Optional[LspClient]:
        ext = os.path.splitext(file_path)[1].lower()
        with self._lock:
            lang_id = self._ext_index.get(ext)
            if lang_id:
                return self.clients.get(lang_id)
        # 退化：按 languageId 推断
        lid = ext_to_language(file_path)
        if lid:
            with self._lock:
                return self.clients.get(lid)
        return None

    def _format_diagnostics(self, file_path: str, diags: List[Dict]) -> str:
        if not diags:
            return f"✅ {file_path}：无诊断（未发现问题）"
        lines = [f"🔍 {file_path}：{len(diags)} 条诊断"]
        for d in diags:
            rng = d.get("range", {}).get("start", {})
            ln = rng.get("line", "?")
            ch = rng.get("character", "?")
            sev = d.get("severity", 1)
            sev_map = {1: "错误", 2: "警告", 3: "信息", 4: "提示"}
            label = sev_map.get(sev, "问题")
            msg = d.get("message", "")
            src = d.get("source", "")
            src_tag = f" [{src}]" if src else ""
            lines.append(f"  {label}{src_tag} L{ln + 1}:{ch + 1} — {msg}")
        return "\n".join(lines)

    def diagnostics(self, file_path: str, wait: float = 2.0) -> str:
        client = self._resolve_client(file_path)
        if client is None:
            return (f"⚠️ 没有为 {file_path} 配置 LSP server"
                    f"（在 config.json 的 lspServers 中添加对应语言）")
        diags = client.diagnostics(file_path, wait)
        text = self._format_diagnostics(file_path, diags)
        # 缓存，供系统提示自动注入
        with self._lock:
            self._workspace_diag[file_path] = diags
        return text

    def hover(self, file_path: str, line: int, character: int) -> str:
        client = self._resolve_client(file_path)
        if client is None:
            return f"⚠️ 没有为 {file_path} 配置 LSP server"
        try:
            res = client.hover(file_path, line, character)
        except Exception as e:
            return f"❌ hover 失败: {e}"
        return self._format_hover(res)

    def definition(self, file_path: str, line: int, character: int) -> str:
        client = self._resolve_client(file_path)
        if client is None:
            return f"⚠️ 没有为 {file_path} 配置 LSP server"
        try:
            res = client.definition(file_path, line, character)
        except Exception as e:
            return f"❌ definition 失败: {e}"
        return self._format_locations(res, "定义")

    def references(self, file_path: str, line: int, character: int) -> str:
        client = self._resolve_client(file_path)
        if client is None:
            return f"⚠️ 没有为 {file_path} 配置 LSP server"
        try:
            res = client.references(file_path, line, character)
        except Exception as e:
            return f"❌ references 失败: {e}"
        return self._format_locations(res, "引用")

    def completion(self, file_path: str, line: int, character: int) -> str:
        client = self._resolve_client(file_path)
        if client is None:
            return f"⚠️ 没有为 {file_path} 配置 LSP server"
        try:
            res = client.completion(file_path, line, character)
        except Exception as e:
            return f"❌ completion 失败: {e}"
        return self._format_completion(res)

    @staticmethod
    def _format_hover(res: Any) -> str:
        if not res:
            return "（无 hover 信息）"
        contents = res.get("contents") if isinstance(res, dict) else None
        if isinstance(contents, str):
            return contents
        if isinstance(contents, dict):
            return contents.get("value", "") or json.dumps(contents, ensure_ascii=False)
        if isinstance(contents, list):
            parts = []
            for c in contents:
                if isinstance(c, str):
                    parts.append(c)
                elif isinstance(c, dict):
                    parts.append(c.get("value", ""))
            return "\n".join(p for p in parts if p)
        return json.dumps(res, ensure_ascii=False)

    @staticmethod
    def _format_locations(res: Any, label: str) -> str:
        if not res:
            return f"（无{label}）"
        items = res if isinstance(res, list) else [res]
        out = [f"{label}位置（{len(items)}）："]
        for it in items:
            loc = it.get("location", it) if isinstance(it, dict) else {}
            uri = loc.get("uri", "")
            rng = loc.get("range", {}).get("start", {})
            ln = rng.get("line", "?")
            ch = rng.get("character", "?")
            fp = uri.replace("file://", "") if uri.startswith("file://") else uri
            out.append(f"  {fp} L{ln + 1}:{ch + 1}")
        return "\n".join(out)

    @staticmethod
    def _format_completion(res: Any) -> str:
        if not res:
            return "（无补全项）"
        if isinstance(res, dict):
            items = res.get("items", [])
        elif isinstance(res, list):
            items = res
        else:
            items = []
        if not items:
            return "（无补全项）"
        out = [f"补全项（{len(items)}）："]
        for it in items[:30]:
            label = it.get("label", "")
            kind = it.get("kind", "")
            detail = it.get("detail", "")
            txt = f"  {label}"
            if detail:
                txt += f" — {detail}"
            elif kind:
                txt += f" (kind={kind})"
            out.append(txt)
        return "\n".join(out)

    def get_context(self) -> str:
        """返回所有已分析文件的诊断汇总，用于注入系统提示"""
        with self._lock:
            entries = list(self._workspace_diag.items())
        if not entries:
            return ""
        blocks = []
        for fp, diags in entries:
            if diags:
                blocks.append(self._format_diagnostics(fp, diags))
        return "\n\n".join(blocks)

    def get_servers(self) -> List[str]:
        with self._lock:
            return list(self.clients.keys())

    def cleanup(self) -> None:
        with self._lock:
            if self._cleaned:
                return
            self._cleaned = True
            clients = list(self.clients.values())
            self.clients.clear()
            self._ext_index.clear()
            self._workspace_diag.clear()
        for client in clients:
            try:
                client.shutdown()
            except Exception:
                pass