import duckdb
import numpy as np
import tempfile
import os
import pandas as pd
import pytest

def get_relation(conn):
    test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
    conn.register("test_df", test_df)
    return conn.from_df(test_df)

class TestRelation(object):
    def test_csv_auto(self, duckdb_cursor):
        conn = duckdb.connect()
        df_rel = get_relation(conn)
        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        test_df.to_csv(temp_file_name, index=False)

        # now create a relation from it
        csv_rel = duckdb.from_csv_auto(temp_file_name)
        assert df_rel.execute().fetchall() == csv_rel.execute().fetchall()

    def test_filter_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.filter('i > 1').execute().fetchall() == [(2, 'two'), (3, 'three'), (4, 'four')]

    def test_projection_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.project('i').execute().fetchall() == [(1,), (2,), (3,), (4,)]

    def test_projection_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.order('j').execute().fetchall() == [(4, 'four'), (1, 'one'), (3, 'three'), (2, 'two')]

    def test_limit_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.limit(2).execute().fetchall() == [(1, 'one'), (2, 'two')]
        assert rel.limit(2, offset=1).execute().fetchall() == [(2, 'two'), (3, 'three')]

    def test_intersect_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4]})
        conn.register("test_df", test_df)
        test_df_2 = pd.DataFrame.from_dict({"i":[3, 4, 5, 6]})
        conn.register("test_df", test_df_2)
        rel = conn.from_df(test_df)
        rel_2 = conn.from_df(test_df_2)

        assert rel.intersect(rel_2).execute().fetchall() == [ (3,), (4,)]

    def test_aggregate_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.aggregate("sum(i)").execute().fetchall() == [(10,)]
        assert rel.aggregate("j, sum(i)").execute().fetchall() == [('one', 1), ('two', 2), ('three', 3), ('four', 4)]

    def test_distinct_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        assert rel.distinct().execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'),(4, 'four')]

    def test_union_operator(self, duckdb_cursor):
        conn = duckdb.connect()
        rel = get_relation(conn)
        print(rel.union(rel).execute().fetchall())
        assert rel.union(rel).execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'), (4, 'four'), (1, 'one'), (2, 'two'), (3, 'three'), (4, 'four')]

    def test_join_operator(self, duckdb_cursor):
        # join rel with itself on i
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = conn.from_df(test_df)
        rel2 = conn.from_df(test_df)
        assert rel.join(rel2, 'i').execute().fetchall()  == [(1, 'one', 'one'), (2, 'two', 'two'), (3, 'three', 'three'), (4, 'four', 'four')]

    def test_except_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = conn.from_df(test_df)
        rel2 = conn.from_df(test_df)
        assert rel.except_(rel2).execute().fetchall() == []

    def test_create_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = conn.from_df(test_df)
        rel.create("test_df")
        assert conn.query("select * from test_df").execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'),(4, 'four')]

    def test_create_view_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = conn.from_df(test_df)
        rel.create_view("test_df")
        assert conn.query("select * from test_df").execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'),(4, 'four')]

    def test_insert_into_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = conn.from_df(test_df)
        rel.create("test_table2")
        # insert the relation's data into an existing table
        conn.execute("CREATE TABLE test_table3 (i INTEGER, j STRING)")
        rel.insert_into("test_table3")

        # Inserting elements into table_3
        print(conn.values([5, 'five']).insert_into("test_table3"))
        rel_3 = conn.table("test_table3")
        rel_3.insert([6,'six'])

        assert rel_3.execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'), (4, 'four'), (5, 'five'), (6, 'six')]

    def test_write_csv_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        df_rel = get_relation(conn)
        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
        df_rel.write_csv(temp_file_name)

        csv_rel = duckdb.from_csv_auto(temp_file_name)
        assert df_rel.execute().fetchall() == csv_rel.execute().fetchall()

    def test_get_attr_operator(self,duckdb_cursor):
        conn = duckdb.connect()
        conn.execute("CREATE TABLE test (i INTEGER)")
        rel = conn.table("test")
        assert rel.alias == "test"
        assert rel.type == "TABLE_RELATION"
        assert rel.columns == ['i']
        assert rel.types == ['INTEGER']

    def test_query_fail(self,duckdb_cursor):
        conn = duckdb.connect()
        conn.execute("CREATE TABLE test (i INTEGER)")
        rel = conn.table("test")
        with pytest.raises(TypeError, match='incompatible function arguments'):
            rel.query("select j from test")

    def test_execute_fail(self,duckdb_cursor):
        conn = duckdb.connect()
        conn.execute("CREATE TABLE test (i INTEGER)")
        rel = conn.table("test")
        with pytest.raises(TypeError, match='incompatible function arguments'):
            rel.execute("select j from test")

    def test_df_proj(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = duckdb.project(test_df, 'i')
        assert rel.execute().fetchall() == [(1,), (2,), (3,), (4,)]

    def test_df_alias(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = duckdb.alias(test_df, 'dfzinho')
        assert rel.alias == "dfzinho"

    def test_df_filter(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = duckdb.filter(test_df, 'i > 1')
        assert rel.execute().fetchall() == [(2, 'two'), (3, 'three'), (4, 'four')]

    def test_df_order_by(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = duckdb.order(test_df, 'j')
        assert rel.execute().fetchall() == [(4, 'four'), (1, 'one'), (3, 'three'), (2, 'two')]

    def test_df_distinct(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        rel = duckdb.distinct(test_df)
        assert rel.execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'),(4, 'four')]

    def test_df_write_csv(self,duckdb_cursor):
        test_df = pd.DataFrame.from_dict({"i":[1, 2, 3, 4], "j":["one", "two", "three", "four"]})
        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
        duckdb.write_csv(test_df, temp_file_name)
        csv_rel = duckdb.from_csv_auto(temp_file_name)
        assert  csv_rel.execute().fetchall() == [(1, 'one'), (2, 'two'), (3, 'three'), (4, 'four')]


    def test_join_types(self, duckdb_cursor):
        test_df1 = pd.DataFrame.from_dict({"i":[1, 2, 3, 4]})
        test_df2 = pd.DataFrame.from_dict({"j":[  3, 4, 5, 6]})
        rel1 = duckdb_cursor.from_df(test_df1)
        rel2 = duckdb_cursor.from_df(test_df2)

        assert rel1.join(rel2, 'i=j', 'inner').aggregate('count()').fetchone()[0] == 2

        assert rel1.join(rel2, 'i=j', 'left').aggregate('count()').fetchone()[0] == 4

    def test_explain(self, duckdb_cursor):
        con = duckdb.connect()
        con.execute("Create table t1 (i integer)")
        con.execute("Create table t2 (j integer)")
        rel1 = con.table('t1')
        rel2 = con.table('t2')
        join = rel1.join(rel2, 'i=j', 'inner').aggregate('count()')
        assert join.explain() == 'Aggregate [count_star()]\n  Join INNER (i = j)\n    Scan Table [t1]\n    Scan Table [t2]'

    def test_fetchnumpy(self, duckdb_cursor):
        start, stop = -1000, 2000
        count = stop - start

        con = duckdb.connect()
        con.execute(f"CREATE table t AS SELECT range AS a FROM range({start}, {stop});")
        rel = con.table("t")

        # empty
        res = rel.limit(0, offset=count + 1).fetchnumpy()
        assert set(res.keys()) == {"a"}
        assert len(res["a"]) == 0

        # < vector_size, == vector_size, > vector_size
        for size in [1000, 1024, 1100]:
            res = rel.project("a").limit(size).fetchnumpy()
            assert set(res.keys()) == {"a"}
            # For some reason, this return a masked array. Shouldn't it be
            # known that there can't be NULLs?
            if isinstance(res, np.ma.MaskedArray):
                assert res.count() == size
                res = res.compressed()
            else:
                assert len(res["a"]) == size
            assert np.all(res["a"] == np.arange(start, start + size))

        with pytest.raises(duckdb.ConversionException, match="Conversion Error.*out of range.*"):
            # invalid conversion of negative integer to UINTEGER
            rel.project("CAST(a as UINTEGER)").fetchnumpy()

