from dataclasses import dataclass

import pytest

from sheerkaql.SheerkaQueryLangage import SheerkaQueryLanguage
from sheerkaql.symbols import flwr_sequence, attribute_value


def oset(x):
    return x


class A(object):
    def __init__(self, q):
        self.q = q

    def __repr__(self):
        return f"A({vars(self)})"

    def __eq__(self, other):
        if not isinstance(other, A):
            return False

        return self.q == other.q

    def __hash__(self):
        return hash(str(self.q))


@dataclass
class BagClass:
    property1: object
    property2: object

    def as_bag(self):
        return {
            "prop1": self.property1,
            "prop2": self.property2,
        }


execute = SheerkaQueryLanguage().execute


class TestSheerkaQueryLanguage:

    def test_i_can_get_the_root_of_a_query(self):
        hello = 'hello world!'
        q = SheerkaQueryLanguage().compile('hello')

        assert q(locals()) == oset([hello])

    def test_i_can_traverse_object(self):
        a = A(A(A("hello world!")))
        q = SheerkaQueryLanguage().compile("a.q.q.q")

        assert q(locals()) == oset(["hello world!"])

    def test_i_can_traverse_list(self):
        lst = [A("one"), A("two"), A("three")]
        a = A(lst)

        assert execute("a.q.q", {"a": a}) == oset(["one", "two", "three"])

    def test_i_can_traverse_list_of_list(self):
        sub_lst_number = [A("1"), A("2"), A("2")]
        sub_lst_letter = [A("a"), A("b"), A("c")]
        lst = [A("one"), A(sub_lst_number), A(sub_lst_letter)]
        a = A(lst)

        res = execute("a.q.q", {"a": a})
        assert res == oset(["one", *sub_lst_number, *sub_lst_letter])

    def test_i_can_traverse_object_when_where_condition_is_a_boolean(self):
        a = A(A(A("hello world!")))
        b_true = A(A(True))
        b_false = A(A(False))
        namespace = {"a": a, "hasattr": hasattr, "func": lambda x: x, "b_true": b_true, "b_false": b_false}

        assert execute("a.q[1 == 1].q.q", namespace) == oset(["hello world!"])
        assert execute("a.q[hasattr(self, 'q')].q.q", namespace) == oset(["hello world!"])

        assert execute("a.q[1 == 2].q.q", namespace) == oset([])
        assert execute("a.q[hasattr(self, 'x')].q.q", namespace) == oset([])

        assert execute("a.q[True].q.q", namespace, allow_builtins=True) == oset(["hello world!"])
        assert execute("a.q[False].q.q", namespace, allow_builtins=True) == oset([])

        assert execute("a.q[func(True)].q.q", namespace) == oset(["hello world!"])
        assert execute("a.q[func(False)].q.q", namespace) == oset([])

        assert execute("a.q[b_true.q.q].q.q", namespace) == oset(["hello world!"])
        assert execute("a.q[b_false.q.q].q.q", namespace) == oset([])

    def test_i_can_request_by_list_index(self):
        lst = [A("one"), A("two"), A("three")]
        a = A(lst)

        assert execute("a.q[1]", {"a": a}) == oset([lst[1]])
        with pytest.raises(IndexError):
            execute("a.q[99]", {"a": a})
        with pytest.raises(TypeError):
            execute("a.q['key']", {"a": a})

    def test_i_can_request_by_dictionary_key(self):
        lst = {"key1": "value1", "key2": "value2"}
        a = A(lst)

        assert execute("a.q['key1']", {"a": a}) == oset([lst["key1"]])
        with pytest.raises(KeyError):
            execute("a.q['key3']", {"a": a})

    def test_i_can_filter(self):
        l = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

        assert execute('l[self < 5]', locals()) == [0, 1, 2, 3, 4]

    def test_i_cannot_traverse_object_when_where_condition_is_not_a_boolean(self):
        a = A("hello world!")
        namespace = {"a": a, "func": lambda x: x, "dictionary": {"key": "value"}, "list": ["value"]}

        with pytest.raises(TypeError):
            execute("a.q[1].q.q", namespace)

        with pytest.raises(TypeError):
            execute("a.q[func(3)].q.q", namespace)

        with pytest.raises(TypeError):
            execute("a.q[dictionary['key']].q.q", namespace)

        with pytest.raises(TypeError):
            execute("a.q[list[0]].q.q", namespace)

    def test_i_can_traverse_object_when_as_bag_is_defined(self):
        a = BagClass(BagClass("sub_value1", BagClass("sub_sub_value1", "sub_sub_value1")), "value2")

        assert execute("a.prop1.prop2.prop1", {"a": a}) == oset(["sub_sub_value1"])

    def test_i_can_traverse_objects_when_where_condition_uses_as_bag(self):
        hash_map = {"sub_sub_value1": "hello world!"}
        b = BagClass(BagClass("sub_value1", BagClass("sub_sub_value1", "sub_sub_value1")), "value2")
        a = A(hash_map)

        assert execute("a.q[b.prop1.prop2.prop1]", {"a": a, "b": b}) == oset(["hello world!"])

    def test_i_can_compute_set_operations(self):
        a = [1, 2, 3, 3]
        b = [2, 4, 3, 4]

        assert execute("a | b", locals()) == [1, 2, 3, 4]
        assert execute("a - b", locals()) == [1]
        assert execute("b - a", locals()) == [4]
        assert execute("a & b", locals()) == [2, 3]

    def test_can_execute_set_expression(self):
        a = [1, 2, 3, 3]
        b = [2, 4, 3, 4]
        c = [1, 2]

        assert execute("return 1 if 1 in a else 0", locals()) == (1,)
        assert execute("return 1 if 1 not in a else 0", locals()) == (0,)
        assert execute("return 1 if 5 in a else 0", locals()) == (0,)
        assert execute("return 1 if 5 not in a else 0", locals()) == (1,)

        assert execute("return 1 if <c> subset <a> else 0", locals()) == (1,)
        assert execute("return 1 if <c> superset <a> else 0", locals()) == (0,)

        assert execute("return 1 if <a> subset <c> else 0", locals()) == (0,)
        assert execute("return 1 if <a> superset <c> else 0", locals()) == (1,)

        l = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

        assert execute('l - l[self < 5]', locals()) == (5, 6, 7, 8, 9)

    def test_i_can_traverse_a_complex_path(self):
        answer = 'o.x.y'

        o = A('top')
        o.x = [A('asdf'), A('123')]
        o.x[0].y = A(answer)
        d = {'hasattr': hasattr, 'o': o}

        assert execute('o/x[hasattr(self,"y")]/y/q', d) == oset([answer])
        assert execute('o/x', d) == oset(o.x)

    def test_i_can_execute_comparison_in_where_clauses(self):
        a = A(5)
        assert execute('a[self.q == 5]', locals()) == oset([a])
        assert execute('a[self.q != 5]', locals()) == oset([])
        assert execute('a[self.q >= 5]', locals()) == oset([a])
        assert execute('a[self.q <= 5]', locals()) == oset([a])
        assert execute('a[self.q > 5]', locals()) == oset([])
        assert execute('a[self.q < 5]', locals()) == oset([])
        assert execute('a[self.q == 7]', locals()) == oset([])
        assert execute('a[self.q != 7]', locals()) == oset([a])
        assert execute('a[self.q >= 7]', locals()) == oset([])
        assert execute('a[self.q <= 7]', locals()) == oset([a])
        assert execute('a[self.q > 7]', locals()) == oset([])
        assert execute('a[self.q < 7]', locals()) == oset([a])
        assert execute('a[self.q == 3]', locals()) == oset([])
        assert execute('a[self.q != 3]', locals()) == oset([a])
        assert execute('a[self.q >= 3]', locals()) == oset([a])
        assert execute('a[self.q <= 3]', locals()) == oset([])
        assert execute('a[self.q > 3]', locals()) == oset([a])
        assert execute('a[self.q < 3]', locals()) == oset([])

    def test_i_can_execute_boolean_expression_in_where_clauses(self):
        a = 'hello'
        true = True
        false = False
        assert execute('a[true]', locals()) == oset([a])
        assert execute('a[false]', locals()) == oset([])
        assert execute('a[not true]', locals()) == oset([])
        assert execute('a[not false]', locals()) == oset([a])
        assert execute('a[true and true]', locals()) == oset([a])
        assert execute('a[false and true]', locals()) == oset([])
        assert execute('a[not true and true]', locals()) == oset([])
        assert execute('a[not false and true]', locals()) == oset([a])
        assert execute('a[true or false]', locals()) == oset([a])
        assert execute('a[true or true]', locals()) == oset([a])
        assert execute('a[false or true]', locals()) == oset([a])
        assert execute('a[false or false]', locals()) == oset([])
        assert execute('a[not true or true]', locals()) == oset([a])
        assert execute('a[not false or false]', locals()) == oset([a])
        assert execute('a[true and true and true and true]', locals()) == oset([a])
        assert execute('a[true and true and true and false]', locals()) == oset([])
        assert execute('a[true and (false or true)]', locals()) == oset([a])
        assert execute('a[true and (false and true)]', locals()) == oset([])
        assert execute('a[true and (true and true)]', locals()) == oset([a])
        assert execute('a[true and (true and (not true or false))]', locals()) == oset([])

    def test_i_can_execute_comparison_in_where_clauses_using_builtin_functions(self):
        a = 'hello'
        true = True
        false = False
        d = locals()
        assert execute('a[1 == 1]', d) == oset([a])
        assert execute('a[-1 == -1]', d) == oset([a])
        assert execute('a[2.2 == 2.2]', d) == oset([a])
        assert execute('a[2.2 == float("2.2")]', d, True) == oset([a])
        assert execute('a[2 == int(2.2)]', d, True) == oset([a])
        assert execute('a["hello" == a]', d, True) == oset([a])
        assert execute('a["HELLO" == a.upper()]', d, True) == oset([a])

    def test_i_can_use_function_in_where_clauses(self):
        a = 'hello'

        def f():
            return 'hello'

        def g(x, y, z):
            return x + y + z

        def h(f, x):
            return f(x)

        def i(x):
            return x ** 2

        def j(f):
            return f

        true = True
        false = False
        d = locals()

        assert execute('a[f() == "hello"]', d) == oset([a])
        assert execute('a[g(1,2,3) == 6]', d) == oset([a])
        assert execute('a[h(i,3) == 9]', d) == oset([a])
        assert execute('a[i(j(j)(j)(j)(h)(i,3)) == 81]', d) == oset([a])
        with pytest.raises(TypeError):
            execute('a[f()]', d)

    def test_i_can_use_lists_in_where_clauses(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('a[l[0] == 1]', d) == oset([a])
        assert execute('a[l[1] == 2]', d) == oset([a])
        assert execute('a[l[7][0] == 1]', d) == oset([a])
        assert execute('a[l[7][1] == 2]', d) == oset([a])
        assert execute('a[l[7][7][0] == 1]', d) == oset([a])
        assert execute('a[l[7][7][1] == 2]', d) == oset([a])
        assert execute('a[l[7][7][7] == 8]', d) == oset([a])

    def test_i_can_use_dicts_in_where_clauses(self):
        a = 'hello'
        l = {"one": 1, "two": 2, "next": {"one": 1, "two": 2, "next": {"one": 1, "two": 2}}}
        d = locals()

        assert execute('a[l["one"] == 1]', d) == oset([a])
        assert execute('a[l["two"] == 2]', d) == oset([a])
        assert execute('a[l["next"]["one"] == 1]', d) == oset([a])
        assert execute('a[l["next"]["two"] == 2]', d) == oset([a])
        assert execute('a[l["next"]["next"]["one"] == 1]', d) == oset([a])
        assert execute('a[l["next"]["next"]["two"] == 2]', d) == oset([a])

    def test_i_can_use_callable_in_where_clauses(self):
        a = 'hello'

        def f():
            return 'hello'

        def g(x, y, z):
            return x + y + z

        def h(f, x):
            return f(x)

        def i(x):
            return x ** 2

        def j(f):
            return f

        m = {"one": 1, "two": 2, "next": [1, 2, 3, 4, 5, 6, 7, j]}
        d = locals()

        assert execute('a[m["next"][7](j)(m["next"][7])(m["next"])[7](i)(m["two"]) == 4]', d) == oset([a])

    def test_i_can_execute_flwr_expression(self):
        def f():
            return 1, 2, 3

        d = locals()
        assert execute('for x in f() return x', d) == (1, 2, 3)
        assert execute('for x in f() let y = f() return x, y', d) == ((1, (1, 2, 3)), (2, (1, 2, 3)), (3, (1, 2, 3)))

    def test_i_can_execute_flwr_with_order_by(self):
        def f():
            return [1, 3, 2]

        d = locals()
        with pytest.raises(SyntaxError):
            execute('for x in f() order by "asdf" asc return x', d)
        with pytest.raises(SyntaxError):
            execute('for x in f() order by 0 asc return "asdf":x', d)
        assert execute('for x in f() order by 0 asc return x', d) == (1, 2, 3)
        assert execute('for x in f() order by 0 desc return x', d) == (3, 2, 1)

    def test_i_can_execute_flwr_with_user_defined_functions(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('''
          for i in l
          let f = function() { 125 }
          return f()
        ''', d) == (125, 125, 125, 125, 125, 125, 125, 125)

        assert execute('''
          for i in l
            let f = function(q) {
              for _ in <a>
              where isinstance(q, list)
              return {
                for j in q
                return f(j)
              }
            }
          return f(i)
        ''', d, True) == ((), (), (), (), (), (), (),
                          (((), (), (), (), (), (), (),
                            (((), (), (), (), (), (), (), ()),)),))

    def test_i_can_execute_flwr_with_if_expression(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        q = True
        d = locals()

        assert execute('''
          for x in <a> return if (q) then 1 else 0
          ''', d) == (1,)

        assert execute('''
          for x in <l> return if (x % 2 == 0) then 1 else 0
          ''', d) == (0, 1, 0, 1, 0, 1, 0, 1, 0, 1)

        assert execute('''
          for x in <l> return if x % 2 == 0 then 1 else 0
          ''', d) == (0, 1, 0, 1, 0, 1, 0, 1, 0, 1)

        assert execute('''
          for x in <l> return 1 if x % 2 == 0 else 0
          ''', d) == (0, 1, 0, 1, 0, 1, 0, 1, 0, 1)

        assert execute('''
          for x in <l> return 1 if (x % 2 == 0) else 0
          ''', d) == (0, 1, 0, 1, 0, 1, 0, 1, 0, 1)

        assert execute('''
            for x in <a> return if (True or X) then 1 else 0
            ''', d, True) == (1,)

    def test_if_short_circuit(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('''
            for x in <a> return if (True or X) then 1 else 0
            ''', d, True) == (1,)

        assert execute('''
            for x in <a> return if (False and false.x) then 1 else 0
            ''', d, True) == (0,)

    def test_i_can_flatten_result(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('''
           for x in l return flatten x
           ''', d, True) == (1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8)

    def test_flattened_return(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('''
            for i in l
            let f = function(l) {
              if (isinstance(l, list))
              then {for j in l return f(j)}
              else l
            }
            return flatten f(i)''', d, True) == (1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8)

        assert execute('''
            for i in l
            let f = function(l) {
              if (isinstance(l, list))
              then {for j in l return f(j)}
              else {a:l}
            }
            return flatten f(i)
            ''', d, True) == tuple({a: i} for i in (1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8))

    def test_i_can_return_none(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        assert execute('for x in l return None', d, True) == (None, None, None, None, None, None, None, None)

    def test_i_can_return_a_class(self):
        l = [1, 2, 3]
        d = {"A": A, "l": l}

        assert execute('for x in l return A(x)', d, True) == (A(1), A(2), A(3))

    def test_i_can_execute_flwr_when_there_is_no_for_statement(self):
        a = 'hello'
        l = [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, [1, 2, 3, 4, 5, 6, 7, 8]]]
        d = locals()

        flwr = flwr_sequence([attribute_value('hello', scalar=True)])
        assert flwr(d) == ('hello',)
        assert execute("return l", d) == (l,)
        assert execute('''
            let f = function(l) {
              if (isinstance(l, list))
              then {for j in l return f(j)}
              else l
            }
            return flatten f(l)
          ''', d, True) == (1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 8)

    def test_can_collect(self):
        l = [1, 2, 3, 4, 5, 6, 7, 3, 4, 5, 6, 7, 3, 4]
        d = locals()

        assert execute('''
            for n in l
            collect n as n with function(prev, next) {
                if prev == None then 1 else prev + 1
            }
          ''', d, True) == {1: 1, 2: 1, 3: 3, 4: 3, 5: 2, 6: 2, 7: 2}

        assert execute('''
            for n in [1,2,3,4,5,6,7,3,4,5,6,7,3,4]
            collect n as n with function(prev, next) {
                if prev == None then 1 else prev + 1
             }
          ''', d, True) == {1: 1, 2: 1, 3: 3, 4: 3, 5: 2, 6: 2, 7: 2}

    def test_i_can_perform_calculations(self):
        d = locals()

        assert execute('''
            for n in [
                4.0*3.0/2.0,
                4.0/3.0*2.0,
                (3.0+9.0)*4.0/8.0,
                ((9.0-3.0)+(5.0-3.0))/2.0 + 2.0,
                5.0 * 4.0 / 2.0 - 10.0 + 5.0 - 2.0 + 3.0,
                5.0 / 4.0 * 2.0 + 10.0 - 5.0 * 2.0 / 3.0
            ]
            return n
          ''', d, True) == (
            4.0 * 3.0 / 2.0,
            4.0 / 3.0 * 2.0,
            (3.0 + 9.0) * 4.0 / 8.0,
            ((9.0 - 3.0) + (5.0 - 3.0)) / 2.0 + 2.0,
            5.0 * 4.0 / 2.0 - 10.0 + 5.0 - 2.0 + 3.0,
            5.0 / 4.0 * 2.0 + 10.0 - 5.0 * 2.0 / 3.0
        )

    def test_i_can_execute_flwr_with_in_and_not_in(self):
        l = [[1, 2, 3], [4, 5, 6], [7, 4, 3], [5, 6, 7], [3, 4]]
        d = locals()

        assert execute('''
            for n in l
            where 4 in n
            return n
          ''', d, True) == ([4, 5, 6], [7, 4, 3], [3, 4])

        assert execute('''
            for n in l
            where 4 not in n
            return n
          ''', d, True) == ([1, 2, 3], [5, 6, 7])

    def test_multi_collect(self):
        l = [1, 2, 3, 4, 5, 6, 7, 3, 4, 5, 6, 7, 3, 4]
        d = locals()

        assert execute('''
            for n in l
            let counter = function(prev, next) {
                if prev == None then 1 else prev + 1
            }
            where 1 in l and 12 not in l
            collect n as n with counter
            collect n as (int(n)/int(2)) with counter
          ''', d, True), (
            {1: 1, 2: 1, 3: 3, 4: 3, 5: 2, 6: 2, 7: 2},
            {0: 1, 1: 4, 2: 5, 3: 4})

    def test_i_can_execute_flwr_on_objects_attributes(self):
        lst = [A(1), A(2), A(3)]
        d = locals()

        assert execute('for x in <lst.q> return x', d) == (1, 2, 3)
        assert execute('for x in lst return x.q', d) == (1, 2, 3)  # another way

    def test_exception_during_return_are_caught(self):
        lst = [A(1), A([2])]
        d = locals()

        res = execute('for x in lst return x.q + [2]', d)
        assert len(res) == 2
        assert isinstance(res[0], TypeError)
        assert res[1] == [2,2]

    def test_i_cannot_execute_flwr_on_attributes_that_does_not_exist(self):
        lst = [A(1), A(2), A(3)]
        d = locals()

        with pytest.raises(AttributeError):
            execute('for x in <lst.x> return x', d)

        with pytest.raises(AttributeError):
            execute('for x in lst.q return x', d)  # AttributeError: 'list' object has no attribute 'q'