"""GBNF — Convert AST to GBNF grammar format for llama.cpp constrained decoding."""

from .grammar import (
    Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
)


def _needs_group(node):
    """Check if a node needs parentheses in GBNF output."""
    return isinstance(node, (Alt, Concat))


def _node_to_gbnf(node):
    """Convert AST node to GBNF fragment string."""
    if isinstance(node, Symbol):
        escaped = node.value.replace('\\', '\\\\').replace('"', '\\"')
        return f'"{escaped}"'
    if isinstance(node, (Epsilon, Empty)):
        return ''
    if isinstance(node, Concat):
        parts = []
        for child in node.parts:
            frag = _node_to_gbnf(child)
            if isinstance(child, Alt):
                frag = f'({frag})'
            parts.append(frag)
        return ' '.join(p for p in parts if p)
    if isinstance(node, Alt):
        parts = [_node_to_gbnf(child) for child in node.parts]
        return ' | '.join(p for p in parts if p)
    if isinstance(node, Plus):
        frag = _node_to_gbnf(node.child)
        if _needs_group(node.child):
            return f'({frag})+'
        return f'{frag}+'
    if isinstance(node, Optional):
        frag = _node_to_gbnf(node.child)
        if _needs_group(node.child):
            return f'({frag})?'
        return f'{frag}?'
    if isinstance(node, Star):
        frag = _node_to_gbnf(node.child)
        if _needs_group(node.child):
            return f'({frag})*'
        return f'{frag}*'
    return ''


def to_gbnf(node):
    """Convert AST node to a GBNF grammar string.

    Args:
        node: Grammar AST node

    Returns:
        GBNF grammar string with a single 'root' rule.
    """
    if node is None or isinstance(node, Empty):
        return 'root ::= ""'
    if isinstance(node, Epsilon):
        return 'root ::= ""'
    frag = _node_to_gbnf(node)
    return f'root ::= {frag}'


def to_gbnf_with_rules(node, name='root'):
    """Convert AST to GBNF with a named rule."""
    if node is None or isinstance(node, Empty):
        return f'{name} ::= ""'
    if isinstance(node, Epsilon):
        return f'{name} ::= ""'
    frag = _node_to_gbnf(node)
    return f'{name} ::= {frag}'


def grammar_structure_score(node):
    """Quantify how structured an AST is (0.0 = flat bag, 1.0 = fully structured)."""
    if node is None or isinstance(node, Empty):
        return 0.0
    if isinstance(node, Symbol):
        return 0.0
    if isinstance(node, Epsilon):
        return 0.0
    if isinstance(node, (Plus, Optional, Star)):
        return 0.5 + 0.5 * grammar_structure_score(node.child)
    if isinstance(node, Alt):
        child_scores = [grammar_structure_score(p) for p in node.parts]
        return sum(child_scores) / max(len(child_scores), 1)
    if isinstance(node, Concat):
        n = len(node.parts)
        if n <= 1:
            return 1.0
        child_scores = [grammar_structure_score(p) for p in node.parts]
        return min(1.0, 0.3 + 0.2 * n + sum(child_scores) / max(len(child_scores), 1))
    return 0.0
