"""Tests for bex/grammar.py — canonical AST."""

import pytest
from bex.grammar import (
    Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
    parse_sore, render_sore, alphabet, count_words, model_cost,
    lang_size, match,
)


# ---------------------------------------------------------------------------
# Node construction and equality
# ---------------------------------------------------------------------------

class TestNodes:
    def test_symbol_eq(self):
        assert Symbol('a') == Symbol('a')
        assert Symbol('a') != Symbol('b')

    def test_concat_eq(self):
        assert Concat([Symbol('a'), Symbol('b')]) == Concat([Symbol('a'), Symbol('b')])
        assert Concat([Symbol('a')]) != Concat([Symbol('b')])

    def test_alt_eq(self):
        assert Alt([Symbol('a'), Symbol('b')]) == Alt([Symbol('a'), Symbol('b')])

    def test_plus_eq(self):
        assert Plus(Symbol('a')) == Plus(Symbol('a'))
        assert Plus(Symbol('a')) != Plus(Symbol('b'))

    def test_optional_eq(self):
        assert Optional(Symbol('a')) == Optional(Symbol('a'))

    def test_star_eq(self):
        assert Star(Symbol('a')) == Star(Symbol('a'))

    def test_epsilon_eq(self):
        assert Epsilon() == Epsilon()
        assert Epsilon() != Empty()

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

    def test_hash(self):
        s = {Symbol('a'), Symbol('b'), Symbol('a')}
        assert len(s) == 2


# ---------------------------------------------------------------------------
# Parse + render roundtrip
# ---------------------------------------------------------------------------

class TestRoundtrip:
    @pytest.mark.parametrize('sore', [
        'a',
        'a.b',
        'a.b.c',
        '(a+b)',
        '(a+b+c)',
        'a+',
        'a?',
        'a*',
        'a.b+',
        'a.b?.c+',
        '(a+b).c',
        '(a+b)+',
        '(a+b)?',
        '(a+b)*',
        'ε',
        '∅',
    ])
    def test_roundtrip(self, sore):
        node = parse_sore(sore)
        rendered = render_sore(node)
        assert rendered == sore, f'{sore} -> {rendered}'

    def test_dotted_symbol_roundtrip(self):
        node = parse_sore('foo\\.bar.baz+')
        rendered = render_sore(node)
        assert rendered == 'foo\\.bar.baz+'
        assert node == Concat([Symbol('foo.bar'), Plus(Symbol('baz'))])

    def test_deeply_nested(self):
        sore = '((a+b)+.c)?'
        node = parse_sore(sore)
        rendered = render_sore(node)
        assert rendered == sore


# ---------------------------------------------------------------------------
# Parse edge cases
# ---------------------------------------------------------------------------

class TestParse:
    def test_empty_string(self):
        assert parse_sore('') == Empty()

    def test_epsilon(self):
        assert parse_sore('ε') == Epsilon()

    def test_empty_symbol(self):
        assert parse_sore('∅') == Empty()

    def test_single_symbol(self):
        assert parse_sore('mockk') == Symbol('mockk')

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

    def test_alt(self):
        assert parse_sore('(a+b+c)') == Alt([Symbol('a'), Symbol('b'), Symbol('c')])

    def test_plus_outside_parens(self):
        assert parse_sore('a+') == Plus(Symbol('a'))

    def test_optional(self):
        assert parse_sore('a?') == Optional(Symbol('a'))

    def test_star(self):
        assert parse_sore('a*') == Star(Symbol('a'))

    def test_compound_quantifier_plus_optional(self):
        # a+? = a* (one or more, optional = zero or more)
        node = parse_sore('a+?')
        assert isinstance(node, Star)
        assert node.child == Symbol('a')

    def test_compound_quantifier_question_plus(self):
        # a?+ = a+ (optional then one or more = one or more)
        node = parse_sore('a?+')
        assert isinstance(node, Plus)
        assert node.child == Symbol('a')

    def test_dotted_symbol(self):
        node = parse_sore('foo\\.bar')
        assert node == Symbol('foo.bar')

    def test_dotted_in_concat(self):
        node = parse_sore('foo\\.bar.baz')
        assert node == Concat([Symbol('foo.bar'), Symbol('baz')])


# ---------------------------------------------------------------------------
# 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_alt(self):
        assert alphabet(Alt([Symbol('a'), Symbol('b'), Symbol('c')])) == {'a', 'b', 'c'}

    def test_plus(self):
        assert alphabet(Plus(Symbol('a'))) == {'a'}

    def test_nested(self):
        g = parse_sore('(a.b+).c?')
        assert alphabet(g) == {'a', 'b', 'c'}

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

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

    def test_dotted_symbol(self):
        assert alphabet(Symbol('foo.bar')) == {'foo.bar'}


# ---------------------------------------------------------------------------
# Match
# ---------------------------------------------------------------------------

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

    def test_concat_match(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_match(self):
        g = Alt([Symbol('a'), Symbol('b')])
        assert match(g, ['a'])
        assert match(g, ['b'])
        assert not match(g, ['c'])

    def test_plus_match(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_match(self):
        g = Optional(Symbol('a'))
        assert match(g, ['a'])
        assert match(g, [])
        assert not match(g, ['b'])

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

    def test_complex_grammar(self):
        g = parse_sore('init+.capability+.(invoke+.request)?')
        assert match(g, ['init', 'capability', 'invoke', 'request'])
        assert match(g, ['init', 'init', 'capability', 'capability'])
        assert not match(g, ['capability'])

    def test_dotted_symbols(self):
        g = parse_sore('foo\\.bar.baz+')
        assert match(g, ['foo.bar', 'baz'])
        assert match(g, ['foo.bar', 'baz', 'baz'])
        assert not match(g, ['foo', 'bar', 'baz'])

    def test_epsilon(self):
        assert match(Epsilon(), [])
        assert not match(Epsilon(), ['a'])

    def test_empty(self):
        assert not match(Empty(), [])
        assert not match(Empty(), ['a'])


# ---------------------------------------------------------------------------
# Count words
# ---------------------------------------------------------------------------

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

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

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

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

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

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

    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


# ---------------------------------------------------------------------------
# Model cost and language size
# ---------------------------------------------------------------------------

class TestScoring:
    def test_model_cost_symbol(self):
        assert model_cost(Symbol('a')) == 1

    def test_model_cost_concat(self):
        assert model_cost(Concat([Symbol('a'), Symbol('b')])) == 2

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

    def test_lang_size_symbol(self):
        # a: words of length 0 = 0, length 1 = 1, total = 1
        assert lang_size(Symbol('a'), 1) == 1

    def test_lang_size_concat(self):
        # a.b: words of length 0 = 0, length 1 = 0, length 2 = 1
        assert lang_size(Concat([Symbol('a'), Symbol('b')]), 2) == 1