import logging
import pathlib
from typing import List, Optional, Tuple
import lsprotocol.types as lsp
from pygls.workspace.text_document import TextDocument
from tclint.lexer import TclSyntaxError
from tclint.format import Formatter, FormatterOpts
from tclint.violations import Violation
from plugins.poco_plugin import commands
from tools import checks, parser
from pygls import server, uris
from tools.completion_items import completion, remove_existing_items, remove_shared_keys
from tools.proc_docs import build_proc_docs


DIAGNOSTIC_SOURCE = "nx-post-support"


class TclLanguageServer(server.LanguageServer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.parser = parser.CustomParser()
        for command in commands:
            self.parser._commands.update(command)
        self.diagnostics = {}
        self.poco_completion: dict = {}
        self.proc_signatures: dict = {}
        self.proc_docs: dict = {}
        # Cache: (uri, version) -> (tree, violations)
        self._ast_cache = {}

    def get_tree(self, document: TextDocument):
        key = (document.uri, document.version)
        cached = self._ast_cache.get(key)
        if cached:
            return cached[0]
        # Parse and cache
        self.parser.violations = []
        tree = self.parser.parse(document.source)
        violations = list(self.parser.violations)
        self._ast_cache[key] = (tree, violations)
        return tree

    def get_tree_and_violations(self, document: TextDocument):
        key = (document.uri, document.version)
        cached = self._ast_cache.get(key)
        if cached:
            return cached
        # Parse and cache
        self.parser.violations = []
        tree = self.parser.parse(document.source)
        violations = list(self.parser.violations)
        self._ast_cache[key] = (tree, violations)
        return tree, violations

    def clear_cache_for_uri(self, uri: str):
        to_delete = [k for k in self._ast_cache.keys() if k[0] == uri]
        for k in to_delete:
            del self._ast_cache[k]

    def update_poco_completion_for_file(self, document: TextDocument):
        """Update poco_completion for a specific file when it changes"""
        filepath = str(pathlib.Path(uris.to_fs_path(document.uri)))

        # Remove existing completion items for this file
        if filepath in self.poco_completion:
            del self.poco_completion[filepath]
        if filepath in self.proc_signatures:
            del self.proc_signatures[filepath]
        if filepath in self.proc_docs:
            del self.proc_docs[filepath]

        # Parse and extract new completion items
        completion.reset()
        try:
            tree = self.get_tree(document)
            tree.accept(completion, recurse=True)
            remove_existing_items(completion.custom_functions, self.poco_completion)
            self.poco_completion[filepath] = completion.custom_functions
            remove_shared_keys(self.proc_signatures, completion.proc_signatures)
            self.proc_signatures[filepath] = completion.proc_signatures
            self.proc_docs[filepath] = build_proc_docs(tree, document.source)
        except Exception as e:
            logging.debug(f"Error parsing {filepath}: {e}")

    def format(
        self,
        document: TextDocument,
        options: lsp.FormattingOptions,
        range: Optional[Tuple[int, int]] = None,
    ):
        # parser = Parser(command_plugins=["nx_plugins.poco_plugin.py"])
        # parser._commands.update(commands)

        indent = "\t" if not options.insert_spaces else " " * options.tab_size
        formatter = Formatter(
            FormatterOpts(
                indent=indent,
                spaces_in_braces=False,
                max_blank_lines=500,
                indent_namespace_eval=True,
            ),
        )

        if range is not None:
            start, end = range
            return formatter.format_partial(document.source[start:end], self.parser)

        return formatter.format_top(document.source, self.parser)

    def linter(
        self,
        document: TextDocument,
    ) -> List[Violation]:
        tree, violations = self.get_tree_and_violations(document)
        for checker in checks.get_checkers():
            violations += checker.check(document.source, tree)
        return violations

    def lint(self, document: TextDocument):
        diagnostics = []

        try:
            violations = self.linter(document)
        except TclSyntaxError as e:
            return [
                lsp.Diagnostic(
                    message=str(e),
                    severity=lsp.DiagnosticSeverity.Error,
                    range=lsp.Range(
                        start=lsp.Position(e.start[0] - 1, e.start[1] - 1),
                        end=lsp.Position(e.end[0] - 1, e.end[1] - 1),
                    ),
                    code="syntax error",
                    source=DIAGNOSTIC_SOURCE,
                )
            ]

        for violation in violations:
            message = violation.message
            severity = lsp.DiagnosticSeverity.Warning
            start = lsp.Position(line=violation.start[0] - 1, character=violation.start[1] - 1)
            end = lsp.Position(line=violation.end[0] - 1, character=violation.end[1] - 1)

            diagnostics.append(
                lsp.Diagnostic(
                    message=message,
                    severity=severity,
                    range=lsp.Range(
                        start=start,
                        end=end,
                    ),
                    code=violation.id,
                    source=DIAGNOSTIC_SOURCE,
                )
            )

        return diagnostics

    def _compute_diagnostics(self, document: TextDocument) -> List[lsp.Diagnostic]:
        return self.lint(document)

    def compute_diagnostics(self, document: TextDocument):
        # `None` sentinel ensures that `diagnostics` gets updated if the URI is not
        # present.
        _, previous = self.diagnostics.get(document, (0, None))

        diagnostics = self._compute_diagnostics(document)

        # Only update if the list has changed
        if previous != diagnostics:
            self.diagnostics[document.uri] = (document.version, diagnostics)