"""
Nexa Language — LSP Server (v26)
JSON-RPC 2.0 over stdio — compatible with VSCode's language-client.

Features (v26):
  - textDocument/hover      — type info + docstring
  - textDocument/completion — keyword + built-in + symbol completions
  - textDocument/definition — cross-file go-to-definition (v26 UPGRADED)
  - workspace/symbol        — all function, class, trait, variable names

Usage:
  nexa lsp           (started automatically by VSCode extension)
  echo '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}' | python -m nexa_lang lsp
"""
from __future__ import annotations
import sys, json, threading, os
from typing import Any, Dict, Optional


# ── Built-in symbol table for completions ─────────────────────────────────────
NEXA_KEYWORDS = [
    "let", "var", "def", "fn", "class", "trait", "impl", "use",
    "if", "else", "for", "while", "loop",
    "return", "break", "continue", "import", "from", "export",
    "true", "false", "null", "new", "self", "try", "catch", "finally",
    "throw", "match", "case", "default", "interface", "implements",
    "abstract", "private", "public", "async", "await", "yield",
    "assert",  # v25
]

NEXA_BUILTINS = [
    "print", "input", "len", "type", "str", "int", "float", "bool",
    "range", "append", "map", "filter", "sorted", "reversed",
    "math", "stats", "fs", "json", "http", "time", "random",
    "nlp", "vision", "automl", "agent", "tensor", "df", "db",
]


def _make_response(id_: Any, result: Any) -> dict:
    return {"jsonrpc": "2.0", "id": id_, "result": result}


def _make_error(id_: Any, code: int, message: str) -> dict:
    return {"jsonrpc": "2.0", "id": id_, "error": {"code": code, "message": message}}


def _send(msg: dict):
    payload = json.dumps(msg)
    sys.stdout.write(f"Content-Length: {len(payload)}\r\n\r\n{payload}")
    sys.stdout.flush()


def _read_message() -> Optional[dict]:
    headers = {}
    while True:
        line = sys.stdin.readline()
        if not line: return None
        line = line.strip()
        if not line: break
        if ":" in line:
            k, v = line.split(":", 1)
            headers[k.strip()] = v.strip()
    length = int(headers.get("Content-Length", 0))
    if not length: return None
    body = sys.stdin.read(length)
    return json.loads(body)


class NexaLSP:
    """LSP server for the Nexa programming language."""

    def __init__(self):
        self._docs: Dict[str, str] = {}          # uri → source text
        self._symbols: Dict[str, list] = {}      # uri → list of (name, line, kind)

    def run(self):
        """Main LSP event loop — reads JSON-RPC messages from stdin."""
        while True:
            msg = _read_message()
            if msg is None: break
            try:
                self._handle(msg)
            except Exception as e:
                _send(_make_error(msg.get("id"), -32603, str(e)))

    def _handle(self, msg: dict):
        method = msg.get("method", "")
        id_ = msg.get("id")
        params = msg.get("params", {})

        if method == "initialize":
            _send(_make_response(id_, {
                "capabilities": {
                    "textDocumentSync": 1,
                    "hoverProvider": True,
                    "completionProvider": {"triggerCharacters": ["."]},
                    "definitionProvider": True,
                    "documentSymbolProvider": True,
                    "workspaceSymbolProvider": True,   # v26
                },
                "serverInfo": {"name": "nexals", "version": "26.0"},
            }))

        elif method == "initialized": pass  # notification, no response

        elif method == "shutdown":
            _send(_make_response(id_, None))

        elif method == "exit":
            sys.exit(0)

        elif method == "textDocument/didOpen":
            doc = params["textDocument"]
            self._docs[doc["uri"]] = doc["text"]
            self._index(doc["uri"], doc["text"])

        elif method == "textDocument/didChange":
            uri = params["textDocument"]["uri"]
            if params.get("contentChanges"):
                self._docs[uri] = params["contentChanges"][-1]["text"]
                self._index(uri, self._docs[uri])

        elif method == "textDocument/hover":
            uri = params["textDocument"]["uri"]
            pos = params["position"]
            result = self._hover(uri, pos["line"], pos["character"])
            _send(_make_response(id_, result))

        elif method == "textDocument/completion":
            uri = params["textDocument"]["uri"]
            pos = params["position"]
            result = self._completions(uri, pos["line"], pos["character"])
            _send(_make_response(id_, result))

        elif method == "textDocument/definition":
            uri = params["textDocument"]["uri"]
            pos = params["position"]
            result = self._definition(uri, pos["line"], pos["character"])
            _send(_make_response(id_, result))

        elif method == "textDocument/documentSymbol":
            uri = params["textDocument"]["uri"]
            result = self._document_symbols(uri)
            _send(_make_response(id_, result))

        elif id_ is not None:
            _send(_make_error(id_, -32601, f"Method not supported: {method}"))

    def _index(self, uri: str, src: str):
        """Parse source and build symbol table (v26: includes var, trait, impl)."""
        symbols = []
        try:
            from nexa_lang.lexer import Lexer
            from nexa_lang.parser import Parser
            tokens = Lexer(src).tokenize()
            ast = Parser(tokens).parse()
            for stmt in ast.statements:
                t = type(stmt).__name__
                line = getattr(stmt, "line", 0)
                if t == "FunctionDef":
                    symbols.append((stmt.name, line, "function"))
                elif t == "ClassDef":
                    symbols.append((stmt.name, line, "class"))
                    for item in stmt.body:
                        if type(item).__name__ == "FunctionDef":
                            symbols.append((f"{stmt.name}.{item.name}",
                                            getattr(item, "line", 0), "method"))
                elif t in ("LetStatement", "VarStatement"):   # v26: VarStatement
                    symbols.append((stmt.name, line, "variable"))
                elif t == "TraitDef":                          # v26: TraitDef
                    symbols.append((stmt.name, line, "trait"))
                    for m in stmt.methods:
                        symbols.append((f"{stmt.name}.{m.name}",
                                        getattr(m, "line", 0), "method"))
                elif t == "TraitImpl":                         # v26: TraitImpl
                    for m in stmt.methods:
                        symbols.append((f"{stmt.class_name}::{m.name}",
                                        getattr(m, "line", 0), "method"))
        except Exception:
            pass
        self._symbols[uri] = symbols

    def _word_at(self, src: str, line: int, col: int) -> str:
        lines = src.splitlines()
        if line >= len(lines): return ""
        row = lines[line]
        start = col
        while start > 0 and (row[start-1].isalnum() or row[start-1] == '_'): start -= 1
        end = col
        while end < len(row) and (row[end].isalnum() or row[end] == '_'): end += 1
        return row[start:end]

    def _hover(self, uri: str, line: int, col: int) -> Optional[dict]:
        src = self._docs.get(uri, "")
        word = self._word_at(src, line, col)
        if not word: return None
        # Look up in symbol table
        for name, sym_line, kind in self._symbols.get(uri, []):
            if name == word or name.endswith("." + word):
                return {"contents": {"kind": "markdown",
                                      "value": f"**{kind}** `{name}` — line {sym_line}"}}
        if word in NEXA_KEYWORDS:
            return {"contents": {"kind": "markdown", "value": f"**keyword** `{word}`"}}
        if word in NEXA_BUILTINS:
            return {"contents": {"kind": "markdown", "value": f"**built-in** `{word}`"}}
        return None

    def _completions(self, uri: str, line: int, col: int) -> dict:
        items = []
        # Add keywords
        for kw in NEXA_KEYWORDS:
            items.append({"label": kw, "kind": 14})  # 14 = keyword
        # Add builtins
        for bi in NEXA_BUILTINS:
            items.append({"label": bi, "kind": 3})   # 3 = function
        # Add user symbols
        for name, _, kind in self._symbols.get(uri, []):
            kind_code = 6 if kind in ("function","method") else 7 if kind == "class" else 6
            items.append({"label": name, "kind": kind_code})
        return {"isIncomplete": False, "items": items}

    def _definition(self, uri: str, line: int, col: int) -> Optional[list]:
        """v26: Cross-file go-to-definition."""
        src = self._docs.get(uri, "")
        word = self._word_at(src, line, col)
        if not word:
            return None

        # 1. Search current document first
        for name, sym_line, _ in self._symbols.get(uri, []):
            bare = name.split(".")[-1].split("::")[-1]
            if bare == word or name == word:
                return [{
                    "uri": uri,
                    "range": {
                        "start": {"line": max(0, sym_line - 1), "character": 0},
                        "end":   {"line": max(0, sym_line - 1), "character": 200},
                    }
                }]

        # 2. Search all other open documents (cross-file)
        for other_uri, sym_list in self._symbols.items():
            if other_uri == uri:
                continue
            for name, sym_line, _ in sym_list:
                bare = name.split(".")[-1].split("::")[-1]
                if bare == word or name == word:
                    return [{
                        "uri": other_uri,
                        "range": {
                            "start": {"line": max(0, sym_line - 1), "character": 0},
                            "end":   {"line": max(0, sym_line - 1), "character": 200},
                        }
                    }]

        # 3. Try to resolve import path from current file
        for raw_line in src.splitlines():
            stripped = raw_line.strip()
            if stripped.startswith("import ") and word in stripped:
                # e.g. import "utils" → try utils.nexa relative to file dir
                import re
                m = re.search(r'import\s+["\']([^"\']+)["\']', stripped)
                if m:
                    mod_path = m.group(1)
                    base_dir = os.path.dirname(uri.replace("file:///", "").replace("file://", ""))
                    candidate = os.path.join(base_dir, mod_path)
                    if not candidate.endswith(".nexa"):
                        candidate += ".nexa"
                    if os.path.exists(candidate):
                        cand_uri = "file:///" + candidate.replace("\\", "/")
                        # Index it if not already indexed
                        if cand_uri not in self._docs:
                            try:
                                with open(candidate, encoding="utf-8") as f:
                                    cand_src = f.read()
                                self._docs[cand_uri] = cand_src
                                self._index(cand_uri, cand_src)
                            except Exception:
                                pass
                        # Search indexed symbols
                        for name, sym_line, _ in self._symbols.get(cand_uri, []):
                            bare = name.split(".")[-1]
                            if bare == word:
                                return [{
                                    "uri": cand_uri,
                                    "range": {
                                        "start": {"line": max(0, sym_line - 1), "character": 0},
                                        "end":   {"line": max(0, sym_line - 1), "character": 200},
                                    }
                                }]
        return None

    def _document_symbols(self, uri: str) -> list:
        out = []
        for name, line, kind in self._symbols.get(uri, []):
            kind_map = {"function": 12, "class": 5, "method": 6, "variable": 13}
            out.append({
                "name": name,
                "kind": kind_map.get(kind, 13),
                "location": {
                    "uri": uri,
                    "range": {
                        "start": {"line": max(0, line-1), "character": 0},
                        "end":   {"line": max(0, line-1), "character": 100},
                    }
                }
            })
        return out


def run_lsp():
    """Entry point: start the LSP server over stdio."""
    NexaLSP().run()
