import unittest
from abc import ABC
from datetime import datetime

from lark.exceptions import VisitError

from slrquerytester.translators.core_transformer import CORETransformer
from slrquerytester.translators.dimensions_transformer import DimensionsTransformer
from slrquerytester.translators.ieee_transformer import IEEETransformer
from slrquerytester.translators.lens_transformer import LENSTransformer
from slrquerytester.translators.openalex_transformer import OpenAlexTransformer
from slrquerytester.translators.scopus_transformer import ScopusTransformer
from slrquerytester.translators.springer_transformer import SpringerTransformerFree, SpringerTransformerPremium
from slrquerytester.translators.wos_transformer import WOSTransformer
from slrquerytester.utils.query_parser import parse_with_semantic_check


class TestTransformer(ABC, unittest.TestCase):
    def _perform_test(self, input_str, expected_str, expected_sy = None, expected_ey = None):
        transformed_query = self._transform_query(input_str)
        if isinstance(transformed_query, str):
            self.assertEqual(expected_str, transformed_query)
        elif isinstance(transformed_query, tuple) and len(transformed_query) == 3:
            self.assertEqual(expected_str, transformed_query[0])
            if expected_sy:
                self.assertEqual(expected_sy, transformed_query[1])
            else:
                self.assertIsNone(transformed_query[1])
            if expected_ey:
                self.assertEqual(expected_ey, transformed_query[2])
            else:
                self.assertIsNone(transformed_query[2])

    @staticmethod
    def _transform_query(query_str: str):
        pass

class TestWOSTransformer(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = WOSTransformer()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "TI=(\"Hello\")")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "TI=(\"Hello\") OR AB=(\"World\")")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "TI=(\"Hello\") AND AB=(\"World String\")")

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           "(TI=(\"Hello\") AND AB=(\"World String\")) OR AU=(\"Smith\")")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           "TI=(\"Hello\") NOT AU=(\"Smith\")")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "TI=(\"Hello\")",
                           2020,
                           datetime.now().year)

    def test_date_expr_le(self):
        self._perform_test("(abstract:\"World String\") AND (title:\"Hello\") AND publication_year<=2010",
                           "(AB=(\"World String\")) AND (TI=(\"Hello\"))",
                           1900,
                           2010)

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "TI=(\"Hello\")",
                           2005,
                           2005)

    def test_multi_date_exprs(self):
        self._perform_test("title:Hello AND publication_year>2010 AND publication_year<2022",
                           "TI=(\"Hello\")",
                           2011,
                           2021)

class TestSpringerTransformerFree(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = SpringerTransformerFree()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "(\"Hello\")")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "(\"Hello\" OR \"World\")")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "(\"Hello\" AND \"World String\")")

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           "((\"Hello\" AND \"World String\") OR \"Smith\")")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           "(\"Hello\" NOT \"Smith\")")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "(\"Hello\" AND datefrom:\"2020-01-01\")")

    def test_date_expr_le(self):
        self._perform_test("abstract:\"World String\" AND publication_year<=2010",
                           "(\"World String\" AND dateto:\"2010-12-31\")")

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "(\"Hello\" AND datefrom:\"2005-01-01\" AND dateto:\"2005-12-31\")")

    def test_multi_date_exprs(self):
        self._perform_test("title:Hello AND publication_year>2010 AND publication_year<2022",
                           "(\"Hello\" AND datefrom:\"2011-01-01\" AND dateto:\"2021-12-31\")")

    def test_complex_parentheses_not_and_dates(self):
        self._perform_test("(title:\"Hello\" OR abstract:\"World\") NOT (author:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
                           "((\"Hello\" OR \"World\") NOT (\"Smith\" OR keyword:\"String\") AND datefrom:\"2015-01-01\" AND dateto:\"2015-12-31\")")

class TestSpringerTransformerPremium(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = SpringerTransformerPremium()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "(title:\"Hello\")")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "(title:\"Hello\" OR \"World\")")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "(title:\"Hello\" AND \"World String\")")

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           "((title:\"Hello\" AND \"World String\") OR name:\"Smith\")")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           "(title:\"Hello\" NOT name:\"Smith\")")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "(title:\"Hello\" AND datefrom:\"2020-01-01\")")

    def test_date_expr_le(self):
        self._perform_test("abstract:\"World String\" AND publication_year<=2010",
                           "(\"World String\" AND dateto:\"2010-12-31\")")

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "(title:\"Hello\" AND datefrom:\"2005-01-01\" AND dateto:\"2005-12-31\")")

    def test_multi_date_exprs(self):
        self._perform_test("address:Hello AND publication_year>2010 AND publication_year<2022",
                           "(orgname:\"Hello\" AND datefrom:\"2011-01-01\" AND dateto:\"2021-12-31\")")

    def test_complex_parentheses_not_and_dates(self):
        self._perform_test("(title:\"Hello\" OR address:\"World\") NOT (author:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
                           "((title:\"Hello\" OR orgname:\"World\") NOT (name:\"Smith\" OR keyword:\"String\") AND datefrom:\"2015-01-01\" AND dateto:\"2015-12-31\")")

class TestScopusTransformer(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = ScopusTransformer()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "TITLE(\"Hello\")")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "TITLE(\"Hello\") OR ABS(\"World\")")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "TITLE(\"Hello\") AND ABS(\"World String\")")

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           "(TITLE(\"Hello\") AND ABS(\"World String\")) OR AUTHOR-NAME(\"Smith\")")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           "TITLE(\"Hello\") AND NOT AUTHOR-NAME(\"Smith\")")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "TITLE(\"Hello\") AND PUBYEAR > 2019")

    def test_date_expr_le(self):
        self._perform_test("abstract:\"World String\" AND publication_year<=2010",
                           "ABS(\"World String\") AND PUBYEAR < 2011")

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "TITLE(\"Hello\") AND PUBYEAR = 2005")

    def test_multi_date_exprs(self):
        self._perform_test("address:Hello AND publication_year>2010 AND publication_year<2022",
                           "AFFIL(\"Hello\") AND PUBYEAR > 2010 AND PUBYEAR < 2022")

    def test_complex_parentheses_not_and_dates(self):
        self._perform_test("(title:\"Hello\" OR address:\"World\") NOT (author:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
                           "(TITLE(\"Hello\") OR AFFIL(\"World\")) AND NOT (AUTHOR-NAME(\"Smith\") OR KEY(\"String\")) AND PUBYEAR = 2015")

class TestOpenAlexTransformer(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = OpenAlexTransformer()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "title.search:\"Hello\"")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "title_and_abstract.search:\"Hello\"|\"World\"")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "title.search:\"Hello\",abstract.search:\"World String\"")

    def test_and_or_expr(self):
        # "(title:Hello AND abstract:\"World String\") || author:Smith"
        # Some might prefer a ValueError if the grammar tries to do OR with an ignored field.
        # But as doc says "any conditions on 'author' are discarded."
        # => result is just "title.search:Hello,abstract.search:\"World String\""
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           "title.search:\"Hello\",abstract.search:\"World String\"")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT abstract:\"Smith\"",
                           "title.search:\"Hello\",abstract.search:!\"Smith\"")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "title.search:\"Hello\",from_publication_date:2020-01-01")

    def test_date_expr_le(self):
        self._perform_test("(abstract:\"World String\") AND (title:\"Hello\") AND publication_year<=2010",
                           "abstract.search:\"World String\",title.search:\"Hello\",to_publication_date:2010-12-31")

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "title.search:\"Hello\",publication_year:2005")

    def test_multi_date_exprs(self):
        # "title:Hello AND publication_year>2010 AND publication_year<2022"
        # => publication_year>2010 => :>2010, <2022 => :<2022
        # So => "title.search:Hello,publication_year:>2010,publication_year:<2022"
        # or if your code merges them into a single set? Typically it won't.
        # Possibly "title.search:Hello,publication_year:>2010,publication_year:<2022"
        self._perform_test("title:Hello AND publication_year>2010 AND publication_year<2022",
                           "title.search:\"Hello\",from_publication_date:2011-01-01,to_publication_date:2021-12-31")

    def test_complex_parentheses_not_and_dates(self):
        # e.g. "(title:\"Hello\" OR address:\"World\") NOT (author:\"Smith\" OR keywords:\"String\") AND publication_year=2015"
        # 1) (title OR address) => if both in OR? But 'address' => raw_affiliation_strings.search
        #    it merges into 'default.search'
        # 2) NOT (keywords or keywords) => effectively just 'keywords'
        self._perform_test("(title:\"Hello\" OR address:\"World\") NOT (keywords:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
                           "default.search:\"Hello\"|\"World\",keywords.keyword:!\"Smith\"|\"String\",publication_year:2015")

class TestCORETransformer(TestTransformer):

    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = CORETransformer()
        return transformer.transform(tree)


    def test_basic_title(self):
        self._perform_test("title:Hello",
                           "title:\"Hello\"")

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           "title:\"Hello\" OR abstract:\"World\"")

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           "title:\"Hello\" AND abstract:\"World String\"")

    def test_and_or_expr(self):
        self._perform_test("(title:\"Hello\" AND abstract:\"World String\") || author:Smith",
                           "(title:\"Hello\" AND abstract:\"World String\") OR authors:\"Smith\"")

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           "title:\"Hello\" AND NOT authors:\"Smith\"")

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           "title:\"Hello\" AND yearPublished >= 2020")

    def test_date_expr_le(self):
        self._perform_test("abstract:\"World String\" AND publication_year<=2010",
                           "abstract:\"World String\" AND yearPublished <= 2010")

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           "title:\"Hello\" AND yearPublished = 2005")

    def test_multi_date_exprs(self):
        with self.assertRaises(VisitError):
            self._transform_query("address:\"Hello\" AND publication_year>2010 AND publication_year<2022")

    def test_complex_parentheses_not_and_dates(self):
        self._perform_test("(title:\"Hello\" OR address:\"World\") NOT (author:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
                           "(title:\"Hello\") AND NOT (authors:\"Smith\" OR fullText:\"String\") AND yearPublished = 2015")

class TestDimensionsTransformer(TestTransformer):
    def _transform_query(self, query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = DimensionsTransformer()
        result = transformer.transform(tree)
        return result

    def test_basic_title(self):
        self._perform_test(
            "title:Hello",
            "search publications in title_only for \"Hello\" return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
        )

    def test_basic_or_expr(self):
        self._perform_test(
            "title:Hello OR abstract:\"World\"",
            (
                "search publications in title_only for \"Hello\" OR in abstract for \"World\" "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_basic_and_expr(self):
        self._perform_test(
            "title:Hello AND abstract:\"World String\"",
            (
                "search publications in title_only for \"Hello\" AND in abstract for \"World String\" "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_and_or_expr(self):
        self._perform_test(
            "(title:Hello AND abstract:\"World String\") || author:Smith",
            (
                "search publications "
                "(in title_only for \"Hello\" AND in abstract for \"World String\") OR in authors for \"Smith\" "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_not_clause(self):
        self._perform_test(
            "title:Hello AND NOT abstract:\"Smith\"",
            (
                "search publications in title_only for \"Hello\" NOT in abstract for \"Smith\""
                " return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_date_expr_ge(self):
        self._perform_test(
            "title:\"Hello\" AND publication_year>=2020",
            (
                "search publications in title_only for \"Hello\" "
                "where year>=2020 "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_date_expr_le(self):
        self._perform_test(
            "(abstract:\"World String\") AND (title:\"Hello\") AND publication_year<=2010",
            (
                "search publications "
                "(in abstract for \"World String\") AND (in title_only for \"Hello\") "
                "where year<=2010 "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_date_expr_equals(self):
        # "title:Hello AND publication_year=2005"
        # => sets start_year=2005, stop_year=2005 => "where year>=2005 and year<=2005"
        self._perform_test(
            "title:Hello AND publication_year=2005",
            (
                "search publications in title_only for \"Hello\" "
                "where year>=2005 and year<=2005 "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_multi_date_exprs(self):
        # "title:Hello AND publication_year>2010 AND publication_year<2022"
        # => effectively sets start_year=2011, stop_year=2021 if we adopt the "year+1" or "year-1" logic
        self._perform_test(
            "title:Hello AND publication_year>2010 AND publication_year<2022",
            (
                "search publications in title_only for \"Hello\" "
                "where year>=2011 and year<=2021 "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

    def test_complex_parentheses_not_and_dates(self):
        # "(title:\"Hello\" OR address:\"World\") NOT (keywords:\"Smith\" OR keywords:\"String\") AND publication_year=2015"
        # Your actual DSL may differ, but here's a plausible shape:
        self._perform_test(
            "(title:\"Hello\" OR address:\"World\") NOT (keywords:\"Smith\" OR keywords:\"String\") AND publication_year=2015",
            (
                "search publications "
                "(in title_only for \"Hello\" OR in research_org_countries for \"World\") NOT (in concepts for \"Smith\" OR in concepts for \"String\") "
                "where year>=2015 and year<=2015 "
                "return publications[title+authors+doi+year+abstract+linkout+concepts+journal+proceedings_title+book_title]"
            )
        )

class TestIEEETransformer(TestTransformer):
    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = IEEETransformer()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           '"Document Title":"Hello"')

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           '"Document Title":"Hello" OR "Abstract":"World"')

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           '"Document Title":"Hello" AND "Abstract":"World String"')

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           '("Document Title":"Hello" AND "Abstract":"World String") OR "Authors":"Smith"')

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           '"Document Title":"Hello" NOT "Authors":"Smith"')

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           '"Document Title":"Hello"',
                           2020,
                           None)

    def test_date_expr_le(self):
        self._perform_test("(abstract:\"World String\") AND (title:\"Hello\") AND publication_year<=2010",
                           '("Abstract":"World String") AND ("Document Title":"Hello")',
                           None,
                           2010)

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           '"Document Title":"Hello"',
                           2005,
                           2005)

    def test_multi_date_exprs(self):
        self._perform_test("title:Hello AND publication_year>2010 AND publication_year<2022",
                           '"Document Title":"Hello"',
                           2011,
                           2021)

class TestLENSTransformer(TestTransformer):
    @staticmethod
    def _transform_query(query_str):
        tree = parse_with_semantic_check(query_str)
        transformer = LENSTransformer()
        return transformer.transform(tree)

    def test_basic_title(self):
        self._perform_test("title:Hello",
                           'title:(Hello)')

    def test_basic_or_expr(self):
        self._perform_test("title:Hello OR abstract:\"World\"",
                           'title:(Hello) OR abstract:(World)')

    def test_basic_and_expr(self):
        self._perform_test("title:Hello AND abstract:\"World String\"",
                           'title:(Hello) AND abstract:(World String)')

    def test_and_or_expr(self):
        self._perform_test("(title:Hello AND abstract:\"World String\") || author:Smith",
                           '(title:(Hello) AND abstract:(World String)) OR author.displayname:(Smith)')

    def test_not_clause(self):
        self._perform_test("title:Hello AND NOT author:\"Smith\"",
                           'title:(Hello) NOT author.displayname:(Smith)')

    def test_date_expr_ge(self):
        self._perform_test("title:\"Hello\" AND publication_year>=2020",
                           'title:(Hello) AND year_published:[2020 TO 2025]')

    def test_date_expr_le(self):
        self._perform_test("(abstract:\"World String\") AND (title:\"Hello\") AND publication_year<=2010",
                           '(abstract:(World String)) AND (title:(Hello)) AND year_published:[0001 TO 2010]')

    def test_date_expr_equals(self):
        self._perform_test("title:Hello AND publication_year=2005",
                           'title:(Hello) AND year_published:2005')

    def test_multi_date_exprs(self):
        self._perform_test("title:Hello AND publication_year>2010 AND publication_year<2022",
                           'title:(Hello) AND year_published:[2011 TO 2025] AND year_published:[0001 TO 2021]')

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