from datetime import datetime
from typing import List, Optional
import unittest

from data_diff.abcs.database_types import FractionalType, TemporalType
from data_diff.databases.base import Database, BaseDialect
from data_diff.utils import CaseInsensitiveDict, CaseSensitiveDict

from data_diff.databases.base import Compiler, CompileError
from data_diff.queries.api import outerjoin, cte, when, coalesce
from data_diff.queries.ast_classes import Random
from data_diff.queries.api import code, this, table


def normalize_spaces(s: str):
    return " ".join(s.split())


class MockDialect(BaseDialect):
    name = "MockDialect"

    PLACEHOLDER_TABLE = None
    ROUNDS_ON_PREC_LOSS = False

    def quote(self, s: str) -> str:
        return s

    def concat(self, l: List[str]) -> str:
        s = ", ".join(l)
        return f"concat({s})"

    def to_comparable(self, s: str) -> str:
        return s

    def to_string(self, s: str) -> str:
        return f"cast({s} as varchar)"

    def is_distinct_from(self, a: str, b: str) -> str:
        return f"{a} is distinct from {b}"

    def random(self) -> str:
        return "random()"

    def current_timestamp(self) -> str:
        return "now()"

    def current_database(self) -> str:
        return "current_database()"

    def current_schema(self) -> str:
        return "current_schema()"

    def limit_select(
        self,
        select_query: str,
        offset: Optional[int] = None,
        limit: Optional[int] = None,
        has_order_by: Optional[bool] = None,
    ) -> str:
        x = offset and f"OFFSET {offset}", limit and f"LIMIT {limit}"
        result = " ".join(filter(None, x))
        return f"SELECT * FROM ({select_query}) AS LIMITED_SELECT {result}"

    def explain_as_text(self, query: str) -> str:
        return f"explain {query}"

    def timestamp_value(self, t: datetime) -> str:
        return f"timestamp '{t}'"

    def set_timezone_to_utc(self) -> str:
        return "set timezone 'UTC'"

    def optimizer_hints(self, s: str):
        return f"/*+ {s} */ "

    def md5_as_int(self, s: str) -> str:
        raise NotImplementedError

    def md5_as_hex(self, s: str) -> str:
        raise NotImplementedError

    def normalize_number(self, value: str, coltype: FractionalType) -> str:
        raise NotImplementedError

    def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
        raise NotImplementedError

    parse_type = NotImplemented


class MockDatabase(Database):
    CONNECT_URI_HELP = "mock://"
    CONNECT_URI_PARAMS = []
    dialect = MockDialect()

    _query = NotImplemented
    query_table_schema = NotImplemented
    select_table_schema = NotImplemented
    _process_table_schema = NotImplemented
    parse_table_name = NotImplemented
    close = NotImplemented
    _normalize_table_path = NotImplemented
    is_autocommit = NotImplemented


class TestQuery(unittest.TestCase):
    def setUp(self):
        pass

    def test_basic(self):
        c = Compiler(MockDatabase())

        t = table("point")
        t2 = t.select(x=this.x + 1, y=t["y"] + this.x)
        assert c.compile(t2) == "SELECT (x + 1) AS x, (y + x) AS y FROM point"

        t = table("point").where(this.x == 1, this.y == 2)
        assert c.compile(t) == "SELECT * FROM point WHERE (x = 1) AND (y = 2)"

        t = table("person").where(this.name == "Albert")
        self.assertEqual(c.compile(t), "SELECT * FROM person WHERE (name = 'Albert')")

    def test_outerjoin(self):
        c = Compiler(MockDatabase())

        a = table("a")
        b = table("b")
        keys = ["x", "y"]
        cols = ["u", "v"]

        j = outerjoin(a, b).on(a[k] == b[k] for k in keys)

        self.assertEqual(
            c.compile(j), "SELECT * FROM a tmp1 FULL OUTER JOIN b tmp2 ON (tmp1.x = tmp2.x) AND (tmp1.y = tmp2.y)"
        )

    def test_schema(self):
        c = Compiler(MockDatabase())
        schema = dict(id="int", comment="varchar")

        # test table
        t = table("a", schema=CaseInsensitiveDict(schema))
        q = t.select(this.Id, t["COMMENT"])
        assert c.compile(q) == "SELECT id, comment FROM a"

        t = table("a", schema=CaseSensitiveDict(schema))
        self.assertRaises(KeyError, t.__getitem__, "Id")
        self.assertRaises(KeyError, t.select, this.Id)

        # test select
        q = t.select(this.id)
        self.assertRaises(KeyError, q.__getitem__, "comment")

        # test join
        s = CaseInsensitiveDict({"x": int, "y": int})
        a = table("a", schema=s)
        b = table("b", schema=s)
        keys = ["x", "y"]
        j = outerjoin(a, b).on(a[k] == b[k] for k in keys).select(a["x"], b["y"], xsum=a["x"] + b["x"])
        j["x"], j["y"], j["xsum"]
        self.assertRaises(KeyError, j.__getitem__, "ysum")

    def test_commutable_select(self):
        # c = Compiler(MockDatabase())

        t = table("a")
        q1 = t.select("a").where("b")
        q2 = t.where("b").select("a")
        assert q1 == q2, (q1, q2)

    def test_cte(self):
        c = Compiler(MockDatabase())

        t = table("a")

        # single cte
        t2 = cte(t.select(this.x))
        t3 = t2.select(this.x)

        expected = "WITH tmp1 AS (SELECT x FROM a) SELECT x FROM tmp1"
        assert normalize_spaces(c.compile(t3)) == expected

        # nested cte
        c = Compiler(MockDatabase())
        t4 = cte(t3).select(this.x)

        expected = "WITH tmp1 AS (SELECT x FROM a), tmp2 AS (SELECT x FROM tmp1) SELECT x FROM tmp2"
        assert normalize_spaces(c.compile(t4)) == expected

        # parameterized cte
        c = Compiler(MockDatabase())
        t2 = cte(t.select(this.x), params=["y"])
        t3 = t2.select(this.y)

        expected = "WITH tmp1(y) AS (SELECT x FROM a) SELECT y FROM tmp1"
        assert normalize_spaces(c.compile(t3)) == expected

    def test_funcs(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.order_by(Random()).limit(10))
        self.assertEqual(q, "SELECT * FROM (SELECT * FROM a ORDER BY random()) AS LIMITED_SELECT LIMIT 10")

        q = c.compile(t.select(coalesce(this.a, this.b)))
        self.assertEqual(q, "SELECT COALESCE(a, b) FROM a")

    def test_select_distinct(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.select(this.b, distinct=True))
        assert q == "SELECT DISTINCT b FROM a"

        # selects merge
        q = c.compile(t.where(this.b > 10).select(this.b, distinct=True))
        self.assertEqual(q, "SELECT DISTINCT b FROM a WHERE (b > 10)")

        # selects stay apart
        q = c.compile(t.limit(10).select(this.b, distinct=True))
        self.assertEqual(q, "SELECT DISTINCT b FROM (SELECT * FROM (SELECT * FROM a) AS LIMITED_SELECT LIMIT 10) tmp1")

        q = c.compile(t.select(this.b, distinct=True).select(distinct=False))
        self.assertEqual(q, "SELECT * FROM (SELECT DISTINCT b FROM a) tmp2")

    def test_select_with_optimizer_hints(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.select(this.b, optimizer_hints="PARALLEL(a 16)"))
        assert q == "SELECT /*+ PARALLEL(a 16) */ b FROM a"

        q = c.compile(t.where(this.b > 10).select(this.b, optimizer_hints="PARALLEL(a 16)"))
        self.assertEqual(q, "SELECT /*+ PARALLEL(a 16) */ b FROM a WHERE (b > 10)")

        q = c.compile(t.limit(10).select(this.b, optimizer_hints="PARALLEL(a 16)"))
        self.assertEqual(
            q, "SELECT /*+ PARALLEL(a 16) */ b FROM (SELECT * FROM (SELECT * FROM a) AS LIMITED_SELECT LIMIT 10) tmp1"
        )

        q = c.compile(t.select(this.a).group_by(this.b).agg(this.c).select(optimizer_hints="PARALLEL(a 16)"))
        self.assertEqual(
            q, "SELECT /*+ PARALLEL(a 16) */ * FROM (SELECT b, c FROM (SELECT a FROM a) tmp2 GROUP BY 1) tmp3"
        )

    def test_table_ops(self):
        c = Compiler(MockDatabase())
        a = table("a").select(this.x)
        b = table("b").select(this.y)

        q = c.compile(a.union(b))
        assert q == "SELECT x FROM a UNION SELECT y FROM b"

        q = c.compile(a.union_all(b))
        assert q == "SELECT x FROM a UNION ALL SELECT y FROM b"

        q = c.compile(a.minus(b))
        assert q == "SELECT x FROM a EXCEPT SELECT y FROM b"

        q = c.compile(a.intersect(b))
        assert q == "SELECT x FROM a INTERSECT SELECT y FROM b"

    def test_ops(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.select(this.b + this.c))
        self.assertEqual(q, "SELECT (b + c) FROM a")

        q = c.compile(t.select(this.b.like(this.c)))
        self.assertEqual(q, "SELECT (b LIKE c) FROM a")

        q = c.compile(t.select(-this.b.sum()))
        self.assertEqual(q, "SELECT (-SUM(b)) FROM a")

    def test_group_by(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.group_by(this.b).agg(this.c))
        self.assertEqual(q, "SELECT b, c FROM a GROUP BY 1")

        q = c.compile(t.where(this.b > 1).group_by(this.b).agg(this.c))
        self.assertEqual(q, "SELECT b, c FROM a WHERE (b > 1) GROUP BY 1")

        self.assertRaises(CompileError, c.compile, t.select(this.b).group_by(this.b))

        q = c.compile(t.select(this.b).group_by(this.b).agg())
        self.assertEqual(q, "SELECT b FROM (SELECT b FROM a) tmp1 GROUP BY 1")

        q = c.compile(t.group_by(this.b, this.c).agg(this.d, this.e))
        self.assertEqual(q, "SELECT b, c, d, e FROM a GROUP BY 1, 2")

        # Having
        q = c.compile(t.group_by(this.b).agg(this.c).having(this.b > 1))
        self.assertEqual(q, "SELECT b, c FROM a GROUP BY 1 HAVING (b > 1)")

        q = c.compile(t.group_by(this.b).having(this.b > 1).agg(this.c))
        self.assertEqual(q, "SELECT b, c FROM a GROUP BY 1 HAVING (b > 1)")

        q = c.compile(t.select(this.b).group_by(this.b).agg().having(this.b > 1))
        self.assertEqual(q, "SELECT b FROM (SELECT b FROM a) tmp2 GROUP BY 1 HAVING (b > 1)")

        # Having sum
        q = c.compile(t.group_by(this.b).agg(this.c, this.d).having(this.b.sum() > 1))
        self.assertEqual(q, "SELECT b, c, d FROM a GROUP BY 1 HAVING (SUM(b) > 1)")

        # Select interaction
        q = c.compile(t.select(this.a).group_by(this.b).agg(this.c).select(this.c + 1))
        self.assertEqual(q, "SELECT (c + 1) FROM (SELECT b, c FROM (SELECT a FROM a) tmp3 GROUP BY 1) tmp4")

    def test_case_when(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.select(when(this.b).then(this.c)))
        self.assertEqual(q, "SELECT CASE WHEN b THEN c END FROM a")

        q = c.compile(t.select(when(this.b).then(this.c).else_(this.d)))
        self.assertEqual(q, "SELECT CASE WHEN b THEN c ELSE d END FROM a")

        q = c.compile(
            t.select(
                when(this.type == "text")
                .then(this.text)
                .when(this.type == "number")
                .then(this.number)
                .else_("unknown type")
            )
        )
        self.assertEqual(
            q,
            "SELECT CASE WHEN (type = 'text') THEN text WHEN (type = 'number') THEN number ELSE 'unknown type' END FROM a",
        )

    def test_code(self):
        c = Compiler(MockDatabase())
        t = table("a")

        q = c.compile(t.select(this.b, code("<x>")).where(code("<y>")))
        self.assertEqual(q, "SELECT b, <x> FROM a WHERE <y>")

        def tablesample(t, size):
            return code("{t} TABLESAMPLE BERNOULLI ({size})", t=t, size=size)

        nonzero = table("points").where(this.x > 0, this.y > 0)

        q = c.compile(tablesample(nonzero, 10))
        self.assertEqual(q, "SELECT * FROM points WHERE (x > 0) AND (y > 0) TABLESAMPLE BERNOULLI (10)")