"""Async JSON-RPC stdio client for LSP servers.

Manages a single server process, sending requests/notifications and
collecting diagnostics + document symbols via the standard LSP wire protocol.
"""

import asyncio
import json
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from otto.log import get_logger

log = get_logger(__name__)

_HEADER_ENCODING = "ascii"
_CONTENT_ENCODING = "utf-8"


@dataclass(frozen=True)
class Diagnostic:
    """A single diagnostic from an LSP server."""

    path: str
    line: int
    character: int
    message: str
    severity: int  # 1=Error, 2=Warning, 3=Info, 4=Hint
    source: str = ""


@dataclass(frozen=True)
class DocumentSymbol:
    """A document symbol extracted via LSP."""

    name: str
    kind: int  # LSP SymbolKind enum
    file: str
    line: int
    signature: str = ""
    parent: str | None = None


# LSP SymbolKind → human-readable name
_SYMBOL_KIND_NAMES: dict[int, str] = {
    1: "file",
    2: "module",
    3: "namespace",
    4: "package",
    5: "class",
    6: "method",
    7: "property",
    8: "field",
    9: "constructor",
    10: "enum",
    11: "interface",
    12: "function",
    13: "variable",
    14: "constant",
    15: "string",
    16: "number",
    17: "boolean",
    18: "array",
    19: "object",
    20: "key",
    21: "null",
    22: "enum_member",
    23: "struct",
    24: "event",
    25: "operator",
    26: "type_parameter",
}


def symbol_kind_name(kind: int) -> str:
    """Return a human-readable name for an LSP SymbolKind."""
    return _SYMBOL_KIND_NAMES.get(kind, "unknown")


class LSPClient:
    """Manages one LSP server subprocess over JSON-RPC stdio."""

    def __init__(self, server_id: str, command: tuple[str, ...], root: Path) -> None:
        self.server_id = server_id
        self.command = command
        self.root = root
        self.last_used: float = time.monotonic()
        self._process: asyncio.subprocess.Process | None = None
        self._request_id = 0
        self._pending: dict[int, asyncio.Future] = {}
        self._diagnostics: dict[str, list[Diagnostic]] = {}
        self._diag_events: dict[str, asyncio.Event] = {}
        self._reader_task: asyncio.Task | None = None
        self._broken = False
        self._initialized = False
        self._open_files: set[str] = set()
        self._file_versions: dict[str, int] = {}
        self._capabilities: dict[str, Any] = {}
        self._server_requests: deque[tuple[int | str, str]] = deque()

    @property
    def broken(self) -> bool:
        return self._broken

    async def initialize(self) -> None:
        """Spawn the server process and send initialize + initialized."""
        if self._broken:
            return
        try:
            self._process = await asyncio.create_subprocess_exec(
                *self.command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        except (FileNotFoundError, PermissionError, OSError) as exc:
            log.info("failed to start %s: %s", self.server_id, exc)
            self._broken = True
            return

        self._reader_task = asyncio.create_task(self._reader_loop())
        self._stderr_task = asyncio.create_task(self._stderr_reader())

        try:
            result = await self._request(
                "initialize",
                {
                    "processId": None,
                    "rootUri": self.root.as_uri(),
                    "rootPath": str(self.root),
                    "capabilities": {
                        "textDocument": {
                            "publishDiagnostics": {"relatedInformation": False},
                            "documentSymbol": {
                                "hierarchicalDocumentSymbolSupport": True,
                            },
                            "synchronization": {
                                "dynamicRegistration": False,
                                "didSave": True,
                            },
                        },
                    },
                    "workspaceFolders": [{"uri": self.root.as_uri(), "name": self.root.name}],
                },
                timeout=15.0,
            )
            if result and isinstance(result, dict):
                self._capabilities = result.get("capabilities", {})
            await self._notify("initialized", {})
            self._initialized = True
        except Exception as exc:
            log.info("initialize failed for %s: %s", self.server_id, exc)
            self._broken = True
            await self._kill()

    async def shutdown(self) -> None:
        """Send shutdown + exit and clean up."""
        if self._process is None:
            return
        try:
            if not self._broken and self._initialized:
                await self._request("shutdown", None, timeout=5.0)
                await self._notify("exit", None)
        except Exception:
            pass
        await self._kill()

    async def did_open(self, path: Path, text: str, lang_id: str) -> None:
        """Notify the server about a newly opened file."""
        self.last_used = time.monotonic()
        if self._broken or not self._initialized:
            return
        uri = path.as_uri()
        key = str(path)
        if key in self._open_files:
            await self.did_change(path, text)
            return
        self._open_files.add(key)
        self._file_versions[key] = 1
        log.debug("did_open: %s lang=%s", uri, lang_id)
        await self._notify(
            "textDocument/didOpen",
            {
                "textDocument": {
                    "uri": uri,
                    "languageId": lang_id,
                    "version": 1,
                    "text": text,
                }
            },
        )

    async def did_change(self, path: Path, text: str) -> None:
        """Notify the server about file content changes."""
        self.last_used = time.monotonic()
        if self._broken or not self._initialized:
            return
        key = str(path)
        version = self._file_versions.get(key, 0) + 1
        self._file_versions[key] = version
        log.debug("did_change: %s v=%d", path.as_uri(), version)
        await self._notify(
            "textDocument/didChange",
            {
                "textDocument": {"uri": path.as_uri(), "version": version},
                "contentChanges": [{"text": text}],
            },
        )

    async def get_diagnostics(self, path: Path, *, timeout: float = 3.0) -> list[Diagnostic]:
        """Wait for publishDiagnostics for the given path.

        Uses a 150ms debounce — waits briefly for the server to finish
        sending diagnostics after a change, then collects.
        """
        self.last_used = time.monotonic()
        if self._broken or not self._initialized:
            return []
        uri = path.as_uri()
        event = self._diag_events.get(uri)
        if event is None:
            event = asyncio.Event()
            self._diag_events[uri] = event

        # Check if diagnostics already arrived (from background analysis)
        existing = self._diagnostics.get(uri)
        if existing:
            return existing

        event.clear()
        try:
            await asyncio.wait_for(event.wait(), timeout=timeout)
            # Debounce: wait a bit more for additional diagnostics
            await asyncio.sleep(0.15)
        except TimeoutError:
            pass

        return self._diagnostics.get(uri, [])

    async def document_symbols(self, path: Path) -> list[DocumentSymbol]:
        """Request textDocument/documentSymbol."""
        self.last_used = time.monotonic()
        if self._broken or not self._initialized:
            return []
        try:
            result = await self._request(
                "textDocument/documentSymbol",
                {"textDocument": {"uri": path.as_uri()}},
                timeout=10.0,
            )
        except Exception:
            return []

        if not result or not isinstance(result, list):
            return []

        rel = str(path)
        try:
            rel = str(path.relative_to(self.root))
        except ValueError:
            pass
        return _parse_document_symbols(result, rel)

    # ------------------------------------------------------------------
    # JSON-RPC wire protocol
    # ------------------------------------------------------------------

    def _next_id(self) -> int:
        self._request_id += 1
        return self._request_id

    async def _request(self, method: str, params: Any, *, timeout: float = 10.0) -> Any:
        """Send a JSON-RPC request and wait for the response."""
        if self._process is None or self._process.stdin is None:
            raise RuntimeError("not connected")

        req_id = self._next_id()
        msg = {"jsonrpc": "2.0", "id": req_id, "method": method}
        if params is not None:
            msg["params"] = params

        future: asyncio.Future = asyncio.get_running_loop().create_future()
        self._pending[req_id] = future

        await self._send(msg)
        try:
            return await asyncio.wait_for(future, timeout=timeout)
        except TimeoutError:
            self._pending.pop(req_id, None)
            raise

    async def _notify(self, method: str, params: Any) -> None:
        """Send a JSON-RPC notification (no response expected)."""
        if self._process is None or self._process.stdin is None:
            return
        msg: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            msg["params"] = params
        await self._send(msg)

    async def _send(self, msg: dict) -> None:
        """Encode and send a JSON-RPC message with Content-Length header."""
        if self._process is None or self._process.stdin is None:
            return
        body = json.dumps(msg).encode(_CONTENT_ENCODING)
        header = f"Content-Length: {len(body)}\r\n\r\n".encode(_HEADER_ENCODING)
        try:
            self._process.stdin.write(header + body)
            await self._process.stdin.drain()
        except (BrokenPipeError, ConnectionResetError, OSError):
            self._broken = True

    async def _reader_loop(self) -> None:
        """Read Content-Length framed JSON-RPC messages from stdout."""
        assert self._process is not None and self._process.stdout is not None
        reader = self._process.stdout
        try:
            while True:
                # Read headers until empty line
                content_length = 0
                while True:
                    line = await reader.readline()
                    if not line:
                        return  # EOF
                    line_str = line.decode(_HEADER_ENCODING).strip()
                    if not line_str:
                        break  # End of headers
                    if line_str.lower().startswith("content-length:"):
                        content_length = int(line_str.split(":", 1)[1].strip())

                if content_length <= 0:
                    continue

                body = await reader.readexactly(content_length)
                try:
                    msg = json.loads(body.decode(_CONTENT_ENCODING))
                except (json.JSONDecodeError, UnicodeDecodeError):
                    continue

                self._handle_message(msg)
                # Respond to any server→client requests
                while self._server_requests:
                    req_id, method = self._server_requests.popleft()
                    log.debug("→ %s respond to %s (id=%s)", self.server_id, method, req_id)
                    await self._send({"jsonrpc": "2.0", "id": req_id, "result": None})
        except (asyncio.IncompleteReadError, ConnectionResetError, OSError):
            pass
        except asyncio.CancelledError:
            pass
        finally:
            self._broken = True
            # Resolve any pending futures
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(ConnectionError("server disconnected"))
            self._pending.clear()

    async def _stderr_reader(self) -> None:
        """Read and log stderr from the server process."""
        if self._process is None or self._process.stderr is None:
            return
        try:
            while True:
                line = await self._process.stderr.readline()
                if not line:
                    break
                text = line.decode("utf-8", errors="replace").rstrip()
                if text:
                    log.debug("stderr[%s]: %s", self.server_id, text[:200])
        except (asyncio.CancelledError, OSError):
            pass

    def _handle_message(self, msg: dict) -> None:
        """Dispatch an incoming JSON-RPC message."""
        method = msg.get("method", "")
        if method:
            log.debug("← %s %s", self.server_id, method)
        if "id" in msg and "method" not in msg:
            # Response to a request we sent
            req_id = msg["id"]
            future = self._pending.pop(req_id, None)
            if future and not future.done():
                if "error" in msg:
                    future.set_exception(
                        RuntimeError(f"LSP error: {msg['error'].get('message', msg['error'])}")
                    )
                else:
                    future.set_result(msg.get("result"))
        elif "id" in msg and "method" in msg:
            # Server→client request — respond with null to avoid blocking
            self._server_requests.append((msg["id"], method))
        elif "method" in msg:
            # Server→client notification
            params = msg.get("params", {})
            if method == "textDocument/publishDiagnostics":
                self._on_diagnostics(params)

    def _on_diagnostics(self, params: dict) -> None:
        """Handle publishDiagnostics notification."""
        uri = params.get("uri", "")
        raw_diags = params.get("diagnostics", [])
        log.debug("publishDiagnostics: uri=%s count=%d", uri, len(raw_diags))

        # Convert URI to path for display
        path_str = uri
        if uri.startswith("file://"):
            path_str = uri[7:]

        diagnostics: list[Diagnostic] = []
        for d in raw_diags:
            rng = d.get("range", {}).get("start", {})
            diagnostics.append(
                Diagnostic(
                    path=path_str,
                    line=rng.get("line", 0) + 1,  # LSP is 0-indexed
                    character=rng.get("character", 0) + 1,
                    message=d.get("message", ""),
                    severity=d.get("severity", 1),
                    source=d.get("source", ""),
                )
            )

        self._diagnostics[uri] = diagnostics
        event = self._diag_events.get(uri)
        if event:
            event.set()

    async def _kill(self) -> None:
        """Kill the server process and clean up."""
        for task in (self._reader_task, getattr(self, "_stderr_task", None)):
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        if self._process:
            try:
                self._process.kill()
                await self._process.wait()
            except (ProcessLookupError, OSError):
                pass
            self._process = None
        self._broken = True


def _parse_document_symbols(
    items: list[dict],
    file_path: str,
    parent_name: str | None = None,
) -> list[DocumentSymbol]:
    """Recursively parse LSP DocumentSymbol response."""
    symbols: list[DocumentSymbol] = []
    for item in items:
        name = item.get("name", "")
        kind = item.get("kind", 0)
        rng = item.get("selectionRange") or item.get("range", {})
        start = rng.get("start", {})
        line = start.get("line", 0) + 1  # 0-indexed → 1-indexed

        detail = item.get("detail", "")
        signature = f"{name}({detail})" if detail else name

        symbols.append(
            DocumentSymbol(
                name=name,
                kind=kind,
                file=file_path,
                line=line,
                signature=signature,
                parent=parent_name,
            )
        )

        # Recurse into children
        children = item.get("children", [])
        if children:
            symbols.extend(_parse_document_symbols(children, file_path, parent_name=name))

    return symbols
