import pytest

_ = pytest.importorskip("duckdb.experimental.spark")

from spark_namespace import USE_ACTUAL_SPARK
from spark_namespace.sql.types import (
    LongType,
    StructType,
    BooleanType,
    StructField,
    StringType,
    IntegerType,
    LongType,
    Row,
    ArrayType,
    MapType,
)
from spark_namespace.sql.functions import col, struct, when
from spark_namespace.sql.column import Column
import duckdb
import re

from spark_namespace.errors import PySparkValueError, PySparkTypeError


def assert_column_objects_equal(col1: Column, col2: Column):
    assert type(col1) == type(col2)
    if not USE_ACTUAL_SPARK:
        assert col1.expr == col2.expr


class TestDataFrame(object):
    def test_dataframe_from_list_of_tuples(self, spark):
        # Valid
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "CA")]
        df = spark.createDataFrame(address, ["id", "address", "state"])
        res = df.collect()
        assert res == [
            Row(id=1, address='14851 Jeffrey Rd', state='DE'),
            Row(id=2, address='43421 Margarita St', state='NY'),
            Row(id=3, address='13111 Siemon Ave', state='CA'),
        ]

        # Tuples of different sizes
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave")]

        if USE_ACTUAL_SPARK:
            from py4j.protocol import Py4JJavaError

            with pytest.raises(Py4JJavaError):
                df = spark.createDataFrame(address, ["id", "address", "state"])
                df.collect()
        else:
            with pytest.raises(PySparkTypeError, match="LENGTH_SHOULD_BE_THE_SAME"):
                df = spark.createDataFrame(address, ["id", "address", "state"])

        # Dataframe instead of list
        with pytest.raises(PySparkTypeError, match="SHOULD_NOT_DATAFRAME"):
            df = spark.createDataFrame(df, ["id", "address", "state"])

        # Not a list
        with pytest.raises(TypeError, match="not iterable"):
            df = spark.createDataFrame(5, ["id", "address", "test"])

        # Empty list
        if not USE_ACTUAL_SPARK:
            # FIXME: Spark raises PySparkValueError [CANNOT_INFER_EMPTY_SCHEMA]
            df = spark.createDataFrame([], ["id", "address", "test"])
            res = df.collect()
            assert res == []

        # Duplicate column names
        address = [(1, "14851 Jeffrey Rd", "DE"), (2, "43421 Margarita St", "NY"), (3, "13111 Siemon Ave", "DE")]
        df = spark.createDataFrame(address, ["id", "address", "id"])
        res = df.collect()
        exptected_res_str = "[Row(id=1, address='14851 Jeffrey Rd', id='DE'), Row(id=2, address='43421 Margarita St', id='NY'), Row(id=3, address='13111 Siemon Ave', id='DE')]"
        if USE_ACTUAL_SPARK:
            # Spark uses string for both ID columns. DuckDB correctly infers the types.
            exptected_res_str = (
                exptected_res_str.replace("id=1", "id='1'").replace("id=2", "id='2'").replace("id=3", "id='3'")
            )
        assert str(res) == exptected_res_str

        # Not enough column names
        if not USE_ACTUAL_SPARK:
            # FIXME: Spark does not raise this error
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                df = spark.createDataFrame(address, ["id", "address"])

        # Empty column names list
        # Columns are filled in with default names
        # TODO: check against Spark behavior
        df = spark.createDataFrame(address, [])
        res = df.collect()
        assert res == [
            Row(col0=1, col1='14851 Jeffrey Rd', col2='DE'),
            Row(col0=2, col1='43421 Margarita St', col2='NY'),
            Row(col0=3, col1='13111 Siemon Ave', col2='DE'),
        ]

        # Too many column names
        if not USE_ACTUAL_SPARK:
            # In Spark, this leads to an IndexError
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                df = spark.createDataFrame(address, ["id", "address", "one", "two", "three"])

        # Column names is not a list (but is iterable)
        if not USE_ACTUAL_SPARK:
            # These things do not work in Spark or throw different errors
            df = spark.createDataFrame(address, {'a': 5, 'b': 6, 'c': 42})
            res = df.collect()
            assert res == [
                Row(a=1, b='14851 Jeffrey Rd', c='DE'),
                Row(a=2, b='43421 Margarita St', c='NY'),
                Row(a=3, b='13111 Siemon Ave', c='DE'),
            ]

            # Column names is not a list (string, becomes a single column name)
            with pytest.raises(PySparkValueError, match="number of columns in the DataFrame don't match"):
                df = spark.createDataFrame(address, 'a')

            with pytest.raises(TypeError, match="must be an iterable, not int"):
                df = spark.createDataFrame(address, 5)

    def test_dataframe(self, spark):
        # Create DataFrame
        df = spark.createDataFrame([("Scala", 25000), ("Spark", 35000), ("PHP", 21000)])
        res = df.collect()
        assert res == [Row(col0='Scala', col1=25000), Row(col0='Spark', col1=35000), Row(col0='PHP', col1=21000)]

    @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="We can't create tables with our Spark test setup")
    def test_writing_to_table(self, spark):
        # Create Hive table & query it.
        spark.sql(
            """
            create table sample_table("_1" bool, "_2" integer)
        """
        )
        spark.sql('insert into sample_table VALUES (True, 42)')
        spark.table("sample_table").write.saveAsTable("sample_hive_table")
        df3 = spark.sql("SELECT _1,_2 FROM sample_hive_table")
        res = df3.collect()
        assert res == [Row(_1=True, _2=42)]
        schema = df3.schema
        assert schema == StructType([StructField('_1', BooleanType(), True), StructField('_2', IntegerType(), True)])

    def test_dataframe_collect(self, spark):
        df = spark.createDataFrame([(42,), (21,)]).toDF('a')
        res = df.collect()
        assert str(res) == '[Row(a=42), Row(a=21)]'

    def test_dataframe_from_rows(self, spark):
        columns = ["language", "users_count"]
        data = [("Java", "20000"), ("Python", "100000"), ("Scala", "3000")]

        rowData = map(lambda x: Row(*x), data)
        df = spark.createDataFrame(rowData, columns)
        res = df.collect()
        assert res == [
            Row(language='Java', users_count='20000'),
            Row(language='Python', users_count='100000'),
            Row(language='Scala', users_count='3000'),
        ]

    def test_empty_df(self, spark):
        schema = StructType(
            [
                StructField('firstname', StringType(), True),
                StructField('middlename', StringType(), True),
                StructField('lastname', StringType(), True),
            ]
        )
        df = spark.createDataFrame([], schema=schema)
        res = df.collect()
        # TODO: assert that the types and column names are correct
        assert res == []

    def test_df_from_pandas(self, spark):
        import pandas as pd

        df = spark.createDataFrame(pd.DataFrame({'a': [42, 21], 'b': [True, False]}))
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_from_struct_type(self, spark):
        schema = StructType([StructField('a', LongType()), StructField('b', BooleanType())])
        df = spark.createDataFrame([(42, True), (21, False)], schema)
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_from_name_list(self, spark):
        df = spark.createDataFrame([(42, True), (21, False)], ['a', 'b'])
        res = df.collect()
        assert res == [Row(a=42, b=True), Row(a=21, b=False)]

    def test_df_creation_coverage(self, spark):
        from spark_namespace.sql.types import StructType, StructField, StringType, IntegerType

        data2 = [
            ("James", "", "Smith", "36636", "M", 3000),
            ("Michael", "Rose", "", "40288", "M", 4000),
            ("Robert", "", "Williams", "42114", "M", 4000),
            ("Maria", "Anne", "Jones", "39192", "F", 4000),
            ("Jen", "Mary", "Brown", "", "F", -1),
        ]

        schema = StructType(
            [
                StructField("firstname", StringType(), True),
                StructField("middlename", StringType(), True),
                StructField("lastname", StringType(), True),
                StructField("id", StringType(), True),
                StructField("gender", StringType(), True),
                StructField("salary", IntegerType(), True),
            ]
        )

        df = spark.createDataFrame(data=data2, schema=schema)
        res = df.collect()
        assert res == [
            Row(firstname='James', middlename='', lastname='Smith', id='36636', gender='M', salary=3000),
            Row(firstname='Michael', middlename='Rose', lastname='', id='40288', gender='M', salary=4000),
            Row(firstname='Robert', middlename='', lastname='Williams', id='42114', gender='M', salary=4000),
            Row(firstname='Maria', middlename='Anne', lastname='Jones', id='39192', gender='F', salary=4000),
            Row(firstname='Jen', middlename='Mary', lastname='Brown', id='', gender='F', salary=-1),
        ]

    def test_df_nested_struct(self, spark):
        structureData = [
            (("James", "", "Smith"), "36636", "M", 3100),
            (("Michael", "Rose", ""), "40288", "M", 4300),
            (("Robert", "", "Williams"), "42114", "M", 1400),
            (("Maria", "Anne", "Jones"), "39192", "F", 5500),
            (("Jen", "Mary", "Brown"), "", "F", -1),
        ]
        structureSchema = StructType(
            [
                StructField(
                    'name',
                    StructType(
                        [
                            StructField('firstname', StringType(), True),
                            StructField('middlename', StringType(), True),
                            StructField('lastname', StringType(), True),
                        ]
                    ),
                ),
                StructField('id', StringType(), True),
                StructField('gender', StringType(), True),
                StructField('salary', IntegerType(), True),
            ]
        )

        df2 = spark.createDataFrame(data=structureData, schema=structureSchema)
        res = df2.collect()
        expected_res = [
            Row(
                name={'firstname': 'James', 'middlename': '', 'lastname': 'Smith'}, id='36636', gender='M', salary=3100
            ),
            Row(
                name={'firstname': 'Michael', 'middlename': 'Rose', 'lastname': ''}, id='40288', gender='M', salary=4300
            ),
            Row(
                name={'firstname': 'Robert', 'middlename': '', 'lastname': 'Williams'},
                id='42114',
                gender='M',
                salary=1400,
            ),
            Row(
                name={'firstname': 'Maria', 'middlename': 'Anne', 'lastname': 'Jones'},
                id='39192',
                gender='F',
                salary=5500,
            ),
            Row(name={'firstname': 'Jen', 'middlename': 'Mary', 'lastname': 'Brown'}, id='', gender='F', salary=-1),
        ]
        if USE_ACTUAL_SPARK:
            expected_res = [Row(name=Row(**r.name), id=r.id, gender=r.gender, salary=r.salary) for r in expected_res]
        assert res == expected_res
        schema = df2.schema
        assert schema == StructType(
            [
                StructField(
                    'name',
                    StructType(
                        [
                            StructField('firstname', StringType(), True),
                            StructField('middlename', StringType(), True),
                            StructField('lastname', StringType(), True),
                        ]
                    ),
                    True,
                ),
                StructField('id', StringType(), True),
                StructField('gender', StringType(), True),
                StructField('salary', IntegerType(), True),
            ]
        )

    def test_df_columns(self, spark):
        from spark_namespace.sql.functions import col, struct, when

        structureData = [
            (("James", "", "Smith"), "36636", "M", 3100),
            (("Michael", "Rose", ""), "40288", "M", 4300),
            (("Robert", "", "Williams"), "42114", "M", 1400),
            (("Maria", "Anne", "Jones"), "39192", "F", 5500),
            (("Jen", "Mary", "Brown"), "", "F", -1),
        ]
        structureSchema = StructType(
            [
                StructField(
                    'name',
                    StructType(
                        [
                            StructField('firstname', StringType(), True),
                            StructField('middlename', StringType(), True),
                            StructField('lastname', StringType(), True),
                        ]
                    ),
                ),
                StructField('id', StringType(), True),
                StructField('gender', StringType(), True),
                StructField('salary', IntegerType(), True),
            ]
        )

        df2 = spark.createDataFrame(data=structureData, schema=structureSchema)
        updatedDF = df2.withColumn(
            "OtherInfo",
            struct(
                col("id").alias("identifier"),
                col("gender").alias("gender"),
                col("salary").alias("salary"),
                when(col("salary").cast(IntegerType()) < 2000, "Low")
                .when(col("salary").cast(IntegerType()) < 4000, "Medium")
                .otherwise("High")
                .alias("Salary_Grade"),
            ),
        ).drop("id", "gender", "salary")

        assert 'OtherInfo' in updatedDF.columns

    def test_array_and_map_type(self, spark):
        """Array & Map"""

        arrayStructureSchema = StructType(
            [
                StructField(
                    'name',
                    StructType(
                        [
                            StructField('firstname', StringType(), True),
                            StructField('middlename', StringType(), True),
                            StructField('lastname', StringType(), True),
                        ]
                    ),
                ),
                StructField('hobbies', ArrayType(StringType()), True),
                StructField('properties', MapType(StringType(), StringType()), True),
            ]
        )

    def test_getitem_filter(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [
            Row(age=56, name="Carol"),
            Row(age=20, name="Alice"),
        ]

        df = df[col("age") > 18]
        assert df.collect() == expected

    def test_getitem_column(self, spark):
        data = [(56, "Carol"), (20, "Alice")]
        df = spark.createDataFrame(data, ["age", "name"])

        assert_column_objects_equal(df["age"], col("age"))

    def test_getitem_dataframe(self, spark):
        data = [(56, "Ben", "Street1"), (20, "Tom", "Street2")]
        df = spark.createDataFrame(data, ["age", "name", "address"])

        res = df[["name", "address"]].collect()
        expected = df.select("name", "address").collect()

        assert res == expected

    def test_getattr_dataframe(self, spark):
        data = [(56, "Ben", "Street1"), (20, "Tom", "Street2")]
        df = spark.createDataFrame(data, ["age", "name", "address"])

        assert_column_objects_equal(df.age, col("age"))

    def test_head_first(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = Row(age=1, name="Ben")

        head = df.orderBy(df.age).head()
        first = df.orderBy(df.age).first()
        assert head == first == expected

    def test_head_take_n(self, spark):
        data = [(56, "Carol"), (20, "Alice"), (3, "Dave"), (3, "Anna"), (1, "Ben")]
        df = spark.createDataFrame(data, ["age", "name"])
        expected = [
            Row(age=1, name="Ben"),
            Row(age=3, name="Anna"),
        ]
        df = df.orderBy(df.age, df.name)
        rows = df.head(2)
        take = df.take(2)
        assert rows == take == expected

    def test_drop(self, spark):
        data = [(1, 2, 3, 4)]
        df = spark.createDataFrame(data, ["one", "two", "three", "four"])
        expected = ["one", "four"]
        assert df.drop("two", "three").columns == expected
        assert df.drop("two", col("three")).columns == expected
        assert df.drop("two", col("three"), col("missing")).columns == expected
