"""AST-structural pattern extraction → GBNF grammar generation.

Instead of flattening tree-sitter captures to sequences and re-inferring
patterns with BEX algorithms, this module walks the AST directly and
extracts structural patterns from function bodies.

The core insight: tree-sitter already understands code structure.
We don't need to throw that away and re-infer it.

Approach:
1. Parse file with tree-sitter
2. For each function, extract the "structural shape" of its body
   (the sequence of AST node types, with text replaced by placeholders)
3. Group functions by their structural shape
4. Generate GBNF rules from each group
"""

import os
import re
from collections import defaultdict
from pathlib import Path

import tree_sitter

from .code import EXTENSION_MAP, _load_grammar, _load_query, _find_method_bodies


def _get_parser(ext):
    """Get a tree-sitter parser for the given extension."""
    lang_name, module_name, func_name = EXTENSION_MAP[ext]
    mod = __import__(module_name)
    lang_obj = getattr(mod, func_name)()
    lang = tree_sitter.Language(lang_obj)
    return tree_sitter.Parser(lang)


def _node_text(node, code_bytes):
    """Get the text of a node, decoded from bytes."""
    return code_bytes[node.start_byte:node.end_byte].decode('utf-8', errors='replace')


def _classify_node(node):
    """Classify an AST node into a GBNF terminal/nonterminal category.

    Returns (category, text) where category is a string like:
    - "return", "if", "loop", "throw" (structural keywords)
    - "call" (function/method call)
    - "new" (constructor call)
    - "member" (member access: a.b)
    - "identifier" (bare identifier)
    - "string", "number", "boolean" (literals)
    - "assign" (assignment)
    - "expr" (complex expression)
    """
    t = node.type

    if t == 'return_statement':
        return ('return', 'return')
    if t in ('if_statement', 'if_else_statement'):
        return ('if', 'if')
    if t in ('for_statement', 'for_in_statement', 'while_statement', 'do_statement'):
        return ('loop', 'loop')
    if t in ('throw_statement', 'raise_statement'):
        return ('throw', 'throw')
    if t == 'try_statement':
        return ('try', 'try')

    if t == 'new_expression':
        ctor = node.child_by_field_name('constructor')
        if ctor:
            return ('new', _node_text(ctor, b''))
        return ('new', 'new')

    if t == 'call_expression':
        func = node.child_by_field_name('function')
        if func and func.type == 'member_expression':
            prop = func.child_by_field_name('property')
            if prop:
                return ('call', _node_text(prop, b''))
        elif func:
            return ('call', _node_text(func, b''))
        return ('call', 'call')

    if t == 'member_expression':
        prop = node.child_by_field_name('property')
        if prop:
            return ('member', _node_text(prop, b''))
        return ('member', 'member')

    if t == 'assignment_expression':
        return ('assign', '=')

    if t in ('identifier', 'property_identifier', 'shorthand_property_identifier'):
        return ('identifier', _node_text(node, b''))

    if t in ('string', 'string_fragment', 'template_string'):
        return ('string', '"..."')

    if t in ('number', 'float', 'integer'):
        return ('number', 'N')

    if t in ('true', 'false'):
        return ('boolean', 'true' if t == 'true' else 'false')

    if t in ('null', 'undefined', 'none'):
        return ('null', 'null')

    if t == 'as_expression':
        return ('cast', 'as')

    return ('expr', t)


def _extract_body_shape(body_node, code_bytes):
    """Extract the structural shape of a function body.

    Returns a list of (category, text) tuples representing the
    high-level structure of the function body.

    For simple functions (single return), this is just one entry.
    For complex functions, it's the sequence of structural statements.
    """
    shape = []
    for child in body_node.children:
        cat, text = _classify_node(child)
        if cat in ('return', 'if', 'loop', 'throw', 'try', 'assign'):
            # Structural statement — always include
            shape.append((cat, text))
        elif cat == 'expr':
            # Expression statement — might be a call
            # Check if it's actually a call expression
            if child.type == 'expression_statement' and child.children:
                inner = child.children[0]
                inner_cat, inner_text = _classify_node(inner)
                if inner_cat in ('call', 'new'):
                    shape.append((inner_cat, inner_text))
        # Skip pure declarations, imports, comments, etc.

    return shape


def _shape_key(shape):
    """Convert a shape to a hashable key for grouping."""
    return tuple(cat for cat, _ in shape)


def _generalize_group shapes):
    """Generalize a group of shapes into a GBNF rule.

    All shapes in the group have the same structural pattern.
    The varying parts are the specific identifiers/names.
    We replace those with wildcards in the GBNF.
    """
    if not shapes:
        return None

    # All shapes have the same categories, so take the first
    first = shapes[0]

    # Check if all shapes in the group are identical
    all_same = all(s == first for s in shapes)

    if all_same:
        # All functions follow the exact same pattern
        # Generate a specific GBNF rule
        tokens = []
        for cat, text in first:
            if cat in ('return', 'if', 'loop', 'throw', 'try'):
                tokens.append(f'"{text}"')
            elif cat == 'call':
                tokens.append(f'CALL')
            elif cat == 'new':
                tokens.append(f'NEW')
            elif cat == 'member':
                tokens.append(f'MEMBER')
            elif cat == 'assign':
                tokens.append(f'"="')
            elif cat == 'cast':
                tokens.append(f'"as"')
            elif cat == 'string':
                tokens.append(f'STRING')
            elif cat == 'number':
                tokens.append(f'NUMBER')
            elif cat == 'boolean':
                tokens.append(f'BOOL')
            elif cat == 'null':
                tokens.append(f'NULL')
            elif cat == 'identifier':
                tokens.append(f'IDENT')
            else:
                tokens.append(f'"{text}"')
        return ' '.join(tokens)
    else:
        # Functions have different specific names but same structure
        # Generate a generalized rule with alternatives
        # For now, use CALL/NEW/IDENT wildcards
        tokens = []
        for cat, text in first:
            if cat in ('return', 'if', 'loop', 'throw', 'try'):
                tokens.append(f'"{text}"')
            elif cat == 'call':
                tokens.append(f'CALL')
            elif cat == 'new':
                tokens.append(f'NEW')
            elif cat == 'member':
                tokens.append(f'MEMBER')
            elif cat == 'assign':
                tokens.append(f'"="')
            elif cat == 'cast':
                tokens.append(f'"as"')
            elif cat == 'string':
                tokens.append(f'STRING')
            elif cat == 'number':
                tokens.append(f'NUMBER')
            elif cat == 'boolean':
                tokens.append(f'BOOL')
            elif cat == 'null':
                tokens.append(f'NULL')
            elif cat == 'identifier':
                tokens.append(f'IDENT')
            else:
                tokens.append(f'"{text}"')
        return ' '.join(tokens)


def extract_patterns_from_file(file_path):
    """Extract structural patterns from a single file.

    Returns a dict: { pattern_key: [function_name, ...] }
    where pattern_key is the GBNF rule string.
    """
    ext = os.path.splitext(file_path)[1].lower()
    if ext not in EXTENSION_MAP:
        return {}

    try:
        parser = _get_parser(ext)
    except Exception:
        return {}

    try:
        with open(file_path, 'rb') as f:
            code_bytes = f.read()
    except Exception:
        return {}

    tree = parser.parse(code_bytes)

    # Find method bodies
    bodies = _find_method_bodies(tree)
    bodies.sort(key=lambda b: b.start_byte)

    # Extract shapes
    patterns = defaultdict(list)
    for body in bodies:
        # Find the parent function name
        parent = body.parent
        name = None
        if parent:
            name_node = parent.child_by_field_name('name')
            if name_node:
                name = _node_text(name_node, code_bytes)

        shape = _extract_body_shape(body, code_bytes)
        if shape:
            key = _shape_key(shape)
            patterns[key].append((name, shape))

    return patterns


def patterns_to_gbnf(patterns):
    """Convert extracted patterns to GBNF grammar rules.

    Args:
        patterns: dict from extract_patterns_from_file

    Returns:
        GBNF grammar string
    """
    rules = []

    # Base rules for wildcards
    rules.append('root ::= function')
    rules.append('function ::= RETURN CALL NEW MEMBER IDENT STRING NUMBER BOOL NULL')
    rules.append('CALL ::= IDENT')
    rules.append('NEW ::= "new" IDENT')
    rules.append('MEMBER ::= IDENT "." IDENT')
    rules.append('IDENT ::= [a-zA-Z_] [a-zA-Z0-9_]*')
    rules.append('STRING ::= "\\"" [^\\"\n]* "\\""')
    rules.append('NUMBER ::= [0-9]+ ("." [0-9]+)?')
    rules.append('BOOL ::= "true" | "false"')
    rules.append('NULL ::= "null" | "undefined"')
    rules.append('RETURN ::= "return"')
    rules.append('IF ::= "if"')
    rules.append('LOOP ::= "for" | "while"')
    rules.append('THROW ::= "throw"')
    rules.append('TRY ::= "try"')
    rules.append('CAST ::= "as"')
    rules.append('"=" ::= "="')

    # Generate rules for each pattern group
    for i, (key, funcs) in enumerate(patterns.items()):
        rule_name = f'pattern_{i}'
        gbnf_body = _generalize_group([shape for _, shape in funcs])
        if gbnf_body:
            rules.append(f'{rule_name} ::= {gbnf_body}')

    return '\n'.join(rules)


if __name__ == '__main__':
    import sys
    if len(sys.argv) < 2:
        print("Usage: python -m bex.ast_to_gbnf <file>")
        sys.exit(1)

    file_path = sys.argv[1]
    patterns = extract_patterns_from_file(file_path)

    print(f"Found {len(patterns)} pattern groups:")
    for key, funcs in patterns.items():
        names = [name for name, _ in funcs]
        print(f"  {key}: {names}")

    print()
    print("GBNF:")
    print(patterns_to_gbnf(patterns))