"""Tests for grammar AST, matching, counting, and GBNF rendering."""

import pytest
from bex.grammar import (
    Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
    alphabet, count_words, lang_size, model_cost, match,
)
from bex.gbnf import to_gbnf, to_gbnf_with_rules, grammar_structure_score


# ── AST Construction ──

class TestASTConstruction:
    def test_symbol(self):
        assert Symbol('a') == Symbol('a')
        assert Symbol('a') != Symbol('b')

    def test_concat(self):
        c = Concat([Symbol('a'), Symbol('b')])
        assert c.parts == [Symbol('a'), Symbol('b')]

    def test_alt(self):
        a = Alt([Symbol('a'), Symbol('b')])
        assert a.parts == [Symbol('a'), Symbol('b')]

    def test_plus(self):
        p = Plus(Symbol('a'))
        assert p.child == Symbol('a')

    def test_optional(self):
        o = Optional(Symbol('a'))
        assert o.child == Symbol('a')

    def test_star(self):
        s = Star(Symbol('a'))
        assert s.child == Symbol('a')

    def test_epsilon(self):
        assert Epsilon() == Epsilon()

    def test_empty(self):
        assert Empty() == Empty()


# ── Alphabet ──

class TestAlphabet:
    def test_symbol(self):
        assert alphabet(Symbol('a')) == {'a'}

    def test_concat(self):
        assert alphabet(Concat([Symbol('a'), Symbol('b')])) == {'a', 'b'}

    def test_nested(self):
        g = Concat([Symbol('a'), Plus(Alt([Symbol('b'), Symbol('c')]))])
        assert alphabet(g) == {'a', 'b', 'c'}


# ── Matching ──

class TestMatch:
    def test_symbol(self):
        assert match(Symbol('a'), ['a'])
        assert not match(Symbol('a'), ['b'])
        assert not match(Symbol('a'), [])

    def test_concat(self):
        g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
        assert match(g, ['a', 'b', 'c'])
        assert not match(g, ['a', 'b'])
        assert not match(g, ['a', 'b', 'c', 'd'])

    def test_alt(self):
        g = Alt([Symbol('a'), Symbol('b')])
        assert match(g, ['a'])
        assert match(g, ['b'])
        assert not match(g, ['c'])

    def test_plus(self):
        g = Plus(Symbol('a'))
        assert match(g, ['a'])
        assert match(g, ['a', 'a'])
        assert match(g, ['a', 'a', 'a'])
        assert not match(g, [])
        assert not match(g, ['b'])

    def test_optional(self):
        g = Optional(Symbol('a'))
        assert match(g, ['a'])
        assert match(g, [])
        assert not match(g, ['b'])

    def test_star(self):
        g = Star(Symbol('a'))
        assert match(g, [])
        assert match(g, ['a'])
        assert match(g, ['a', 'a'])
        assert not match(g, ['b'])

    def test_complex_grammar(self):
        g = Concat([Plus(Symbol('init')), Plus(Symbol('capability')),
                     Concat([Plus(Symbol('invoke')), Symbol('request')])])
        assert match(g, ['init', 'capability', 'invoke', 'request'])
        assert match(g, ['init', 'init', 'capability', 'capability', 'invoke', 'request'])
        assert not match(g, ['init'])


# ── Counting ──

class TestCountWords:
    def test_symbol(self):
        assert count_words(Symbol('a'), 1) == 1
        assert count_words(Symbol('a'), 0) == 0

    def test_epsilon(self):
        assert count_words(Epsilon(), 0) == 1
        assert count_words(Epsilon(), 1) == 0

    def test_empty(self):
        assert count_words(Empty(), 0) == 0

    def test_concat(self):
        g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
        assert count_words(g, 3) == 1
        assert count_words(g, 2) == 0

    def test_alt(self):
        g = Alt([Symbol('a'), Symbol('b'), Symbol('c')])
        assert count_words(g, 1) == 3

    def test_plus(self):
        g = Plus(Symbol('a'))
        assert count_words(g, 1) == 1
        assert count_words(g, 2) == 1
        assert count_words(g, 0) == 0

    def test_optional(self):
        g = Optional(Symbol('a'))
        assert count_words(g, 0) == 1
        assert count_words(g, 1) == 1

    def test_star(self):
        g = Star(Symbol('a'))
        assert count_words(g, 0) == 1
        assert count_words(g, 1) == 1

    def test_disjunction_plus(self):
        g = Plus(Alt([Symbol('a'), Symbol('b'), Symbol('c')]))
        assert count_words(g, 1) == 3
        assert count_words(g, 2) == 9

    def test_info_plus(self):
        g = Plus(Symbol('info'))
        for l in range(1, 8):
            assert count_words(g, l) == 1


# ── Model cost ──

class TestModelCost:
    def test_symbol(self):
        assert model_cost(Symbol('a')) == 1

    def test_concat(self):
        assert model_cost(Concat([Symbol('a'), Symbol('b'), Symbol('c')])) == 3

    def test_plus(self):
        assert model_cost(Plus(Symbol('a'))) == 1

    def test_nested(self):
        g = Concat([Symbol('a'), Plus(Alt([Symbol('b'), Symbol('c')]))])
        assert model_cost(g) == 3


# ── GBNF Rendering ──

def test_to_gbnf():
    g = Plus(Symbol('init'))
    result = to_gbnf(g)
    assert result == 'root ::= "init"+'


def test_to_gbnf_with_rules():
    g = Alt([Symbol('a'), Symbol('b')])
    result = to_gbnf_with_rules(g, name='choice')
    assert result == 'choice ::= "a" | "b"'


def test_count_words_concat_is_memoized():
    """Regression: _count_concat must be memoized.

    An uncached implementation revisits (remaining_parts, length) states
    exponentially; a long sequence over a multi-part Concat explodes into
    millions of calls. With memoization it completes instantly.
    """
    import time
    # length-L string of (a|b) alternatives -> exactly 2^L words of length L
    L = 20
    g = Concat([Alt([Symbol('a'), Symbol('b')]) for _ in range(L)])
    t0 = time.time()
    n = count_words(g, L)
    dt = time.time() - t0
    assert dt < 2.0, f"count_words was too slow ({dt:.2f}s) — memoization lost"
    assert n == 2 ** L, f"expected 2^{L}={2**L} words, got {n}"


# ── Structure score ──

class TestStructureScore:
    def test_empty(self):
        assert grammar_structure_score(Empty()) == 0.0

    def test_symbol(self):
        assert grammar_structure_score(Symbol('a')) == 0.0

    def test_concat(self):
        g = Concat([Symbol('a'), Symbol('b'), Symbol('c')])
        assert grammar_structure_score(g) > 0.0

    def test_plus(self):
        g = Plus(Symbol('a'))
        assert grammar_structure_score(g) > 0.0