"""
Syntax highlighting using tree-sitter.

Supports:
- Bundled languages (Python, C, C++) via pip packages
- Installing additional languages via pip
- Incremental parsing for fast updates
- Fallback to simple highlighting for huge files

Uses tree-sitter 0.25+ API with pip-installable language packages.
"""

import json
import subprocess
import sys
from pathlib import Path
from typing import Optional, Dict, List, Tuple, Any

# Try to import tree-sitter
try:
    from tree_sitter import Language, Parser, Tree, Query, QueryCursor
    TREE_SITTER_AVAILABLE = True
except ImportError:
    TREE_SITTER_AVAILABLE = False
    Language = None
    Parser = None
    Tree = None
    Query = None
    QueryCursor = None

# Standard highlight groups -> curses color indices
# These map to curses.init_pair() calls in the editor
HIGHLIGHT_GROUPS = {
    'keyword': 1,        # Blue
    'string': 2,         # Green
    'comment': 3,        # Gray/dim
    'function': 4,       # Yellow
    'type': 5,           # Cyan
    'number': 6,         # Magenta
    'operator': 7,       # White
    'variable': 8,       # Default
    'constant': 9,       # Red
    'punctuation': 0,    # Default
    'property': 4,       # Yellow (like function)
    'parameter': 8,      # Default
    'builtin': 1,        # Blue (like keyword)
}

# Mapping from language name to pip package and module name
LANGUAGE_PACKAGES = {
    'python': ('tree-sitter-python', 'tree_sitter_python'),
    'c': ('tree-sitter-c', 'tree_sitter_c'),
    'cpp': ('tree-sitter-cpp', 'tree_sitter_cpp'),
    'javascript': ('tree-sitter-javascript', 'tree_sitter_javascript'),
    'typescript': ('tree-sitter-typescript', 'tree_sitter_typescript'),
    'rust': ('tree-sitter-rust', 'tree_sitter_rust'),
    'go': ('tree-sitter-go', 'tree_sitter_go'),
    'java': ('tree-sitter-java', 'tree_sitter_java'),
    'json': ('tree-sitter-json', 'tree_sitter_json'),
    'bash': ('tree-sitter-bash', 'tree_sitter_bash'),
    'markdown': ('tree-sitter-markdown', 'tree_sitter_markdown'),
    'sql': ('tree-sitter-sql', 'tree_sitter_sql'),
    'json': ('tree-sitter-json', 'tree_sitter_json'),
    'yaml': ('tree-sitter-yaml', 'tree_sitter_yaml'),
}

# Simple fallback patterns for huge files (regex-based)
SIMPLE_PATTERNS = {
    'python': {
        'keywords': ['def', 'class', 'if', 'elif', 'else', 'for', 'while', 'try',
                     'except', 'finally', 'with', 'as', 'import', 'from', 'return',
                     'yield', 'raise', 'pass', 'break', 'continue', 'and', 'or',
                     'not', 'in', 'is', 'lambda', 'True', 'False', 'None', 'async', 'await'],
        'comment': '#',
        'strings': ['"', "'", '"""', "'''"],
    },
    'c': {
        'keywords': ['if', 'else', 'for', 'while', 'do', 'switch', 'case', 'default',
                     'break', 'continue', 'return', 'goto', 'typedef', 'struct',
                     'union', 'enum', 'sizeof', 'static', 'extern', 'const', 'volatile',
                     'register', 'auto', 'void', 'int', 'char', 'short', 'long',
                     'float', 'double', 'signed', 'unsigned'],
        'comment': '//',
        'block_comment': ('/*', '*/'),
        'strings': ['"'],
    },
    'cpp': {
        'keywords': ['if', 'else', 'for', 'while', 'do', 'switch', 'case', 'default',
                     'break', 'continue', 'return', 'goto', 'typedef', 'struct',
                     'union', 'enum', 'sizeof', 'static', 'extern', 'const', 'volatile',
                     'class', 'public', 'private', 'protected', 'virtual', 'override',
                     'template', 'typename', 'namespace', 'using', 'new', 'delete',
                     'try', 'catch', 'throw', 'nullptr', 'true', 'false', 'this',
                     'void', 'int', 'char', 'short', 'long', 'float', 'double', 'bool',
                     'auto', 'constexpr', 'inline', 'explicit', 'friend', 'mutable'],
        'comment': '//',
        'block_comment': ('/*', '*/'),
        'strings': ['"'],
    },
    'sql': {
        'keywords': ['SELECT', 'FROM', 'WHERE', 'AND', 'OR', 'NOT', 'IN', 'IS', 'NULL',
                     'INSERT', 'INTO', 'VALUES', 'UPDATE', 'SET', 'DELETE', 'CREATE',
                     'TABLE', 'INDEX', 'VIEW', 'DROP', 'ALTER', 'JOIN', 'LEFT', 'RIGHT',
                     'INNER', 'OUTER', 'ON', 'GROUP', 'BY', 'ORDER', 'ASC', 'DESC',
                     'HAVING', 'LIMIT', 'UNION', 'ALL', 'DISTINCT', 'AS', 'LIKE',
                     'BETWEEN', 'EXISTS', 'CASE', 'WHEN', 'THEN', 'ELSE', 'END',
                     'BEGIN', 'COMMIT', 'ROLLBACK', 'PRIMARY', 'KEY', 'FOREIGN',
                     'select', 'from', 'where', 'and', 'or', 'insert', 'into', 'values',
                     'update', 'set', 'delete', 'create', 'table', 'drop', 'alter'],
        'comment': '--',
        'block_comment': ('/*', '*/'),
        'strings': ["'"],
    },
    'ttl': {
        'keywords': ['@prefix', '@base', 'a', 'true', 'false'],
        'comment': '#',
        'strings': ['"'],
    },
    'yaml': {
        'keywords': ['true', 'false', 'null', 'yes', 'no', 'on', 'off'],
        'comment': '#',
        'strings': ['"', "'"],
    },
    'bash': {
        'keywords': ['if', 'then', 'else', 'elif', 'fi', 'for', 'while', 'do', 'done',
                     'case', 'esac', 'in', 'function', 'return', 'exit', 'export',
                     'local', 'readonly', 'declare', 'source', 'alias', 'unset'],
        'comment': '#',
        'strings': ['"', "'"],
    },
}


class SyntaxHighlighter:
    """Manages syntax highlighting for the editor."""

    # Threshold for falling back to simple highlighting
    LARGE_FILE_THRESHOLD = 50000  # lines

    def __init__(self, syntax_dir: Optional[str] = None):
        self.syntax_dir = Path(syntax_dir) if syntax_dir else Path(__file__).parent
        self.queries_dir = self.syntax_dir / 'queries'
        self.registry_path = self.syntax_dir / 'registry.json'

        self.languages: Dict[str, Language] = {}
        self.parsers: Dict[str, Parser] = {}
        self.queries: Dict[str, Any] = {}  # Highlight queries
        self.registry: Dict[str, dict] = {}

        self.current_tree: Optional[Tree] = None
        self.current_lang: Optional[str] = None
        self.enabled = TREE_SITTER_AVAILABLE
        self.use_simple = False  # Fallback mode for large files

        self._load_registry()

    def _load_registry(self) -> None:
        """Load the language registry."""
        if self.registry_path.exists():
            with open(self.registry_path) as f:
                self.registry = json.load(f)

    def detect_language(self, filename: str) -> Optional[str]:
        """Detect language from filename extension."""
        if not filename:
            return None
        ext = Path(filename).suffix.lower()
        for lang, info in self.registry.items():
            if ext in info.get('extensions', []):
                return lang
        return None

    def is_available(self, lang: str) -> bool:
        """Check if a language grammar is available (installed via pip or has simple patterns)."""
        if lang in self.languages:
            return True
        # Check for tree-sitter package
        if lang in LANGUAGE_PACKAGES:
            _, module_name = LANGUAGE_PACKAGES[lang]
            try:
                __import__(module_name)
                return True
            except ImportError:
                pass
        # Fall back to simple patterns
        if lang in SIMPLE_PATTERNS:
            return True
        return False

    def load_language(self, lang: str) -> bool:
        """Load a language grammar. Returns True if successful."""
        if lang in self.languages:
            return True

        # Try tree-sitter first
        if self.enabled and lang in LANGUAGE_PACKAGES:
            _, module_name = LANGUAGE_PACKAGES[lang]
            try:
                # Import the language module
                lang_module = __import__(module_name)

                # Create Language from the module's language() function
                language = Language(lang_module.language())
                self.languages[lang] = language

                # Create parser with the language
                parser = Parser(language)
                self.parsers[lang] = parser

                # Try to load highlight query
                self._load_query(lang)
                return True
            except:
                pass  # Fall through to simple patterns

        # Fall back to simple patterns
        if lang in SIMPLE_PATTERNS:
            return True

        return False

    def _load_query(self, lang: str) -> None:
        """Load highlight query for a language."""
        query_file = self.queries_dir / f'{lang}.scm'
        if query_file.exists():
            with open(query_file) as f:
                query_text = f.read()
            try:
                # tree-sitter 0.25+ uses Query constructor
                self.queries[lang] = Query(self.languages[lang], query_text)
            except Exception as e:
                pass  # Query syntax may be incompatible

    def install_language(self, lang: str, callback=None) -> bool:
        """Install a language grammar via pip."""
        if lang not in LANGUAGE_PACKAGES:
            if callback:
                callback(f"Unknown language: {lang}")
            return False

        pip_package, _ = LANGUAGE_PACKAGES[lang]

        if callback:
            callback(f"Installing {pip_package}...")

        try:
            result = subprocess.run(
                [sys.executable, '-m', 'pip', 'install', pip_package],
                capture_output=True, text=True, timeout=120
            )
            if result.returncode != 0:
                if callback:
                    callback(f"Failed to install: {result.stderr}")
                return False

            if callback:
                callback(f"Installed {lang} grammar")

            return self.load_language(lang)
        except subprocess.TimeoutExpired:
            if callback:
                callback(f"Installation timed out")
            return False
        except Exception as e:
            if callback:
                callback(f"Failed to install: {e}")
            return False

    def parse(self, source: str, lang: str) -> Optional[Tree]:
        """Parse source code. Returns syntax tree."""
        if not self.enabled or lang not in self.parsers:
            return None

        try:
            self.current_tree = self.parsers[lang].parse(bytes(source, 'utf-8'))
            self.current_lang = lang
            return self.current_tree
        except Exception:
            return None

    def parse_incremental(self, source: str, edit_start: int, old_end: int, new_end: int,
                          start_point: Tuple[int, int], old_end_point: Tuple[int, int],
                          new_end_point: Tuple[int, int]) -> Optional[Tree]:
        """Incrementally update parse tree after an edit."""
        if not self.enabled or not self.current_tree or not self.current_lang:
            return None

        try:
            self.current_tree.edit(
                start_byte=edit_start,
                old_end_byte=old_end,
                new_end_byte=new_end,
                start_point=start_point,
                old_end_point=old_end_point,
                new_end_point=new_end_point,
            )
            self.current_tree = self.parsers[self.current_lang].parse(
                bytes(source, 'utf-8'),
                self.current_tree
            )
            return self.current_tree
        except Exception:
            return None

    def get_highlights(self, start_line: int, end_line: int, lines: List[str]) -> Dict[int, List[Tuple[int, int, int]]]:
        """
        Get highlights for a range of lines.

        Returns: {line_num: [(start_col, end_col, color_index), ...]}
        """
        # Use simple patterns if forced or no tree-sitter available
        if self.current_lang and (self.use_simple or not self.current_tree):
            if self.current_lang in SIMPLE_PATTERNS:
                return self._get_simple_highlights(start_line, end_line, lines)

        if not self.current_tree or not self.current_lang:
            return {}

        result: Dict[int, List[Tuple[int, int, int]]] = {}

        # Use highlight query if available
        if self.current_lang in self.queries:
            try:
                # tree-sitter 0.25+ uses QueryCursor
                cursor = QueryCursor(self.queries[self.current_lang])
                cursor.set_point_range((start_line, 0), (end_line + 1, 0))

                # matches() returns iterator of (pattern_idx, captures_dict)
                for pattern_idx, captures_dict in cursor.matches(self.current_tree.root_node):
                    for capture_name, nodes in captures_dict.items():
                        # Map capture name to highlight group
                        group = capture_name.split('.')[-1]  # e.g., "keyword.function" -> "function"
                        color = HIGHLIGHT_GROUPS.get(group, 0)

                        for node in nodes:
                            start_row, start_col = node.start_point
                            end_row, end_col = node.end_point

                            for row in range(start_row, end_row + 1):
                                if row < start_line or row > end_line:
                                    continue

                                if row not in result:
                                    result[row] = []

                                s = start_col if row == start_row else 0
                                e = end_col if row == end_row else len(lines[row - start_line]) if (row - start_line) < len(lines) else 0

                                result[row].append((s, e, color))
            except Exception:
                pass
        else:
            # No query, use node types for basic highlighting
            return self._get_node_highlights(start_line, end_line, lines)

        return result

    def _get_node_highlights(self, start_line: int, end_line: int,
                              lines: List[str]) -> Dict[int, List[Tuple[int, int, int]]]:
        """Get highlights based on node types (fallback when no query)."""
        if not self.current_tree:
            return {}

        result: Dict[int, List[Tuple[int, int, int]]] = {}

        # Map common node types to highlight groups
        TYPE_COLORS = {
            'string': HIGHLIGHT_GROUPS['string'],
            'string_literal': HIGHLIGHT_GROUPS['string'],
            'comment': HIGHLIGHT_GROUPS['comment'],
            'number': HIGHLIGHT_GROUPS['number'],
            'integer': HIGHLIGHT_GROUPS['number'],
            'float': HIGHLIGHT_GROUPS['number'],
            'identifier': 0,  # default
            'keyword': HIGHLIGHT_GROUPS['keyword'],
            'type_identifier': HIGHLIGHT_GROUPS['type'],
            'primitive_type': HIGHLIGHT_GROUPS['type'],
            'function_definition': HIGHLIGHT_GROUPS['function'],
            'call_expression': HIGHLIGHT_GROUPS['function'],
        }

        def visit(node):
            row, col = node.start_point
            end_row, end_col = node.end_point

            if end_row < start_line or row > end_line:
                return

            node_type = node.type
            if node_type in TYPE_COLORS and TYPE_COLORS[node_type] != 0:
                for r in range(max(row, start_line), min(end_row, end_line) + 1):
                    if r not in result:
                        result[r] = []
                    s = col if r == row else 0
                    e = end_col if r == end_row else len(lines[r - start_line]) if (r - start_line) < len(lines) else 0
                    result[r].append((s, e, TYPE_COLORS[node_type]))

            for child in node.children:
                visit(child)

        visit(self.current_tree.root_node)
        return result

    def _get_simple_highlights(self, start_line: int, end_line: int,
                                lines: List[str]) -> Dict[int, List[Tuple[int, int, int]]]:
        """Simple regex-based highlighting for large files."""
        if self.current_lang not in SIMPLE_PATTERNS:
            return {}

        import re
        patterns = SIMPLE_PATTERNS[self.current_lang]
        keywords = set(patterns.get('keywords', []))
        comment_char = patterns.get('comment', '')
        result: Dict[int, List[Tuple[int, int, int]]] = {}

        for i, line in enumerate(lines):
            row = start_line + i
            highlights = []
            working_line = line

            # Find comments
            if comment_char and comment_char in line:
                idx = line.find(comment_char)
                highlights.append((idx, len(line), HIGHLIGHT_GROUPS['comment']))
                working_line = line[:idx]  # Don't highlight keywords in comments

            # Find strings (simple - doesn't handle escapes)
            in_string = False
            string_char = None
            string_start = 0
            for j, ch in enumerate(working_line):
                if not in_string and ch in ('"', "'"):
                    in_string = True
                    string_char = ch
                    string_start = j
                elif in_string and ch == string_char:
                    highlights.append((string_start, j + 1, HIGHLIGHT_GROUPS['string']))
                    in_string = False

            # Find keywords (word boundaries)
            for keyword in keywords:
                for match in re.finditer(r'\b' + re.escape(keyword) + r'\b', working_line):
                    highlights.append((match.start(), match.end(), HIGHLIGHT_GROUPS['keyword']))

            if highlights:
                result[row] = highlights

        return result

    def set_large_file_mode(self, num_lines: int) -> None:
        """Enable/disable simple mode based on file size."""
        self.use_simple = num_lines > self.LARGE_FILE_THRESHOLD


# Singleton instance
_highlighter: Optional[SyntaxHighlighter] = None

def get_highlighter() -> SyntaxHighlighter:
    """Get or create the global highlighter instance."""
    global _highlighter
    if _highlighter is None:
        _highlighter = SyntaxHighlighter()
    return _highlighter
