"""
LSP 客户端实现

使用 asyncio.subprocess 与语言服务器进行 JSON-RPC 通信。
"""

import asyncio
import json
import logging
import os
from typing import Any, Callable, Dict, List, Optional, Union

from .protocol import (
    Position,
    Range,
    Location,
    LocationLink,
    DocumentSymbol,
    SymbolInformation,
    Hover,
    Diagnostic,
    CallHierarchyItem,
    CallHierarchyIncomingCall,
    CallHierarchyOutgoingCall,
    JsonRpcRequest,
    JsonRpcResponse,
    JsonRpcNotification,
    TextDocumentIdentifier,
    TextDocumentPositionParams,
)

logger = logging.getLogger(__name__)


class LSPClient:
    """
    LSP 客户端，管理与语言服务器的通信

    使用 JSON-RPC 2.0 协议与语言服务器进行通信。
    """

    def __init__(self, server_id: str, command: List[str], env: Optional[Dict[str, str]] = None):
        """
        初始化 LSP 客户端

        Args:
            server_id: 服务器标识
            command: 启动服务器的命令
            env: 环境变量
        """
        self.server_id = server_id
        self.command = command
        self.env = env

        self._process: Optional[asyncio.subprocess.Process] = None
        self._request_id = 0
        self._pending_requests: Dict[int, asyncio.Future] = {}
        self._initialized = False
        self._root_path: Optional[str] = None
        self._reader_task: Optional[asyncio.Task] = None

        # 诊断回调
        self._diagnostics: Dict[str, List[Diagnostic]] = {}
        self._on_diagnostics: Optional[Callable[[str, List[Diagnostic]], None]] = None

        # 已打开的文件
        self._open_files: set = set()

    async def initialize(self, root_path: str) -> bool:
        """
        初始化 LSP 连接

        Args:
            root_path: 项目根目录路径

        Returns:
            是否初始化成功
        """
        if self._initialized:
            return True

        self._root_path = root_path

        try:
            # 启动语言服务器进程
            env = os.environ.copy()
            if self.env:
                env.update(self.env)

            self._process = await asyncio.create_subprocess_exec(
                *self.command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
                cwd=root_path
            )

            # 启动消息读取任务
            self._reader_task = asyncio.create_task(self._read_messages())

            # 发送 initialize 请求
            root_uri = f"file://{root_path}"
            init_params = {
                "processId": os.getpid(),
                "rootPath": root_path,
                "rootUri": root_uri,
                "capabilities": {
                    "textDocument": {
                        "definition": {"dynamicRegistration": False, "linkSupport": True},
                        "references": {"dynamicRegistration": False},
                        "hover": {"dynamicRegistration": False, "contentFormat": ["markdown", "plaintext"]},
                        "documentSymbol": {"dynamicRegistration": False, "hierarchicalDocumentSymbolSupport": True},
                        "callHierarchy": {"dynamicRegistration": False},
                        "implementation": {"dynamicRegistration": False, "linkSupport": True},
                        "publishDiagnostics": {"relatedInformation": True},
                    },
                    "workspace": {
                        "symbol": {"dynamicRegistration": False},
                        "workspaceFolders": True,
                    }
                },
                "workspaceFolders": [{"uri": root_uri, "name": os.path.basename(root_path)}]
            }

            result = await self._send_request("initialize", init_params)
            if result is None:
                logger.error(f"[{self.server_id}] Initialize failed: no result")
                return False

            # 发送 initialized 通知
            await self._send_notification("initialized", {})
            self._initialized = True
            logger.info(f"[{self.server_id}] LSP client initialized for {root_path}")
            return True

        except Exception as e:
            logger.error(f"[{self.server_id}] Failed to initialize: {e}")
            await self.shutdown()
            return False

    async def shutdown(self) -> None:
        """关闭 LSP 连接"""
        if self._reader_task:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass

        if self._process and self._process.returncode is None:
            try:
                await self._send_request("shutdown", None, timeout=2.0)
                await self._send_notification("exit", None)
            except Exception:
                pass

            try:
                self._process.terminate()
                await asyncio.wait_for(self._process.wait(), timeout=2.0)
            except asyncio.TimeoutError:
                self._process.kill()
            except Exception:
                pass

        self._initialized = False
        self._open_files.clear()
        self._pending_requests.clear()
        logger.info(f"[{self.server_id}] LSP client shutdown")

    async def _send_request(
        self,
        method: str,
        params: Optional[Any],
        timeout: float = 60.0
    ) -> Optional[Any]:
        """
        发送 JSON-RPC 请求并等待响应

        Args:
            method: 方法名
            params: 参数
            timeout: 超时时间（秒）

        Returns:
            响应结果
        """
        if not self._process or not self._process.stdin:
            logger.error(f"[{self.server_id}] Process not running")
            return None

        self._request_id += 1
        request_id = self._request_id

        request = JsonRpcRequest(id=request_id, method=method, params=params)
        future: asyncio.Future = asyncio.get_event_loop().create_future()
        self._pending_requests[request_id] = future

        try:
            await self._write_message(request.to_dict())
            result = await asyncio.wait_for(future, timeout=timeout)
            return result
        except asyncio.TimeoutError:
            logger.error(f"[{self.server_id}] Request {method} timed out")
            return None
        except Exception as e:
            logger.error(f"[{self.server_id}] Request {method} failed: {e}")
            return None
        finally:
            self._pending_requests.pop(request_id, None)

    async def _send_notification(self, method: str, params: Optional[Any]) -> None:
        """发送 JSON-RPC 通知（不需要响应）"""
        if not self._process or not self._process.stdin:
            return

        notification = JsonRpcNotification(method=method, params=params)
        await self._write_message(notification.to_dict())

    async def _write_message(self, message: Dict[str, Any]) -> None:
        """写入 LSP 消息"""
        if not self._process or not self._process.stdin:
            return

        content = json.dumps(message)
        content_bytes = content.encode("utf-8")
        header = f"Content-Length: {len(content_bytes)}\r\n\r\n"

        self._process.stdin.write(header.encode("utf-8"))
        self._process.stdin.write(content_bytes)
        await self._process.stdin.drain()

    async def _read_messages(self) -> None:
        """持续读取服务器消息"""
        if not self._process or not self._process.stdout:
            return

        try:
            while True:
                # 读取 header
                headers = {}
                while True:
                    line = await self._process.stdout.readline()
                    if not line:
                        return
                    line_str = line.decode("utf-8").strip()
                    if not line_str:
                        break
                    if ":" in line_str:
                        key, value = line_str.split(":", 1)
                        headers[key.strip().lower()] = value.strip()

                # 读取 content
                content_length = int(headers.get("content-length", 0))
                if content_length == 0:
                    continue

                content = await self._process.stdout.readexactly(content_length)
                message = json.loads(content.decode("utf-8"))
                await self._handle_message(message)

        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"[{self.server_id}] Error reading messages: {e}")

    async def _handle_message(self, message: Dict[str, Any]) -> None:
        """处理收到的消息"""
        if "id" in message and "method" not in message:
            # 响应消息
            response = JsonRpcResponse.from_dict(message)
            future = self._pending_requests.get(response.id)
            if future and not future.done():
                if response.error:
                    logger.error(f"[{self.server_id}] RPC error: {response.error}")
                    future.set_result(None)
                else:
                    future.set_result(response.result)
        elif "method" in message:
            # 通知或请求
            method = message.get("method", "")
            params = message.get("params")

            if method == "textDocument/publishDiagnostics":
                self._handle_diagnostics(params)
            elif method == "window/logMessage":
                level = params.get("type", 3)
                msg = params.get("message", "")
                if level <= 2:  # Error or Warning
                    logger.warning(f"[{self.server_id}] {msg}")

    def _handle_diagnostics(self, params: Dict[str, Any]) -> None:
        """处理诊断信息"""
        uri = params.get("uri", "")
        diagnostics_data = params.get("diagnostics", [])
        diagnostics = [Diagnostic.from_dict(d) for d in diagnostics_data]
        self._diagnostics[uri] = diagnostics

        if self._on_diagnostics:
            self._on_diagnostics(uri, diagnostics)

    def on_diagnostics(self, callback: Callable[[str, List[Diagnostic]], None]) -> None:
        """注册诊断回调"""
        self._on_diagnostics = callback

    def get_diagnostics(self, file_path: str) -> List[Diagnostic]:
        """获取文件的诊断信息"""
        uri = f"file://{file_path}"
        return self._diagnostics.get(uri, [])

    async def open_file(self, file_path: str) -> None:
        """打开文件（通知服务器）"""
        if file_path in self._open_files:
            return

        uri = f"file://{file_path}"
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                content = f.read()
        except Exception as e:
            logger.error(f"[{self.server_id}] Failed to read file {file_path}: {e}")
            return

        # 根据扩展名猜测语言 ID
        ext = os.path.splitext(file_path)[1].lower()
        language_ids = {
            ".py": "python",
            ".pyi": "python",
            ".js": "javascript",
            ".jsx": "javascriptreact",
            ".ts": "typescript",
            ".tsx": "typescriptreact",
            ".go": "go",
            ".rs": "rust",
            ".java": "java",
            ".c": "c",
            ".cpp": "cpp",
            ".h": "c",
            ".hpp": "cpp",
            ".rb": "ruby",
            ".php": "php",
            ".swift": "swift",
        }
        language_id = language_ids.get(ext, "plaintext")

        await self._send_notification("textDocument/didOpen", {
            "textDocument": {
                "uri": uri,
                "languageId": language_id,
                "version": 1,
                "text": content
            }
        })
        self._open_files.add(file_path)

    async def close_file(self, file_path: str) -> None:
        """关闭文件（通知服务器）"""
        if file_path not in self._open_files:
            return

        uri = f"file://{file_path}"
        await self._send_notification("textDocument/didClose", {
            "textDocument": {"uri": uri}
        })
        self._open_files.discard(file_path)

    async def text_document_definition(
        self,
        file_path: str,
        line: int,
        character: int
    ) -> List[Location]:
        """
        获取定义位置

        Args:
            file_path: 文件路径
            line: 行号（1-based，会转换为 0-based）
            character: 字符偏移（1-based，会转换为 0-based）

        Returns:
            定义位置列表
        """
        await self.open_file(file_path)

        params = TextDocumentPositionParams(
            textDocument=TextDocumentIdentifier(uri=f"file://{file_path}"),
            position=Position(line=line - 1, character=character - 1)
        )

        result = await self._send_request("textDocument/definition", params.to_dict())
        return self._parse_locations(result)

    async def text_document_references(
        self,
        file_path: str,
        line: int,
        character: int,
        include_declaration: bool = True
    ) -> List[Location]:
        """
        获取引用位置

        Args:
            file_path: 文件路径
            line: 行号（1-based）
            character: 字符偏移（1-based）
            include_declaration: 是否包含声明

        Returns:
            引用位置列表
        """
        await self.open_file(file_path)

        params = {
            "textDocument": {"uri": f"file://{file_path}"},
            "position": {"line": line - 1, "character": character - 1},
            "context": {"includeDeclaration": include_declaration}
        }

        result = await self._send_request("textDocument/references", params)
        return self._parse_locations(result)

    async def text_document_hover(
        self,
        file_path: str,
        line: int,
        character: int
    ) -> Optional[Hover]:
        """
        获取悬停信息

        Args:
            file_path: 文件路径
            line: 行号（1-based）
            character: 字符偏移（1-based）

        Returns:
            悬停信息
        """
        await self.open_file(file_path)

        params = TextDocumentPositionParams(
            textDocument=TextDocumentIdentifier(uri=f"file://{file_path}"),
            position=Position(line=line - 1, character=character - 1)
        )

        result = await self._send_request("textDocument/hover", params.to_dict())
        if result:
            return Hover.from_dict(result)
        return None

    async def document_symbol(self, file_path: str) -> List[DocumentSymbol]:
        """
        获取文档符号

        Args:
            file_path: 文件路径

        Returns:
            文档符号列表
        """
        await self.open_file(file_path)

        params = {"textDocument": {"uri": f"file://{file_path}"}}
        result = await self._send_request("textDocument/documentSymbol", params)

        if not result:
            return []

        # 结果可能是 DocumentSymbol[] 或 SymbolInformation[]
        if result and isinstance(result[0], dict):
            if "range" in result[0] and "selectionRange" in result[0]:
                return [DocumentSymbol.from_dict(s) for s in result]
            elif "location" in result[0]:
                # SymbolInformation 转换为 DocumentSymbol
                symbols = []
                for s in result:
                    info = SymbolInformation.from_dict(s)
                    symbols.append(DocumentSymbol(
                        name=info.name,
                        kind=info.kind,
                        range=info.location.range,
                        selectionRange=info.location.range
                    ))
                return symbols

        return []

    async def workspace_symbol(self, query: str) -> List[SymbolInformation]:
        """
        搜索工作区符号

        Args:
            query: 搜索查询

        Returns:
            符号列表
        """
        params = {"query": query}
        result = await self._send_request("workspace/symbol", params)

        if not result:
            return []

        return [SymbolInformation.from_dict(s) for s in result]

    async def text_document_implementation(
        self,
        file_path: str,
        line: int,
        character: int
    ) -> List[Location]:
        """
        获取实现位置

        Args:
            file_path: 文件路径
            line: 行号（1-based）
            character: 字符偏移（1-based）

        Returns:
            实现位置列表
        """
        await self.open_file(file_path)

        params = TextDocumentPositionParams(
            textDocument=TextDocumentIdentifier(uri=f"file://{file_path}"),
            position=Position(line=line - 1, character=character - 1)
        )

        result = await self._send_request("textDocument/implementation", params.to_dict())
        return self._parse_locations(result)

    async def prepare_call_hierarchy(
        self,
        file_path: str,
        line: int,
        character: int
    ) -> List[CallHierarchyItem]:
        """
        准备调用层次

        Args:
            file_path: 文件路径
            line: 行号（1-based）
            character: 字符偏移（1-based）

        Returns:
            调用层次项列表
        """
        await self.open_file(file_path)

        params = TextDocumentPositionParams(
            textDocument=TextDocumentIdentifier(uri=f"file://{file_path}"),
            position=Position(line=line - 1, character=character - 1)
        )

        result = await self._send_request("textDocument/prepareCallHierarchy", params.to_dict())

        if not result:
            return []

        return [CallHierarchyItem.from_dict(item) for item in result]

    async def incoming_calls(self, item: CallHierarchyItem) -> List[CallHierarchyIncomingCall]:
        """
        获取传入调用

        Args:
            item: 调用层次项

        Returns:
            传入调用列表
        """
        params = {"item": item.to_dict()}
        result = await self._send_request("callHierarchy/incomingCalls", params)

        if not result:
            return []

        return [CallHierarchyIncomingCall.from_dict(call) for call in result]

    async def outgoing_calls(self, item: CallHierarchyItem) -> List[CallHierarchyOutgoingCall]:
        """
        获取传出调用

        Args:
            item: 调用层次项

        Returns:
            传出调用列表
        """
        params = {"item": item.to_dict()}
        result = await self._send_request("callHierarchy/outgoingCalls", params)

        if not result:
            return []

        return [CallHierarchyOutgoingCall.from_dict(call) for call in result]

    def _parse_locations(self, result: Optional[Any]) -> List[Location]:
        """解析位置结果（可能是 Location, Location[], LocationLink[]）"""
        if not result:
            return []

        if isinstance(result, dict):
            # 单个 Location 或 LocationLink
            if "targetUri" in result:
                link = LocationLink.from_dict(result)
                return [Location(uri=link.targetUri, range=link.targetSelectionRange)]
            else:
                return [Location.from_dict(result)]

        if isinstance(result, list):
            locations = []
            for item in result:
                if "targetUri" in item:
                    link = LocationLink.from_dict(item)
                    locations.append(Location(uri=link.targetUri, range=link.targetSelectionRange))
                else:
                    locations.append(Location.from_dict(item))
            return locations

        return []
