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

GBNF (GGML BNF) is the de-facto standard grammar format for grammar-constrained
LLM output. It supports: literals, concatenation, alternation, repetition (+, *, ?),
and grouping.

SORE operators map directly:
  SORE `.` (concat)  → GBNF implicit concat (space)
  SORE `|` (alt)     → GBNF `|`
  SORE `+` (plus)    → GBNF `+`
  SORE `?` (optional) → GBNF `?`
  SORE `*` (star)    → GBNF `*`
  SORE `()` (group)  → GBNF `()`
  SORE literal       → GBNF `"literal"`
"""


# ---------------------------------------------------------------------------
# SORE tokenizer
# ---------------------------------------------------------------------------

_TOKEN_CHARS = set('.|+?*()')


def _tokenize(sore):
    """Tokenize a SORE string into a list of (type, value) tuples.

    Literals are sequences of non-operator characters.
    """
    tokens = []
    i = 0
    n = len(sore)
    while i < n:
        ch = sore[i]
        if ch == '.':
            tokens.append(('DOT', ch))
            i += 1
        elif ch == '|':
            tokens.append(('PIPE', ch))
            i += 1
        elif ch == '+':
            tokens.append(('PLUS', ch))
            i += 1
        elif ch == '?':
            tokens.append(('QUESTION', ch))
            i += 1
        elif ch == '*':
            tokens.append(('STAR', ch))
            i += 1
        elif ch == '(':
            tokens.append(('LPAREN', ch))
            i += 1
        elif ch == ')':
            tokens.append(('RPAREN', ch))
            i += 1
        elif ch == 'ε':
            tokens.append(('EPSILON', 'ε'))
            i += 1
        elif ch == '∅':
            tokens.append(('EMPTY', '∅'))
            i += 1
        else:
            # Collect literal characters (until next operator or paren)
            start = i
            while i < n and sore[i] not in _TOKEN_CHARS and sore[i] not in 'ε∅':
                i += 1
            lit = sore[start:i]
            if lit:
                # Strip newlines/extra whitespace from symbol names
                lit = ' '.join(lit.split())
                tokens.append(('LITERAL', lit))
    return tokens


# ---------------------------------------------------------------------------
# SORE parser — produces a tree
# ---------------------------------------------------------------------------

class _Node:
    """AST node for SORE."""
    pass


class _Literal(_Node):
    def __init__(self, value):
        self.value = value
    def __repr__(self):
        return f'Lit({self.value!r})'


class _Concat(_Node):
    def __init__(self, parts):
        self.parts = parts
    def __repr__(self):
        return f'Concat({self.parts})'


class _Alt(_Node):
    def __init__(self, parts):
        self.parts = parts
    def __repr__(self):
        return f'Alt({self.parts})'


class _Plus(_Node):
    def __init__(self, child):
        self.child = child
    def __repr__(self):
        return f'Plus({self.child})'


class _Optional(_Node):
    def __init__(self, child):
        self.child = child
    def __repr__(self):
        return f'Optional({self.child})'


class _Star(_Node):
    def __init__(self, child):
        self.child = child
    def __repr__(self):
        return f'Star({self.child})'


class _Epsilon(_Node):
    def __repr__(self):
        return 'Epsilon()'


class _Empty(_Node):
    def __repr__(self):
        return 'Empty()'


class _Parser:
    """Recursive descent parser for SOREs.

    Handles the overloaded + operator:
      - (a+b+c)  → disjunction (inside parens)
      - r+        → one-or-more repetition (outside parens)
    """

    def __init__(self, tokens):
        self.tokens = tokens
        self.pos = 0
        self.paren_depth = 0

    def peek(self):
        if self.pos < len(self.tokens):
            return self.tokens[self.pos]
        return ('EOF', '')

    def consume(self, expected_type=None):
        tok = self.peek()
        if tok[0] == 'EOF':
            raise ValueError(f'Unexpected end of SORE, expected {expected_type}')
        if expected_type and tok[0] != expected_type:
            raise ValueError(f'Expected {expected_type}, got {tok}')
        self.pos += 1
        return tok

    def parse(self):
        """Parse full SORE."""
        node = self.parse_alternation()
        return node

    def parse_alternation(self):
        """Parse: concat (('|' | '+') concat)*  — + is alternation inside parens"""
        parts = [self.parse_concat()]
        while self.peek()[0] in ('PIPE', 'PLUS'):
            if self.peek()[0] == 'PLUS' and self.paren_depth == 0:
                break  # + outside parens is repetition, not alternation
            self.consume()
            parts.append(self.parse_concat())
        if len(parts) == 1:
            return parts[0]
        return _Alt(parts)

    def parse_concat(self):
        """Parse: repetition (('.' | LPAREN) repetition)*  — implicit concat"""
        parts = [self.parse_repetition()]
        while self.peek()[0] in ('DOT', 'LPAREN'):
            if self.peek()[0] == 'DOT':
                self.consume('DOT')
            # LPAREN = implicit concat (no separator)
            parts.append(self.parse_repetition())
        if len(parts) == 1:
            return parts[0]
        return _Concat(parts)

    def parse_repetition(self):
        """Parse: atom ('+' | '?' | '*')?

        Inside parens, + is alternation (consumed by parse_alternation),
        not repetition. Outside parens, always consume + as repetition.
        Handles compound: +?, +*, ?+, *+ etc.
        """
        node = self.parse_atom()
        if self.peek()[0] in ('PLUS', 'QUESTION', 'STAR'):
            if self.peek()[0] == 'PLUS' and self.paren_depth > 0:
                return node  # + inside parens is alternation, handled by caller
            op = self.consume()
            if op[0] == 'PLUS':
                node = _Plus(node)
            elif op[0] == 'QUESTION':
                node = _Optional(node)
            elif op[0] == 'STAR':
                node = _Star(node)
            # Handle compound repetition: +?, +*, ?+ etc.
            # Normalize: Optional(Plus(x)) → Star(x)
            if self.peek()[0] in ('PLUS', 'QUESTION', 'STAR'):
                if self.peek()[0] == 'PLUS' and self.paren_depth > 0:
                    return node
                op2 = self.consume()
                if op2[0] == 'QUESTION':
                    if isinstance(node, _Plus):
                        node = _Star(node.child)
                    else:
                        node = _Optional(node)
                elif op2[0] == 'STAR':
                    node = _Star(node.child if isinstance(node, (_Plus, _Optional)) else node)
                elif op2[0] == 'PLUS':
                    if isinstance(node, _Optional):
                        node = _Plus(node.child)
                    elif isinstance(node, (_Plus, _Star)):
                        node = node  # ++ is idempotent
                    else:
                        node = _Plus(node)
        return node

    def parse_atom(self):
        """Parse: literal | '(' alternation ')' | 'ε' | '∅'"""
        tok = self.peek()
        if tok[0] == 'EOF':
            return _Epsilon()
        if tok[0] == 'LITERAL':
            self.consume()
            return _Literal(tok[1])
        if tok[0] == 'EPSILON':
            self.consume()
            return _Epsilon()
        if tok[0] == 'EMPTY':
            self.consume()
            return _Empty()
        if tok[0] == 'LPAREN':
            self.consume('LPAREN')
            self.paren_depth += 1
            node = self.parse_alternation()
            self.consume('RPAREN')
            self.paren_depth -= 1
            return node
        raise ValueError(f'Unexpected token: {tok}')


def _parse_sore(sore):
    """Parse a SORE string into an AST."""
    tokens = _tokenize(sore)
    parser = _Parser(tokens)
    return parser.parse()


# ---------------------------------------------------------------------------
# AST → GBNF
# ---------------------------------------------------------------------------

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


def _node_to_gbnf(node, rule_counter):
    """Convert AST node to GBNF fragment string.

    Returns (gbnf_string, new_rule_counter, new_rules_list).
    new_rules_list contains any helper rules needed.
    """
    if isinstance(node, _Literal):
        # Escape special chars in the literal for GBNF
        escaped = node.value.replace('\\', '\\\\').replace('"', '\\"')
        return f'"{escaped}"', rule_counter, []

    if isinstance(node, _Epsilon):
        return '', rule_counter, []

    if isinstance(node, _Empty):
        return '', rule_counter, []

    if isinstance(node, _Concat):
        parts = []
        all_new_rules = []
        for child in node.parts:
            frag, rule_counter, new_rules = _node_to_gbnf(child, rule_counter)
            if isinstance(child, _Alt):
                frag = f'({frag})'
            parts.append(frag)
            all_new_rules.extend(new_rules)
        return ' '.join(p for p in parts if p), rule_counter, all_new_rules

    if isinstance(node, _Alt):
        parts = []
        all_new_rules = []
        for child in node.parts:
            frag, rule_counter, new_rules = _node_to_gbnf(child, rule_counter)
            parts.append(frag)
            all_new_rules.extend(new_rules)
        return ' | '.join(p for p in parts if p), rule_counter, all_new_rules

    if isinstance(node, _Plus):
        frag, rule_counter, new_rules = _node_to_gbnf(node.child, rule_counter)
        if _needs_group(node.child):
            return f'({frag})+', rule_counter, new_rules
        return f'{frag}+', rule_counter, new_rules

    if isinstance(node, _Optional):
        frag, rule_counter, new_rules = _node_to_gbnf(node.child, rule_counter)
        if _needs_group(node.child):
            return f'({frag})?', rule_counter, new_rules
        return f'{frag}?', rule_counter, new_rules

    if isinstance(node, _Star):
        frag, rule_counter, new_rules = _node_to_gbnf(node.child, rule_counter)
        if _needs_group(node.child):
            return f'({frag})*', rule_counter, new_rules
        return f'{frag}*', rule_counter, new_rules

    raise ValueError(f'Unknown node type: {type(node)}')


def to_gbnf(sore):
    """Convert a SORE string to a GBNF grammar string.

    Args:
        sore: A SORE string like 'raise.(ValueError)+' or 'mockk'

    Returns:
        A GBNF grammar string with a single 'root' rule.

    Example:
        >>> to_gbnf('mockk')
        'root ::= "mockk"'
        >>> to_gbnf('raise.(ValueError)+')
        'root ::= "raise" ("ValueError")+'
        >>> to_gbnf('assertEquals.(of.(assertFailsWith)?)+')
        'root ::= "assertEquals" ("of" ("assertFailsWith")?)+'
    """
    if not sore or sore == '∅':
        return 'root ::= ""'
    if sore == 'ε':
        return 'root ::= ""'

    tree = _parse_sore(sore)
    frag, _, _ = _node_to_gbnf(tree, 0)
    return f'root ::= {frag}'


def to_gbnf_with_rules(sore, name='root'):
    """Convert a SORE to GBNF with a named rule.

    Args:
        sore: A SORE string
        name: Rule name (default: 'root')

    Returns:
        GBNF rule string like 'my-rule ::= "foo" ("bar")+'
    """
    if not sore or sore == '∅':
        return f'{name} ::= ""'
    if sore == 'ε':
        return f'{name} ::= ""'

    tree = _parse_sore(sore)
    frag, _, _ = _node_to_gbnf(tree, 0)
    return f'{name} ::= {frag}'


def validate_sore(sore):
    """Check if a SORE string is parseable. Returns (True, None) or (False, error_msg)."""
    if not sore or sore in ('∅', 'ε'):
        return True, None
    try:
        _parse_sore(sore)
        return True, None
    except ValueError as e:
        return False, str(e)


def grammar_structure_score(sore):
    """Quantify how structured a SORE is (0.0 = flat bag, 1.0 = fully structured).

    Structured grammars have ordering (concatenation, optional, repetition)
    that tells you the SEQUENCE things happen. Flat bags just list symbols.
    """
    import re
    if not sore or sore in ('∅', 'ε'):
        return 0.0
    depth = 0
    dots = 0
    pluses_outside = 0
    questions = 0
    stars = 0
    for ch in sore:
        if ch == '(':
            depth += 1
        elif ch == ')':
            depth -= 1
        elif ch == '.' and depth == 0:
            dots += 1
        elif ch == '+' and depth == 0:
            pluses_outside += 1
        elif ch == '?' and depth == 0:
            questions += 1
        elif ch == '*' and depth == 0:
            stars += 1
    disj_parts = 0
    for m in re.finditer(r'\(([^)]+)\)', sore):
        inner = m.group(1)
        if '+' in inner:
            disj_parts = max(disj_parts, inner.count('+') + 1)
    symbols_only = re.sub(r'[.?*+()]', '', sore)
    sym_len = len(symbols_only)
    if sym_len == 0:
        return 0.0
    struct_ops = dots + questions + stars + pluses_outside
    struct_ratio = struct_ops / max(sym_len, 1)
    disj_ratio = disj_parts / max(sym_len, 1)
    score = min(1.0, struct_ratio * 3)
    if disj_ratio > 0.5 and dots == 0:
        score *= 0.3
    return score