"""Custom widgets for the TactiTerm TUI with auto-completion engine and StdinInput."""
import asyncio
from typing import List, Dict
from textual import events
from textual.message import Message
from textual.widgets import TextArea, Input

# Comprehensive auto-completion syntax catalog per language for TUI
TUI_COMPLETIONS: Dict[str, List[str]] = {
    "python": [
        "int", "float", "bool", "str", "list", "dict", "set", "tuple", "bytes", "None",
        "True", "False", "print()", "len()", "range()", "enumerate()", "zip()", "map()",
        "filter()", "sorted()", "isinstance()", "type()", "open()", "input()", ".append()",
        ".split()", ".join()", ".get()", "def ", "class ", "return ", "import ", "from ",
        "if ", "elif ", "else:", "try:", "except ", "finally:", "raise ", "with ", "as ",
        "yield ", "async ", "await ", "lambda ", "pass", "break", "continue", "assert "
    ],
    "rust": [
        "bool", "i32", "i64", "u32", "u64", "usize", "f32", "f64", "char", "str", "String",
        "Vec", "Option", "Result", "Some", "None", "Ok", "Err", "println!()", "print!()",
        "format!()", "vec![]", "Vec::push()", "String::from()", "Option::unwrap()",
        "Result::expect()", ".iter()", ".collect()", "Box::new()", "fn ", "let ", "let mut ",
        "pub ", "struct ", "enum ", "impl ", "trait ", "match ", "if ", "else ", "loop ",
        "while ", "for ", "in ", "return ", "use std::", "async ", "await "
    ],
    "cpp": [
        "int", "float", "double", "bool", "char", "void", "size_t", "std::string",
        "std::vector<", "std::map<", "std::set<", "std::unique_ptr<", "std::shared_ptr<",
        "std::cout << ", "std::cin >> ", "std::endl", ".push_back()", "std::sort()",
        "std::make_unique<", "std::make_shared<", "class ", "struct ", "public:", "private:",
        "protected:", "virtual ", "override", "const ", "constexpr ", "auto ", "template<typename T>",
        "namespace ", "using namespace std;", "return ", "if ", "else ", "for ", "while "
    ],
    "go": [
        "int", "int64", "float64", "bool", "string", "byte", "rune", "error", "nil", "true",
        "false", "fmt.Println()", "fmt.Printf()", "fmt.Sprintf()", "make()", "append()",
        "len()", "cap()", "delete()", "strings.Split()", "os.Open()", "func ", "var ", "type ",
        "struct ", "interface ", "package main", "import ", "return ", "if ", "else ", "for ",
        "range ", "switch ", "select ", "chan ", "go func()"
    ],
    "javascript": [
        "number", "string", "boolean", "bool", "any", "void", "null", "undefined", "Array",
        "Object", "Promise", "console.log()", "console.error()", "console.warn()", ".map()",
        ".filter()", ".reduce()", ".find()", ".push()", ".includes()", "Object.keys()",
        "JSON.parse()", "fetch()", "const ", "let ", "var ", "function ", "class ", "return ",
        "if ", "else ", "for ", "while ", "switch ", "try ", "catch ", "async ", "await ", "new "
    ],
    "typescript": [
        "number", "string", "boolean", "bool", "any", "void", "unknown", "never", "null",
        "undefined", "Array<", "Record<", "Partial<", "Readonly<", "Pick<", "Omit<",
        "Promise<", "console.log()", ".map()", ".filter()", "interface ", "type ", "const ",
        "let ", "function ", "class ", "export ", "import ", "async ", "await "
    ],
    "java": [
        "int", "long", "double", "float", "boolean", "bool", "char", "byte", "short", "void",
        "String", "List<", "ArrayList<", "Map<", "HashMap<", "Set<", "HashSet<",
        "System.out.println()", ".add()", ".get()", ".size()", "public ", "private ",
        "protected ", "static ", "final ", "class ", "interface ", "extends ", "implements ",
        "new ", "return ", "if ", "else ", "for ", "while ", "try ", "catch "
    ],
    "csharp": [
        "int", "long", "double", "float", "bool", "char", "string", "object", "void", "var",
        "List<", "Dictionary<", "HashSet<", "Task<", "Console.WriteLine()", "Console.ReadLine()",
        ".Add()", "Enumerable.Where()", "public ", "private ", "protected ", "internal ",
        "static ", "class ", "struct ", "interface ", "record ", "namespace ", "using System;",
        "new ", "return ", "if ", "else ", "foreach ", "while ", "async ", "await "
    ],
    "lua": [
        "number", "int", "string", "boolean", "bool", "table", "nil", "function", "print()",
        "type()", "tostring()", "tonumber()", "table.insert()", "string.sub()", "local ",
        "if ", "then", "else", "elseif ", "end", "for ", "while ", "do", "return ", "require()"
    ],
    "html": [
        "<!DOCTYPE html>", "<html>", "<head>", "<title>", "<body>", "<h1>", "<h2>", "<h3>",
        "<p>", "<span>", "<div>", "<a href=\"\">", "<img src=\"\">", "<form>", "<input type=\"text\">",
        "<button>", "<select>", "<option>", "<ul>", "<ol>", "<li>", "<table>", "<tr>", "<td>", "<th>",
        "<script>", "<link rel=\"stylesheet\">"
    ],
}


# Duplicate csharp completion list for c_sharp key
TUI_COMPLETIONS["c_sharp"] = TUI_COMPLETIONS["csharp"]


def get_tui_language(language: str) -> str:
    """Map display/challenge language names to Textual's expected TextArea language keys."""
    if not language:
        return "python"
    norm = language.lower().strip()
    mapping = {
        "c++": "cpp",
        "cpp": "cpp",
        "cxx": "cpp",
        "c#": "c_sharp",
        "cs": "c_sharp",
        "csharp": "c_sharp",
        "c_sharp": "c_sharp",
        "js": "javascript",
        "javascript": "javascript",
        "ts": "typescript",
        "typescript": "typescript",
        "py": "python",
        "python": "python",
        "rs": "rust",
        "rust": "rust",
        "golang": "go",
        "go": "go",
        "html": "html",
        "java": "java",
        "lua": "lua",
    }
    return mapping.get(norm, norm)


def register_custom_tree_sitter_languages(editor: TextArea) -> None:
    """Register tree-sitter language grammars with rich query definitions for cpp, c_sharp, typescript, lua if available."""
    langs = {
        "cpp": (
            "tree_sitter_cpp",
            "language",
            [
                '["if" "else" "for" "while" "return" "class" "struct" "namespace" "using" "public" "private" "protected" "virtual" "const" "inline" "static" "template" "typename" "new" "delete" "catch" "try" "throw"] @keyword',
                '(primitive_type) @type',
                '(type_identifier) @type',
                '(comment) @comment',
                '(string_literal) @string',
                '(system_lib_string) @string',
                '(number_literal) @number',
                '(field_identifier) @property',
                '(function_declarator declarator: (identifier) @function)',
                '(call_expression function: (identifier) @function.call)',
                '(preproc_include) @include',
                '(preproc_def) @keyword',
                '(true) @boolean',
                '(false) @boolean',
                '(null) @constant.builtin',
            ],
        ),
        "c_sharp": (
            "tree_sitter_c_sharp",
            "language",
            [
                '["if" "else" "for" "foreach" "while" "return" "class" "struct" "interface" "public" "private" "protected" "internal" "static" "async" "await" "using" "namespace" "new" "get" "set" "try" "catch" "throw"] @keyword',
                '(predefined_type) @type',
                '(comment) @comment',
                '(string_literal) @string',
                '(verbatim_string_literal) @string',
                '(integer_literal) @number',
                '(real_literal) @number',
                '(identifier) @variable',
                '(method_declaration name: (identifier) @function)',
                '(invocation_expression function: (identifier) @function.call)',
                '(boolean_literal) @boolean',
                '(null_literal) @constant.builtin',
            ],
        ),
        "typescript": (
            "tree_sitter_typescript",
            "language_typescript",
            [
                '["if" "else" "for" "while" "return" "function" "class" "interface" "type" "const" "let" "var" "import" "from" "export" "async" "await" "new" "try" "catch" "throw" "switch" "case"] @keyword',
                '(predefined_type) @type',
                '(type_identifier) @type',
                '(comment) @comment',
                '(string) @string',
                '(template_string) @string',
                '(number) @number',
                '(property_identifier) @property',
                '(function_declaration name: (identifier) @function)',
                '(call_expression function: (identifier) @function.call)',
                '(true) @boolean',
                '(false) @boolean',
                '(null) @constant.builtin',
                '(undefined) @constant.builtin',
            ],
        ),
        "lua": (
            "tree_sitter_lua",
            "language",
            [
                '["if" "then" "else" "elseif" "end" "function" "return" "while" "for" "do" "local" "repeat" "until" "and" "or" "not" "in"] @keyword',
                '(break_statement) @keyword',
                '(comment) @comment',
                '(string) @string',
                '(number) @number',
                '(function_declaration name: (identifier) @function)',
                '(function_call name: (identifier) @function.call)',
                '(dot_index_expression field: (identifier) @property)',
                '(true) @boolean',
                '(false) @boolean',
                '(nil) @constant.builtin',
            ],
        ),
    }
    for key, (mod_name, func_name, sample_queries) in langs.items():
        if key not in editor.available_languages:
            try:
                import importlib
                import textual._tree_sitter as ts
                import textual.widgets._text_area as ta

                mod = importlib.import_module(mod_name)
                func = getattr(mod, func_name)
                lang_obj = ts.Language(func())

                query_str = ""
                for q in sample_queries:
                    try:
                        document = ta.SyntaxAwareDocument("", lang_obj)
                        document.prepare_query(q)
                        query_str += ("\n" if query_str else "") + q
                    except Exception:
                        pass

                editor.register_language(key, lang_obj, query_str)
            except Exception:
                pass


import re

class CodebaseSymbolExtractor:
    """Extracts symbols, variables, functions, and classes from the active code buffer."""

    @staticmethod
    def extract_symbols(code: str) -> List[str]:
        if not code:
            return []
        words = re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*", code)
        unique = set()
        reserved = {
            "if", "else", "for", "while", "return", "import", "from", "def", "class",
            "public", "private", "protected", "static", "void", "int", "double", "float",
            "bool", "boolean", "char", "const", "let", "var", "func", "package", "struct",
        }
        for w in words:
            if len(w) > 2 and w not in reserved:
                unique.add(w)
        return sorted(list(unique))


METHOD_CATALOGS: Dict[str, Dict[str, List[str]]] = {
    "java": {
        "map": ["put(key, value)", "get(key)", "containsKey(key)", "containsValue(val)", "size()", "isEmpty()", "keySet()", "values()", "entrySet()", "remove(key)", "clear()", "getOrDefault(key, default)"],
        "list": ["add(element)", "get(index)", "size()", "remove(index)", "contains(element)", "indexOf(element)", "isEmpty()", "clear()", "set(index, element)"],
        "set": ["add(element)", "remove(element)", "contains(element)", "size()", "isEmpty()", "clear()"],
        "string": ["length()", "substring(beginIndex)", "charAt(index)", "toLowerCase()", "toUpperCase()", "trim()", "split(regex)", "contains(str)", "startsWith(prefix)", "equals(obj)"],
        "system": ["println(value)", "printf(format, args)", "print(value)"],
    },
    "python": {
        "map": ["get(key)", "keys()", "values()", "items()", "update(dict)", "pop(key)", "clear()"],
        "list": ["append(item)", "extend(iterable)", "insert(index, item)", "remove(item)", "pop()", "sort()", "reverse()", "clear()", "count(item)", "index(item)"],
        "string": ["split(sep)", "join(iterable)", "lower()", "upper()", "strip()", "replace(old, new)", "startswith(prefix)", "endswith(suffix)", "find(sub)"],
    },
    "cpp": {
        "map": ["insert({key, val})", "find(key)", "count(key)", "size()", "empty()", "clear()", "at(key)"],
        "list": ["push_back(val)", "pop_back()", "size()", "empty()", "clear()", "begin()", "end()", "at(idx)"],
        "string": ["length()", "size()", "substr(pos, len)", "append(str)", "find(str)", "c_str()", "empty()"],
        "std": ["cout << value << std::endl;", "cin >> var;", "vector<int>", "map<string, int>", "sort(begin, end)"],
    },
    "c_sharp": {
        "map": ["Add(key, val)", "ContainsKey(key)", "TryGetValue(key, out val)", "Remove(key)", "Count", "Clear()"],
        "list": ["Add(item)", "Remove(item)", "RemoveAt(index)", "Contains(item)", "Count", "Clear()"],
        "string": ["Length", "Substring(startIndex)", "ToLower()", "ToUpper()", "Trim()", "Split(sep)", "Replace(old, new)"],
        "console": ["WriteLine(value)", "ReadLine()", "Write(value)"],
    },
    "javascript": {
        "map": ["set(key, val)", "get(key)", "has(key)", "delete(key)", "clear()", "size"],
        "list": ["map(x => x)", "filter(x => true)", "push(item)", "pop()", "slice(start, end)", "includes(item)", "length"],
        "console": ["log(msg)", "error(err)", "warn(msg)"],
    },
    "typescript": {
        "map": ["set(key, val)", "get(key)", "has(key)", "delete(key)", "clear()", "size"],
        "list": ["map(x => x)", "filter(x => true)", "push(item)", "pop()", "slice(start, end)", "includes(item)", "length"],
        "console": ["log(msg)", "error(err)", "warn(msg)"],
    },
    "rust": {
        "map": ["insert(key, val)", "get(&key)", "contains_key(&key)", "remove(&key)", "len()", "is_empty()"],
        "list": ["push(val)", "pop()", "len()", "is_empty()", "contains(&val)", "iter()", "collect()"],
    },
    "go": {
        "fmt": ["Println(v)", "Printf(format, v)", "Sprintf(format, v)"],
        "strings": ["Split(s, sep)", "Join(a, sep)", "ToLower(s)", "ToUpper(s)"],
    },
    "lua": {
        "table": ["insert(t, val)", "remove(t, pos)", "concat(t, sep)", "sort(t)"],
        "string": ["sub(s, i, j)", "lower(s)", "upper(s)", "len(s)"],
    },
}


class CodeEditor(TextArea):
    """A code editor with auto-completion dropdown trigger."""

    DEFAULT_CSS = """
    CodeEditor {
        height: 1fr;
        min-height: 8;
        border: solid $accent;
    }
    CodeEditor:focus {
        border: heavy $primary;
    }
    """

    def __init__(
        self,
        language: str = "python",
        theme: str = "vscode_dark",
        id: str | None = None,
        classes: str | None = None,
    ) -> None:
        super().__init__(
            language=None,
            theme=theme,
            soft_wrap=False,
            show_line_numbers=True,
            id=id,
            classes=classes,
        )
        register_custom_tree_sitter_languages(self)
        self.language = get_tui_language(language)

    def on_mount(self) -> None:
        self.tab_behavior = "indent"
        self.cursor_style = "bold blue"

    def get_code(self) -> str:
        return self.text

    def get_completions_at_cursor(self) -> tuple[str, List[str], bool]:
        """Returns (prefix_to_replace, matching_candidates, is_dot_access)."""
        cursor_row, cursor_col = self.cursor_location
        lines = self.text.split("\n")
        if cursor_row >= len(lines):
            return ("", [], False)

        current_line = lines[cursor_row][:cursor_col]
        lang = get_tui_language(self.language or "python")

        # 1. Check dot member access (e.g. stockMap. or stockMap.p or System.out.)
        dot_match = re.search(r"([a-zA-Z_][a-zA-Z0-9_]*)\.([a-zA-Z0-9_]*)$", current_line)
        if dot_match:
            obj_name = dot_match.group(1)
            member_prefix = dot_match.group(2)
            obj_lower = obj_name.lower()

            lang_methods = METHOD_CATALOGS.get(lang, METHOD_CATALOGS.get(lang.replace("_", ""), METHOD_CATALOGS["java"]))
            target_type = "map"
            if "map" in obj_lower or "dict" in obj_lower:
                target_type = "map"
            elif "list" in obj_lower or "arr" in obj_lower or "items" in obj_lower:
                target_type = "list"
            elif "set" in obj_lower:
                target_type = "set"
            elif "str" in obj_lower or "text" in obj_lower or "name" in obj_lower or "msg" in obj_lower:
                target_type = "string"
            elif obj_lower in lang_methods:
                target_type = obj_lower

            candidates = lang_methods.get(target_type, lang_methods.get("map", []))
            matches = [c for c in candidates if c.lower().startswith(member_prefix.lower())]
            return (member_prefix, matches, True)

        # 2. Standalone prefix word completion
        word_match = re.search(r"([a-zA-Z_][a-zA-Z0-9_]*)$", current_line)
        if word_match:
            prefix = word_match.group(1)
            if len(prefix) >= 1:
                static_candidates = TUI_COMPLETIONS.get(lang, TUI_COMPLETIONS.get(lang.replace("_", ""), TUI_COMPLETIONS["python"]))
                extracted_symbols = CodebaseSymbolExtractor.extract_symbols(self.text)
                all_candidates = list(dict.fromkeys(static_candidates + extracted_symbols))
                matches = [c for c in all_candidates if c.lower().startswith(prefix.lower()) and c.lower() != prefix.lower()]
                return (prefix, matches, False)

        return ("", [], False)

    def on_text_area_changed(self, event: TextArea.Changed) -> None:
        """Trigger instant as-you-type completion update on any text mutation."""
        if hasattr(self.app, "update_as_you_type_completion"):
            self.app.update_as_you_type_completion(self)

    def _on_key(self, event: events.Key) -> None:
        """Handle completion key navigation, selection, escape, and Ctrl+E toggle."""
        if event.key == "ctrl+e":
            event.prevent_default()
            event.stop()
            if hasattr(self.app, "action_toggle_expand_output"):
                self.app.action_toggle_expand_output()
            return

        is_open = getattr(self.app, "is_completion_open", lambda: False)()

        if is_open:
            if event.key in ("down", "up"):
                event.prevent_default()
                event.stop()
                if hasattr(self.app, "navigate_completion"):
                    self.app.navigate_completion(1 if event.key == "down" else -1)
                return
            elif event.key in ("enter", "tab"):
                event.prevent_default()
                event.stop()
                if hasattr(self.app, "insert_selected_completion"):
                    self.app.insert_selected_completion()
                return
            elif event.key == "escape":
                event.prevent_default()
                event.stop()
                if hasattr(self.app, "hide_completion_popup"):
                    self.app.hide_completion_popup()
                return

        super()._on_key(event)


class StdinInput(Input):
    """Input control for stdin that forwards Ctrl+E to screen output toggler."""

    def _on_key(self, event: events.Key) -> None:
        if event.key == "ctrl+e":
            event.prevent_default()
            event.stop()
            if hasattr(self.app, "action_toggle_expand_output"):
                self.app.action_toggle_expand_output()
            return
        res = super()._on_key(event)
        if asyncio.iscoroutine(res):
            asyncio.create_task(res)


class MentorInputTextArea(TextArea):
    """Multi-line question input for Socratic Mentor. Submits on Enter, inserts newline on Shift+Enter."""

    class SubmitQuestion(Message):
        """Event posted when Enter is pressed."""

        def __init__(self, question: str) -> None:
            self.question = question
            super().__init__()

    def _on_key(self, event: events.Key) -> None:
        if event.key == "ctrl+e":
            event.prevent_default()
            event.stop()
            if hasattr(self.app, "action_toggle_expand_output"):
                self.app.action_toggle_expand_output()
            return
        elif event.key == "enter":
            event.prevent_default()
            event.stop()
            self.post_message(self.SubmitQuestion(self.text))
        elif event.key in ["shift+enter", "ctrl+j"]:
            self.insert("\n")
            event.prevent_default()
        else:
            super()._on_key(event)
