"""Universal tree-sitter tag preprocessor.

Usage:
    python -m bex.tag_preprocessor.code <file>

Emits an ordered sequence of behavioral tokens using community highlights.scm
queries from nvim-treesitter. One code path for all languages.
"""

import importlib
import os
import re
import sys
from pathlib import Path



from tree_sitter import Language, Parser, Query, QueryCursor

QUERIES_DIR = Path(__file__).parent / "queries"

# (query_name, module_name, func_name)
EXTENSION_MAP = {
    ".py": ("python", "tree_sitter_python", "language"),
    ".js": ("javascript", "tree_sitter_javascript", "language"),
    ".mjs": ("javascript", "tree_sitter_javascript", "language"),
    ".cjs": ("javascript", "tree_sitter_javascript", "language"),
    ".ts": ("typescript", "tree_sitter_typescript", "language_typescript"),
    ".tsx": ("typescript", "tree_sitter_typescript", "language_tsx"),
    ".rb": ("ruby", "tree_sitter_ruby", "language"),
    ".go": ("go", "tree_sitter_go", "language"),
    ".rs": ("rust", "tree_sitter_rust", "language"),
    ".java": ("java", "tree_sitter_java", "language"),
    ".kt": ("kotlin", "tree_sitter_kotlin", "language"),
    ".kts": ("kotlin", "tree_sitter_kotlin", "language"),
    ".c": ("c", "tree_sitter_c", "language"),
    ".h": ("c", "tree_sitter_c", "language"),
    ".cpp": ("cpp", "tree_sitter_cpp", "language"),
    ".cc": ("cpp", "tree_sitter_cpp", "language"),
    ".cxx": ("cpp", "tree_sitter_cpp", "language"),
    ".hpp": ("cpp", "tree_sitter_cpp", "language"),
}

BEHAVIORAL_PREFIXES = (
    "definition.",
    "reference.",
    "keyword.",
    "function",
    "attribute",
    "constructor",
    "label",
    "type.definition",
    "module",
    # Bare captures from Kotlin highlights.scm
    # (Kotlin uses "conditional" not "keyword.conditional", etc.)
    # Excluded: "variable" (too noisy — every identifier reference)
    "conditional",
    "exception",
    "repeat",
    "property",
    "type",
)

CALL_PREFIXES = ("function.call", "function.method.call", "reference.call", "reference.class", "constructor", "function")


def coarsen_token(capname):
    """Map a tree-sitter capture name to a language-agnostic category.

    Uses capture names from highlights.scm — same across all languages.
    Function calls are kept as raw text (they ARE the behavioral content).
    Only the 4 highest-signal structural tokens are coarsened:
    RETURN, IF, EXCEPTION, LOOP. Everything else kept raw.
    Handles both dotted (keyword.conditional) and bare (conditional) captures.
    """
    if capname.startswith("function.call") or capname.startswith("reference.call"):
        return None  # keep raw
    if capname.startswith("function.method.call"):
        return None  # keep raw
    if capname.startswith("constructor"):
        return None  # keep raw
    if capname.startswith("function"):
        return None  # keep raw (bare @function in Kotlin etc.)
    if capname.startswith("keyword.return") or capname == "return":
        return "RETURN"
    if capname.startswith("keyword.raise") or capname.startswith("keyword.throw") or capname == "raise":
        return "RAISE"
    if capname.startswith("keyword.yield") or capname == "yield":
        return "YIELD"
    if capname.startswith("keyword.exception") or capname == "exception":
        return "EXCEPTION"
    if capname.startswith("keyword.conditional") or capname == "conditional":
        return "IF"
    if capname.startswith("keyword.repeat") or capname == "repeat":
        return "LOOP"
    return None  # keep everything else raw


def _extract_coarsened_tokens(seq):
    """Extract coarsened tokens from a method sequence.

    Function calls are kept as raw text (they are the behavioral content).
    Structural tokens (keywords, types) are mapped to language-agnostic
    categories via coarsen_token().
    Returns list of strings: either raw text or category label.
    """
    result = []
    for capname, text, _ in seq:
        for prefix in BEHAVIORAL_PREFIXES:
            if capname.startswith(prefix):
                cat = coarsen_token(capname)
                result.append(text if cat is None else cat)
                break
    return result

def _extract_call_tokens(seq):
    """Extract ordered call-like tokens from a method sequence.

    Filters to captures representing function calls, constructors,
    or references — the 'what happens in what order'.
    Returns list of text values.
    """
    result = []
    for capname, text, _ in seq:
        for prefix in CALL_PREFIXES:
            if capname.startswith(prefix):
                result.append(text)
                break
    return result


ARG_LITERAL_TYPES = {
    "string", "string_literal", "number", "integer", "float",
    "decimal", "true", "false", "null", "nil", "none",
}

LAMBDA_TYPES = {
    "lambda", "do_block", "block",
}


def _classify_arg_node(node):
    t = node.type
    if t in ARG_LITERAL_TYPES:
        return "lit"
    if t in LAMBDA_TYPES or "ambda" in t or "block" in t:
        return "lambda"
    if t == "keyword_argument" or t.endswith("named_argument"):
        return "kwarg"
    if t.endswith("call_expression") or t in ("call", "method_invocation"):
        return "call"
    if t.endswith("identifier") or t.endswith("name"):
        return "var"
    if "subscript" in t:
        return "subscript"
    if "binary" in t or "unary" in t or "ternary" in t or "operator" in t:
        return "expr"
    if t in ("interpolation", "template_string"):
        return "template"
    return "other"


def _find_arglist_node(parent):
    args = parent.child_by_field_name("arguments")
    if args:
        return args
    for child in parent.children:
        if child.type in ("argument_list", "arguments"):
            return child
    return None


def _iterate_arg_nodes(arglist):
    for child in arglist.children:
        if child.is_named:
            yield child


def extract_arg_info(file_path, code):
    ext = os.path.splitext(file_path)[1].lower()
    lang, query_name = _load_grammar(ext)
    query_src = _load_query(query_name)

    code_bytes = code.encode()
    parser = Parser(lang)
    tree = parser.parse(code_bytes)
    query = Query(lang, query_src)

    cursor = QueryCursor(query)
    captures = cursor.captures(tree.root_node)

    info = {}
    for capname, nodes in captures.items():
        if not capname.startswith(BEHAVIORAL_PREFIXES):
            continue
        for node in nodes:
            text = code_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace")
            parent = node.parent
            if not parent:
                continue
            arglist = _find_arglist_node(parent)
            if not arglist:
                continue
            types = [_classify_arg_node(c) for c in _iterate_arg_nodes(arglist)]
            info.setdefault(text, []).append((len(types), tuple(types)))
    return info


def _summarize_arg_info(info):
    from collections import Counter
    summary = {}
    for call_name, observations in info.items():
        counts = [c for c, _ in observations]
        pattern_counts = Counter(observations)
        top_patterns = pattern_counts.most_common(5)
        summary[call_name] = {
            "occurrences": len(observations),
            "arg_count": {
                "min": min(counts),
                "max": max(counts),
                "common": max(set(counts), key=counts.count),
            },
            "patterns": [
                {"count": c, "args": n, "types": list(t)}
                for (n, t), c in top_patterns
            ],
        }
    return summary


_grammar_cache = {}
_query_cache = {}


def _load_grammar(ext):
    entry = EXTENSION_MAP.get(ext)
    if entry is None:
        raise ValueError(f"Unsupported extension: {ext}")
    query_name, module_name, func_name = entry

    cache_key = f"{module_name}.{func_name}"
    if cache_key in _grammar_cache:
        return _grammar_cache[cache_key], query_name

    mod = importlib.import_module(module_name)
    lang = Language(getattr(mod, func_name)())
    _grammar_cache[cache_key] = lang
    return lang, query_name


INHERIT_RE = re.compile(r"^;\s*inherits:\s*(.+)$", re.MULTILINE)


def _resolve_inherits(src, query_name, seen=None):
    if seen is None:
        seen = set()
    if query_name in seen:
        return ""
    seen.add(query_name)

    qpath = QUERIES_DIR / f"{query_name}.scm"
    if not qpath.exists():
        return ""

    content = qpath.read_text()
    m = INHERIT_RE.search(content)
    if m:
        parents = m.group(1)
        parent_parts = []
        for parent in parents.split(","):
            parent = parent.strip().strip("()")
            if parent:
                parent_parts.append(_resolve_inherits(parent, parent, seen))
        parent_src = "\n".join(p for p in parent_parts if p)
        body = INHERIT_RE.sub("", content)
        return parent_src + "\n" + body if parent_src else body
    return content


def _load_query(query_name):
    if query_name in _query_cache:
        return _query_cache[query_name]

    src = _resolve_inherits(query_name, query_name)
    if not src:
        raise FileNotFoundError(f"Query file not found: {query_name}")

    _query_cache[query_name] = src
    return src


def _find_method_bodies(tree):
    """Find function/method body nodes across all languages.

    Uses tree-sitter's named field 'body' (universal across 9/10 grammars).
    Kotlin fallback: scan children for body-like types when the field is absent.
    Parent filter: 'function' or 'method' in type name avoids class/loop bodies.
    """
    bodies = []

    def walk(node):
        body = node.child_by_field_name("body")
        if not body:
            for child in node.children:
                ctype = child.type.lower()
                if "body" in ctype or "block" in ctype or ctype == "compound_statement":
                    body = child
                    break
        if body:
            ptype = node.type.lower()
            if "function" in ptype or "method" in ptype:
                bodies.append(body)
        for child in node.children:
            walk(child)

    walk(tree.root_node)
    return bodies


def preprocess_by_method(file_path: str, code: str):
    """Preprocess and group behavioral tokens by enclosing method body.

    Returns list of sequences, one per function/method found.
    Each sequence is [(capture_name, text, line_number), ...].
    """
    ext = os.path.splitext(file_path)[1].lower()
    lang, query_name = _load_grammar(ext)
    query_src = _load_query(query_name)

    code_bytes = code.encode()
    parser = Parser(lang)
    tree = parser.parse(code_bytes)

    query = Query(lang, query_src)

    cursor = QueryCursor(query)
    captures = cursor.captures(tree.root_node)

    items = []
    for capname, nodes in captures.items():
        if not capname.startswith(BEHAVIORAL_PREFIXES):
            continue
        for node in nodes:
            text = code_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace")
            if text:
                items.append((node.start_byte, capname, node, text))

    items.sort(key=lambda x: x[0])

    method_bodies = _find_method_bodies(tree)
    method_bodies.sort(key=lambda b: b.start_byte)

    sequences = []
    for body_node in method_bodies:
        seq = []
        for start, capname, node, text in items:
            if body_node.start_byte <= start < body_node.end_byte:
                seq.append((capname, text, code_bytes[:start].count(b"\n") + 1))
        if seq:
            sequences.append(seq)

    return sequences


def preprocess(file_path: str, code: str):
    ext = os.path.splitext(file_path)[1].lower()
    lang, query_name = _load_grammar(ext)
    query_src = _load_query(query_name)

    code_bytes = code.encode()
    parser = Parser(lang)
    tree = parser.parse(code_bytes)

    try:
        query = Query(lang, query_src)
    except Exception as e:
        print(f"Query error for {query_name}: {e}", file=sys.stderr)
        return

    cursor = QueryCursor(query)
    captures = cursor.captures(tree.root_node)

    items = []
    for capname, nodes in captures.items():
        if not capname.startswith(BEHAVIORAL_PREFIXES):
            continue
        for node in nodes:
            text = code_bytes[node.start_byte:node.end_byte].decode("utf-8", errors="replace")
            if text:
                items.append((node.start_byte, capname, node, text))

    items.sort(key=lambda x: x[0])

    return [(capname, text, code_bytes[:start].count(b"\n") + 1) for start, capname, _, text in items]


def main():
    if len(sys.argv) < 2:
        print("Usage: python -m bex.tag-preprocessor.code <file>", file=sys.stderr)
        sys.exit(1)

    file_path = sys.argv[1]

    if not os.path.exists(file_path):
        print(f"File not found: {file_path}", file=sys.stderr)
        sys.exit(1)

    try:
        with open(file_path) as f:
            code = f.read()
    except Exception as e:
        print(f"Error reading {file_path}: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        seq = preprocess(file_path, code)
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except FileNotFoundError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error preprocessing {file_path}: {e}", file=sys.stderr)
        sys.exit(1)

    if not seq:
        print("(no behavioral tokens)")
        return

    print(f"{'CAPTURE':35s} {'TEXT':50s} LINE")
    print("-" * 88)
    for capname, text, line in seq:
        print(f"{capname:35s} '{text[:48]:48s}' L{line}")

    print(f"\nSequence ({len(seq)} tokens):")
    print(" -> ".join(c for c, _, _ in seq))


if __name__ == "__main__":
    main()
