import unittest
from lark import UnexpectedToken, UnexpectedCharacters, Token, Tree, UnexpectedEOF

from slrquerytester.utils.query_parser import parse_with_semantic_check
from slrquerytester.translators.base_transformer import BaseTransformer


class TestQueryParseTreeStructure(unittest.TestCase):

    def test_pattern_1_single_term(self):
        def test_pattern_1_single_term(self):
            query = "title:Hello"
            tree = parse_with_semantic_check(query)

            # start -> or_expr
            self.assertIsInstance(tree, Tree)
            self.assertEqual(tree.data, 'start')
            self.assertEqual(len(tree.children), 1)
            or_expr_node = tree.children[0]
            self.assertEqual(or_expr_node.data, 'or_expr')

            # or_expr -> single child and_expr
            self.assertEqual(len(or_expr_node.children), 1)
            and_expr_node = or_expr_node.children[0]
            self.assertEqual(and_expr_node.data, 'and_expr')

            # and_expr -> single child field_term
            self.assertEqual(len(and_expr_node.children), 1)
            field_term_node = and_expr_node.children[0]
            self.assertEqual(field_term_node.data, 'field_term')

            # field_term -> single child normal_expr
            self.assertEqual(len(field_term_node.children), 1)
            normal_expr_node = field_term_node.children[0]
            self.assertEqual(normal_expr_node.data, 'normal_expr')

            # normal_expr -> single child normal_term
            self.assertEqual(len(normal_expr_node.children), 1)
            normal_term_node = normal_expr_node.children[0]
            self.assertEqual(normal_term_node.data, 'normal_term')

            # normal_term -> 3 tokens: NORMAL_FIELD, TEXT_OP, VALUE
            self.assertEqual(len(normal_term_node.children), 3)
            field_token = normal_term_node.children[0]
            colon_token = normal_term_node.children[1]
            value_token = normal_term_node.children[2]

            self.assertIsInstance(field_token, Token)
            self.assertEqual(field_token.type, 'NORMAL_FIELD')
            self.assertEqual(field_token.value, 'title')

            self.assertIsInstance(colon_token, Token)
            self.assertEqual(colon_token.type, 'TEXT_OP')
            self.assertEqual(colon_token.value, ':')

            self.assertIsInstance(value_token, Token)
            self.assertEqual(value_token.type, 'VALUE')
            self.assertEqual(value_token.value, 'Hello')

    def test_pattern_2_quoted_phrase(self):
        query = 'title:"Hello World"'
        tree = parse_with_semantic_check(query)

        # start -> or_expr -> [and_expr] -> [field_term] -> [normal_expr] -> [normal_term]
        self.assertEqual(tree.data, 'start')
        normal_term_node = tree.children[0].children[0].children[0].children[0].children[0]

        # normal_term -> 3 tokens: NORMAL_FIELD, TEXT_OP, VALUE
        value_token = normal_term_node.children[2]

        self.assertEqual(value_token.type, 'VALUE')
        self.assertEqual(value_token.value, '"Hello World"')

    def test_pattern_3_and_expression(self):
        query = "title:Hello AND author:Bob"
        tree = parse_with_semantic_check(query)

        # start -> or_expr -> single and_expr
        self.assertIsInstance(tree, Tree)
        self.assertEqual(tree.data, 'start')
        self.assertEqual(len(tree.children), 1)
        or_expr_node = tree.children[0]
        self.assertEqual(or_expr_node.data, 'or_expr')
        self.assertEqual(len(or_expr_node.children), 1)
        and_expr_node = or_expr_node.children[0]
        self.assertEqual(and_expr_node.data, 'and_expr')
        self.assertEqual(len(and_expr_node.children), 3)

        # and_expr -> field_term, AND token, field_term
        first_field_term = and_expr_node.children[0]
        and_tok = and_expr_node.children[1]
        second_field_term = and_expr_node.children[2]

        self.assertEqual(first_field_term.data, 'field_term')
        self.assertIsInstance(and_tok, Token)
        self.assertEqual(and_tok.type, 'AND')
        self.assertEqual(second_field_term.data, 'field_term')

        # Each field_term -> normal_expr -> normal_term
        first_normal_expr = first_field_term.children[0]
        second_normal_expr = second_field_term.children[0]
        first_normal_term = first_normal_expr.children[0]
        second_normal_term = second_normal_expr.children[0]

        self.assertEqual(first_normal_term.data, 'normal_term')
        self.assertEqual(second_normal_term.data, 'normal_term')

        # normal_term -> [NORMAL_FIELD, TEXT_OP, VALUE]
        self.assertEqual(first_normal_term.children[0].value, 'title')
        self.assertEqual(first_normal_term.children[1].value, ':')
        self.assertEqual(first_normal_term.children[2].value, 'Hello')

        self.assertEqual(second_normal_term.children[0].value, 'author')
        self.assertEqual(second_normal_term.children[1].value, ':')
        self.assertEqual(second_normal_term.children[2].value, 'Bob')

    def test_pattern_4_or_expression(self):
        query = 'title:"Hello World" OR abstract:World'
        tree = parse_with_semantic_check(query)

        # start -> or_expr
        self.assertIsInstance(tree, Tree)
        self.assertEqual(tree.data, 'start')
        self.assertEqual(len(tree.children), 1)
        or_expr_node = tree.children[0]
        self.assertEqual(or_expr_node.data, 'or_expr')

        # or_expr -> and_expr, OR token, and_expr
        self.assertEqual(len(or_expr_node.children), 3)
        first_and_expr = or_expr_node.children[0]
        or_tok = or_expr_node.children[1]
        second_and_expr = or_expr_node.children[2]
        self.assertEqual(first_and_expr.data, 'and_expr')
        self.assertIsInstance(or_tok, Token)
        self.assertEqual(or_tok.type, 'OR')
        self.assertEqual(second_and_expr.data, 'and_expr')

        # each and_expr -> single field_term
        first_field_term = first_and_expr.children[0]
        second_field_term = second_and_expr.children[0]

        # field_term -> normal_expr -> normal_term
        first_normal_expr = first_field_term.children[0]
        second_normal_expr = second_field_term.children[0]
        first_normal_term = first_normal_expr.children[0]
        second_normal_term = second_normal_expr.children[0]
        self.assertEqual(first_normal_term.data, 'normal_term')
        self.assertEqual(second_normal_term.data, 'normal_term')

        self.assertEqual(first_normal_term.children[0].value, 'title')
        self.assertEqual(first_normal_term.children[1].value, ':')
        self.assertEqual(first_normal_term.children[2].value, '"Hello World"')

        self.assertEqual(second_normal_term.children[0].value, 'abstract')
        self.assertEqual(second_normal_term.children[1].value, ':')
        self.assertEqual(second_normal_term.children[2].value, 'World')

    def test_pattern_5_parentheses_nesting(self):
        query = '(title:Hello OR abstract:World) AND keywords:"Deep Learning"'
        tree = parse_with_semantic_check(query)

        # start -> or_expr -> single and_expr
        self.assertIsInstance(tree, Tree)
        self.assertEqual(tree.data, 'start')
        self.assertEqual(len(tree.children), 1)
        or_expr_node = tree.children[0]
        self.assertEqual(or_expr_node.data, 'or_expr')

        # or_expr has one child: and_expr (because there's an AND at top level)
        self.assertEqual(len(or_expr_node.children), 1)
        and_expr_node = or_expr_node.children[0]
        self.assertEqual(and_expr_node.data, 'and_expr')

        # and_expr -> field_term, AND token, field_term
        self.assertEqual(len(and_expr_node.children), 3)
        first_field_term = and_expr_node.children[0]
        and_tok = and_expr_node.children[1]
        second_field_term = and_expr_node.children[2]

        self.assertEqual(first_field_term.data, 'field_term')
        self.assertEqual(second_field_term.data, 'field_term')
        self.assertEqual(and_tok.type, 'AND')

        # first_field_term is the parenthesized expression: '(' or_expr ')'
        # so inside field_term, we expect a single child: an or_expr node
        self.assertEqual(len(first_field_term.children), 1)
        paren_expr_node = first_field_term.children[0]
        self.assertEqual(paren_expr_node.data, 'paren_expr')

        # sub_or_expr_node -> sub_or_expr_node[and_expr, OR, and_expr]
        sub_or_expr_node = paren_expr_node.children[0]
        self.assertEqual(len(sub_or_expr_node.children), 3)
        first_sub_and_expr = sub_or_expr_node.children[0]
        or_tok = sub_or_expr_node.children[1]
        second_sub_and_expr = sub_or_expr_node.children[2]
        self.assertEqual(first_sub_and_expr.data, 'and_expr')
        self.assertEqual(second_sub_and_expr.data, 'and_expr')
        self.assertEqual(or_tok.type, 'OR')

        # each of those and_expr -> single field_term -> normal_expr -> normal_term
        first_sub_ft = first_sub_and_expr.children[0]
        second_sub_ft = second_sub_and_expr.children[0]
        first_sub_normal_expr = first_sub_ft.children[0]
        second_sub_normal_expr = second_sub_ft.children[0]
        first_sub_nt = first_sub_normal_expr.children[0]
        second_sub_nt = second_sub_normal_expr.children[0]

        self.assertEqual(first_sub_nt.data, 'normal_term')
        self.assertEqual(first_sub_nt.children[0].value, 'title')
        self.assertEqual(first_sub_nt.children[1].value, ':')
        self.assertEqual(first_sub_nt.children[2].value, 'Hello')

        self.assertEqual(second_sub_nt.data, 'normal_term')
        self.assertEqual(second_sub_nt.children[0].value, 'abstract')
        self.assertEqual(second_sub_nt.children[1].value, ':')
        self.assertEqual(second_sub_nt.children[2].value, 'World')

        # second_field_term -> normal_expr -> normal_term
        self.assertEqual(len(second_field_term.children), 1)
        second_normal_expr = second_field_term.children[0]
        self.assertEqual(len(second_normal_expr.children), 1)
        second_nt = second_normal_expr.children[0]
        self.assertEqual(second_nt.data, 'normal_term')

        self.assertEqual(second_nt.children[0].value, 'keywords')
        self.assertEqual(second_nt.children[1].value, ':')
        self.assertEqual(second_nt.children[2].value, '"Deep Learning"')

    def test_pattern_6_not_clause(self):
        query = "title:Hello AND NOT abstract:World"
        tree = parse_with_semantic_check(query)

        # start -> or_expr -> [and_expr]
        self.assertIsInstance(tree, Tree)
        self.assertEqual(tree.data, 'start')
        self.assertEqual(len(tree.children), 1)
        or_expr_node = tree.children[0]
        self.assertEqual(or_expr_node.data, 'or_expr')

        # or_expr -> single child: and_expr
        self.assertEqual(len(or_expr_node.children), 1)
        and_expr_node = or_expr_node.children[0]
        self.assertEqual(and_expr_node.data, 'and_expr')

        # "title:Hello AND NOT abstract:World"
        # The grammar says and_expr -> [field_term],
        # but there's no top-level AND between separate field_terms in this example.
        # Instead, it's a single field_term that has: normal_expr (NOT normal_expr)
        self.assertEqual(len(and_expr_node.children), 1)
        field_term_node = and_expr_node.children[0]
        self.assertEqual(field_term_node.data, 'field_term')

        # field_term -> normal_expr, NOT token, normal_expr
        # So length = 3 if you have "A AND NOT B" style
        self.assertEqual(len(field_term_node.children), 3)

        left_normal_expr = field_term_node.children[0]
        not_token = field_term_node.children[1]
        right_normal_expr = field_term_node.children[2]

        self.assertIsInstance(not_token, Token)
        self.assertEqual(not_token.type, 'NOT')
        self.assertEqual(not_token.value, 'AND NOT')

        # left_normal_expr -> single normal_term
        left_nt = left_normal_expr.children[0]
        self.assertEqual(left_nt.data, 'normal_term')
        self.assertEqual(left_nt.children[0].value, 'title')
        self.assertEqual(left_nt.children[1].value, ':')
        self.assertEqual(left_nt.children[2].value, 'Hello')

        # right_normal_expr -> single normal_term
        right_nt = right_normal_expr.children[0]
        self.assertEqual(right_nt.data, 'normal_term')
        self.assertEqual(right_nt.children[0].value, 'abstract')
        self.assertEqual(right_nt.children[1].value, ':')
        self.assertEqual(right_nt.children[2].value, 'World')

    def test_pattern_7_date_clause(self):
        query = "title:Hello AND publication_year=2023"
        tree = parse_with_semantic_check(query)

        # start -> or_expr -> single child: and_expr
        self.assertEqual(tree.data, 'start')
        self.assertEqual(len(tree.children), 1)
        or_expr = tree.children[0]
        self.assertEqual(or_expr.data, 'or_expr')
        self.assertEqual(len(or_expr.children), 1)
        and_expr = or_expr.children[0]
        self.assertEqual(and_expr.data, 'and_expr')

        # and_expr should have 3 children: [ field_term, "AND", field_term ]
        self.assertEqual(len(and_expr.children), 3)
        first_field_term = and_expr.children[0]
        and_tok = and_expr.children[1]
        second_field_term = and_expr.children[2]

        self.assertEqual(first_field_term.data, 'field_term')
        self.assertIsInstance(and_tok, Token)
        self.assertEqual(and_tok.type, 'AND')
        self.assertEqual(second_field_term.data, 'field_term')

        # 1) First field_term -> normal_expr -> normal_term => "title:Hello"
        self.assertEqual(len(first_field_term.children), 1)
        normal_expr_node = first_field_term.children[0]
        self.assertEqual(normal_expr_node.data, 'normal_expr')

        self.assertEqual(len(normal_expr_node.children), 1)
        normal_term_node = normal_expr_node.children[0]
        self.assertEqual(normal_term_node.data, 'normal_term')
        self.assertEqual(normal_term_node.children[0].value, 'title')
        self.assertEqual(normal_term_node.children[1].value, ':')
        self.assertEqual(normal_term_node.children[2].value, 'Hello')

        # 2) Second field_term -> single child => date_expr => "publication_year=2023"
        self.assertEqual(len(second_field_term.children), 1)
        date_expr_node = second_field_term.children[0]
        self.assertEqual(date_expr_node.data, 'date_expr')

        self.assertEqual(len(date_expr_node.children), 3)
        self.assertEqual(date_expr_node.children[0].type, 'DATE_FIELD')
        self.assertEqual(date_expr_node.children[0].value, 'publication_year')
        self.assertEqual(date_expr_node.children[1].type, 'DATE_OP')
        self.assertEqual(date_expr_node.children[1].value, '=')
        self.assertEqual(date_expr_node.children[2].type, 'INT')
        self.assertEqual(date_expr_node.children[2].value, '2023')


class TestQueryParserError(unittest.TestCase):
    def test_unknown_field(self):
        query = "unknownfield:Something"
        with self.assertRaises((UnexpectedToken, UnexpectedCharacters)):
            parse_with_semantic_check(query)

    def test_missing_value(self):
        query = "title:"
        with self.assertRaises(UnexpectedEOF):
            parse_with_semantic_check(query)

    def test_unbalanced_parentheses(self):
        query = "(title:Hello OR abstract:World"
        with self.assertRaises(UnexpectedEOF):
            parse_with_semantic_check(query)

    def test_not_without_left_term(self):
        query = "AND NOT title:Hello"
        with self.assertRaises((UnexpectedToken, UnexpectedCharacters)):
            parse_with_semantic_check(query)

    def test_date_without_normal_expr(self):
        query = "publication_year=2020"
        with self.assertRaises(ValueError):
            parse_with_semantic_check(query)

    def test_invalid_operator(self):
        query = "title>Hello"
        with self.assertRaises((UnexpectedToken, UnexpectedCharacters)):
            parse_with_semantic_check(query)

    def test_invalid_date_operator(self):
        query = "title:Hello AND publication_year<>2020"
        with self.assertRaises((UnexpectedToken, UnexpectedCharacters)):
            parse_with_semantic_check(query)


class TestQueryParserRoundtrip(unittest.TestCase):
    base_transformer = BaseTransformer()

    def _parse_roundtrip(self, query):
        parse_tree = parse_with_semantic_check(query)
        round_trip_str = self.base_transformer.transform(parse_tree)
        return round_trip_str

    def test_single_normal_term(self):
        query = 'title:Test'
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_single_normal_term_with_word(self):
        query = 'abstract:"Hello World"'
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_multiple_terms_and(self):
        query = "(title:\"Hello World\") AND (abstract:Python)"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_multiple_terms_or(self):
        query = "(title:\"Hello\") OR (abstract:\"World\")"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_mixed_and_or(self):
        query = "(title:Alpha) AND (abstract:Beta) OR (keywords:\"Gamma\")"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_parentheses_nesting(self):
        query = "(title:Hello OR abstract:World) AND (keywords:\"Test Phrase\")"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_single_date_expr(self):
        query = "title:Hello AND publication_year>=2020"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_multiple_date_exprs(self):
        query = "title:\"AI\" AND publication_year>2018 AND publication_year<2024"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_date_expr_with_phrases_and_andor(self):
        query = "(title:\"Deep Learning\" OR author:\"Smith\") AND publication_year=2021"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_basic_and_not_clause(self):
        query = "title:Hello AND NOT abstract:World"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_multiple_and_not_clauses(self):
        query = "title:Foo AND NOT abstract:Bar AND author:Baz"
        self.assertEqual(query, self._parse_roundtrip(query))

    def test_complex_query(self):
        # Mix everything: OR, AND, parentheses, date expressions, AND NOT
        query = "(title:\"Deep Learning\" OR abstract:\"Neural Networks\") AND author:\"Goodfellow\" AND publication_year=2016 AND NOT keywords:\"GANs\""
        self.assertEqual(query, self._parse_roundtrip(query))

if __name__ == '__main__':
    unittest.main()
