from __future__ import annotations

import copy
import re
from collections.abc import Iterable, Sequence
from enum import Enum

import lsprotocol.types as lsp
from tclint.syntax_tree import BareWord, Command, List, Visitor

from common.load_data import standard_items
from tools.navigation import FileSymbolIndex
from tools.variable_names import variable_name
from tools.tcl_command_completion import (
    DynamicCompletionKind,
    line_prefix_at_position,
    tcl_argument_completion,
)

BUILTIN_VAR_LABELS = {ci.label for ci in standard_items.nx_variables}
BUILTIN_PROC_LABELS = {ci.label for ci in standard_items.nx_procs}


class CompletionContext(Enum):
    VARIABLE = "variable"
    COMMAND = "command"
    GENERAL = "general"


VARIABLE_KINDS = {
    lsp.CompletionItemKind.Variable,
    lsp.CompletionItemKind.Field,
    lsp.CompletionItemKind.Constant,
}
COMMAND_KINDS = {
    lsp.CompletionItemKind.Class,
    lsp.CompletionItemKind.Function,
    lsp.CompletionItemKind.Method,
    lsp.CompletionItemKind.Constructor,
    lsp.CompletionItemKind.Keyword,
    lsp.CompletionItemKind.Snippet,
}

_VARIABLE_PREFIX_RE = re.compile(r"(?<!\\)\$(?:\{)?[A-Za-z0-9_:]*$")
_COMMAND_PREFIX_RE = re.compile(r"(?:^|[;\[\{])\s*[^\s;\[\]\{\}]*$")
_ARRAY_PREFIX_RE = re.compile(
    r"(?P<name>(?:::)?[A-Za-z_][A-Za-z0-9_:]*)\((?P<key>[^()\n]*)$"
)


def array_element_completions(
    source_lines: Sequence[str],
    position: lsp.Position,
    indexes: Iterable[FileSymbolIndex],
    current_path: str,
) -> list[lsp.CompletionItem] | None:
    """Complete literal array keys, replacing only the text inside parentheses."""
    prefix = line_prefix_at_position(source_lines, position)
    if prefix is None:
        return None
    match = _ARRAY_PREFIX_RE.search(prefix)
    if match is None:
        return None
    before = prefix[:match.start()]
    if not before.endswith("$"):
        argument = tcl_argument_completion(source_lines, position)
        if argument is None or argument.dynamic_kind != DynamicCompletionKind.VARIABLE:
            return None
    name, key_prefix = match.group("name", "key")
    part_index = key_prefix.count(",")
    key_prefix = key_prefix.rsplit(",", 1)[-1]
    if any(char in key_prefix for char in "$[]{}\\"):
        return None
    indexes = list(indexes)
    local_scope = None
    for index in indexes:
        if index.path != current_path:
            continue
        for occurrence in index.occurrences:
            span = occurrence.declaration_range
            if (
                occurrence.identity.kind == "proc"
                and occurrence.is_definition
                and span is not None
                and span.start.line <= position.line <= span.end.line
            ):
                local_scope = f"{index.path}::proc::{occurrence.identity.name}"
    keys: dict[str, set[tuple[str | None, ...]]] = {}
    for index in indexes:
        for occurrence in index.occurrences:
            parts = occurrence.array_parts
            key = parts[part_index] if part_index < len(parts) else None
            if (
                key
                and (
                    occurrence.identity.scope is None
                    or occurrence.identity.scope == local_scope
                )
                and key.startswith(key_prefix)
                and not any(char in key for char in "$[]\\")
                and occurrence.identity.name.removeprefix("::") == name.removeprefix("::")
            ):
                keys.setdefault(key, set()).add(occurrence.array_template_parts[part_index + 1:])
    line = source_lines[position.line]
    suffix = line[len(prefix):]
    remaining = re.match(r"[^(),\s$\[\]{}]*", suffix).group()
    has_close = suffix[len(remaining):].startswith((")", ","))
    start = position.character - len(key_prefix.encode("utf-16-le")) // 2
    end = position.character + len(remaining.encode("utf-16-le")) // 2
    items = []
    for key in sorted(keys):
        new_text = key
        # Fill missing index components only; keep an existing comma and suffix.
        if not suffix[len(remaining):].startswith(","):
            tails = sorted(keys[key], key=lambda tail: (len(tail), repr(tail)))
            tail = tails[0]
            if all(part is not None for part in tail) and any("$" in part for part in tail):
                new_text = ",".join([key, *tail])
        new_text += "" if has_close else ")"
        items.append(lsp.CompletionItem(
            label=key,
            kind=lsp.CompletionItemKind.Field,
            detail=f"{name}({key})",
            insert_text_format=lsp.InsertTextFormat.PlainText,
            text_edit=lsp.TextEdit(
                range=lsp.Range(
                    start=lsp.Position(line=position.line, character=start),
                    end=lsp.Position(line=position.line, character=end),
                ),
                new_text=new_text,
            ),
        ))
    return items


def completion_context(
    source_lines: Sequence[str], position: lsp.Position
) -> CompletionContext:
    prefix = line_prefix_at_position(source_lines, position)
    if prefix is None:
        return CompletionContext.GENERAL
    if _VARIABLE_PREFIX_RE.search(prefix):
        return CompletionContext.VARIABLE
    if _COMMAND_PREFIX_RE.search(prefix):
        return CompletionContext.COMMAND
    return CompletionContext.GENERAL


def _is_allowed(item: lsp.CompletionItem, context: CompletionContext) -> bool:
    if context == CompletionContext.VARIABLE:
        return item.kind in VARIABLE_KINDS
    if context == CompletionContext.COMMAND:
        return item.kind in COMMAND_KINDS
    return True


def ranked_completion_items(
    candidates: Iterable[tuple[int, lsp.CompletionItem]],
    context: CompletionContext,
) -> list[lsp.CompletionItem]:
    """Filter, de-duplicate, and rank completion candidates for one request."""
    items = []
    seen: set[tuple[str, lsp.CompletionItemKind | None]] = set()
    for sequence, (priority, item) in enumerate(candidates):
        if not _is_allowed(item, context):
            continue

        key = (item.label, item.kind)
        if key in seen:
            continue
        seen.add(key)

        ranked = copy.copy(item)
        ranked.sort_text = f"{priority:03d}:{item.label.casefold()}:{sequence:06d}"
        items.append(ranked)
    return items


class CompletionItems:
    def __init__(self):
        self._custom_functions: list[lsp.CompletionItem] = []

    @property
    def custom_functions(self) -> list[lsp.CompletionItem]:
        return self._custom_functions

    @custom_functions.setter
    def custom_functions(self, value: lsp.CompletionItem):
        self._custom_functions.append(value)


class CompletionCollector(Visitor):
    def __init__(self):
        super().__init__()
        self._custom_functions: list[lsp.CompletionItem] = []
        self._custom_function_keys: set[tuple[str, lsp.CompletionItemKind | None]] = set()
        self._proc_signatures = {}

    @property
    def custom_functions(self) -> list[lsp.CompletionItem]:
        return self._custom_functions

    @property
    def proc_signatures(self):
        return self._proc_signatures

    def _append_unique(self, item: lsp.CompletionItem):
        # Avoid duplicate labels within the same file scan
        key = (item.label, item.kind)
        if key in self._custom_function_keys:
            return
        self._custom_function_keys.add(key)
        self._custom_functions.append(item)

    def visit_command(self, command: Command):
        routine = command.routine

        # Collect custom proc names and their signatures
        if routine.contents == "proc" and command.args:
            first_arg = command.args[0]
            if not getattr(first_arg, "value", None):
                return

            if first_arg.value in BUILTIN_PROC_LABELS:
                return

            # Record proc name as a completion item
            self._append_unique(lsp.CompletionItem(label=first_arg.value, kind=lsp.CompletionItemKind.Function))
            if len(command.args) < 2:
                return

            param_list_node = command.args[1]
            if not hasattr(param_list_node, "children"):
                return

            param_names = []
            for arg in param_list_node.children:
                if isinstance(arg, BareWord):
                    param_names.append(arg.value)
                elif isinstance(arg, List) and len(arg.children) >= 1:
                    first = arg.children[0]
                    if isinstance(first, BareWord):
                        param_names.append(first.value)

            self._proc_signatures[first_arg.value] = param_names

        # Collect global variables declared with: global var1 var2 ...
        elif routine.contents == "global" and command.args:
            for arg in command.args:
                if (
                    isinstance(arg, BareWord)
                    and getattr(arg, "value", None)
                    and arg.value not in BUILTIN_VAR_LABELS
                ):
                    self._append_unique(
                        lsp.CompletionItem(
                            label=arg.value,
                            kind=lsp.CompletionItemKind.Variable,
                        )
                    )

        # Collect variables set with explicit global namespace: set ::var_name ...
        elif routine.contents == "set" and command.args:
            first = command.args[0]
            var_name = variable_name(first)
            if var_name:
                if var_name.startswith("::"):
                    base_name = var_name.split("(", 1)[0]
                    clean_name = base_name[2:]  # remove leading '::' for completion display
                    if clean_name not in BUILTIN_VAR_LABELS:
                        self._append_unique(lsp.CompletionItem(label=clean_name, kind=lsp.CompletionItemKind.Variable))