#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import time
import unittest
from decimal import Decimal

import pyspark.sql.functions as F
from pyspark.errors import (
    PySparkTypeError,
    PySparkValueError,
)
from pyspark.sql import Row
from pyspark.sql.types import (
    DateType,
    DecimalType,
    IntegerType,
    StringType,
    StructField,
    StructType,
    TimestampNTZType,
    TimestampType,
    TimeType,
)
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import (
    have_pandas,
    have_pyarrow,
    pandas_requirement_message,
    pyarrow_requirement_message,
)


class DataFrameCreationTestsMixin:
    def test_create_str_from_dict(self):
        data = [
            {"broker": {"teamId": 3398, "contactEmail": "abc.xyz@123.ca"}},
        ]

        for schema in [
            StructType([StructField("broker", StringType())]),
            "broker: string",
        ]:
            df = self.spark.createDataFrame(data, schema=schema)
            self.assertEqual(
                df.first().broker,
                """{'teamId': 3398, 'contactEmail': 'abc.xyz@123.ca'}""",
            )

    def test_create_dataframe_from_array_of_long(self):
        import array

        data = [Row(longarray=array.array("l", [-9223372036854775808, 0, 9223372036854775807]))]
        df = self.spark.createDataFrame(data)
        self.assertEqual(df.first(), Row(longarray=[-9223372036854775808, 0, 9223372036854775807]))

    def test_create_dataframe_from_datetime_time(self):
        import datetime

        df = self.spark.createDataFrame(
            [
                (datetime.time(1, 2, 3),),
                (datetime.time(4, 5, 6),),
                (datetime.time(7, 8, 9),),
            ],
            ["t"],
        )
        self.assertIsInstance(df.schema["t"].dataType, TimeType)
        self.assertEqual(df.count(), 3)

    @unittest.skipIf(not have_pandas, pandas_requirement_message)
    def test_create_dataframe_from_pandas_with_timestamp(self):
        from datetime import datetime

        import pandas as pd

        pdf = pd.DataFrame(
            {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]},
            columns=["d", "ts"],
        )
        # test types are inferred correctly without specifying schema
        df = self.spark.createDataFrame(pdf)
        self.assertIsInstance(df.schema["ts"].dataType, TimestampType)
        self.assertIsInstance(df.schema["d"].dataType, DateType)
        # test with schema will accept pdf as input
        df = self.spark.createDataFrame(pdf, schema="d date, ts timestamp")
        self.assertIsInstance(df.schema["ts"].dataType, TimestampType)
        self.assertIsInstance(df.schema["d"].dataType, DateType)
        df = self.spark.createDataFrame(pdf, schema="d date, ts timestamp_ntz")
        self.assertIsInstance(df.schema["ts"].dataType, TimestampNTZType)
        self.assertIsInstance(df.schema["d"].dataType, DateType)

    @unittest.skipIf(have_pandas, "Required Pandas was found.")
    def test_create_dataframe_required_pandas_not_found(self):
        with self.quiet():
            with self.assertRaisesRegex(
                ImportError, "(Pandas >= .* must be installed|No module named '?pandas'?)"
            ):
                from datetime import datetime

                import pandas as pd

                pdf = pd.DataFrame(
                    {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]}
                )
                self.spark.createDataFrame(pdf)

    # Regression test for SPARK-23360
    @unittest.skipIf(not have_pandas, pandas_requirement_message)
    def test_create_dataframe_from_pandas_with_dst(self):
        from datetime import datetime

        import pandas as pd
        from pandas.testing import assert_frame_equal

        pdf = pd.DataFrame({"time": [datetime(2015, 10, 31, 22, 30)]})

        df = self.spark.createDataFrame(pdf)
        assert_frame_equal(pdf, df.toPandas())

        orig_env_tz = os.environ.get("TZ", None)
        try:
            tz = "America/Los_Angeles"
            os.environ["TZ"] = tz
            time.tzset()
            with self.sql_conf({"spark.sql.session.timeZone": tz}):
                df = self.spark.createDataFrame(pdf)
                assert_frame_equal(pdf, df.toPandas())
        finally:
            del os.environ["TZ"]
            if orig_env_tz is not None:
                os.environ["TZ"] = orig_env_tz
            time.tzset()

    @unittest.skipIf(not have_pandas, pandas_requirement_message)
    def test_create_dataframe_from_pandas_with_day_time_interval(self):
        # SPARK-37277: Test DayTimeIntervalType in createDataFrame without Arrow.
        from datetime import timedelta

        import pandas as pd

        df = self.spark.createDataFrame(pd.DataFrame({"a": [timedelta(microseconds=123)]}))
        self.assertEqual(df.toPandas().a.iloc[0], timedelta(microseconds=123))

    # test for SPARK-36337
    def test_create_nan_decimal_dataframe(self):
        self.assertEqual(
            self.spark.createDataFrame(data=[Decimal("NaN")], schema="decimal").collect(),
            [Row(value=None)],
        )

    def test_check_decimal_nan(self):
        data = [Row(dec=Decimal("NaN"))]
        schema = StructType([StructField("dec", DecimalType(), False)])
        with self.assertRaises(PySparkValueError):
            self.spark.createDataFrame(data=data, schema=schema)

    def test_decimal_round(self):
        df = self.spark.createDataFrame([(Decimal(1.234),)], ["d"])
        rounded = df.first().d
        self.assertEqual(rounded, Decimal("1.233999999999999986"))

    def test_invalid_argument_create_dataframe(self):
        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.createDataFrame([(1, 2)], schema=123)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_EXPECTED_TYPE",
            messageParameters={
                "expected_type": "list, None or StructType",
                "arg_name": "schema",
                "arg_type": "int",
            },
        )

        with self.assertRaises(PySparkTypeError) as pe:
            self.spark.createDataFrame(self.spark.range(1))

        self.check_error(
            exception=pe.exception,
            errorClass="INVALID_TYPE",
            messageParameters={"arg_name": "data", "arg_type": "DataFrame"},
        )

    def test_partial_inference_failure(self):
        with self.assertRaises(PySparkValueError) as pe:
            self.spark.createDataFrame([(None, 1)])

        self.check_error(
            exception=pe.exception,
            errorClass="CANNOT_DETERMINE_TYPE",
            messageParameters={},
        )

    @unittest.skipIf(
        not have_pandas or not have_pyarrow,
        pandas_requirement_message or pyarrow_requirement_message,
    )
    def test_schema_inference_from_pandas_with_dict(self):
        # SPARK-47543: test for verifying if inferring `dict` as `MapType` work properly.
        import pandas as pd

        pdf = pd.DataFrame({"str_col": ["second"], "dict_col": [{"first": 0.7, "second": 0.3}]})

        with self.sql_conf(
            {
                "spark.sql.execution.arrow.pyspark.enabled": True,
                "spark.sql.execution.arrow.pyspark.fallback.enabled": False,
                "spark.sql.execution.pandas.inferPandasDictAsMap": True,
            }
        ):
            sdf = self.spark.createDataFrame(pdf)
            self.assertEqual(
                sdf.withColumn("test", F.col("dict_col")[F.col("str_col")]).collect(),
                [Row(str_col="second", dict_col={"first": 0.7, "second": 0.3}, test=0.3)],
            )

            # Empty dict should fail
            pdf_empty_struct = pd.DataFrame({"str_col": ["second"], "dict_col": [{}]})

            with self.assertRaises(PySparkValueError) as pe:
                self.spark.createDataFrame(pdf_empty_struct)

            self.check_error(
                exception=pe.exception,
                errorClass="CANNOT_INFER_EMPTY_SCHEMA",
                messageParameters={},
            )

            # Dict has different types of values should fail
            pdf_different_type = pd.DataFrame(
                {"str_col": ["second"], "dict_col": [{"first": 0.7, "second": "0.3"}]}
            )
            self.assertRaises(
                PySparkValueError, lambda: self.spark.createDataFrame(pdf_different_type)
            )

        with self.sql_conf(
            {
                "spark.sql.execution.arrow.pyspark.enabled": False,
                "spark.sql.execution.pandas.inferPandasDictAsMap": True,
            }
        ):
            sdf = self.spark.createDataFrame(pdf)
            self.assertEqual(
                sdf.withColumn("test", F.col("dict_col")[F.col("str_col")]).collect(),
                [Row(str_col="second", dict_col={"first": 0.7, "second": 0.3}, test=0.3)],
            )

    def test_empty_dataframe_with_struct_type(self):
        schema = StructType(
            [
                StructField("name", StringType(), True),
                StructField("age", IntegerType(), True),
            ]
        )
        df = self.spark.emptyDataFrame(schema)
        self.assertEqual(df.schema, schema)
        self.assertEqual(df.count(), 0)
        self.assertEqual(df.collect(), [])

    def test_empty_dataframe_with_ddl_string(self):
        df = self.spark.emptyDataFrame("name STRING, age INT")
        self.assertEqual(len(df.schema.fields), 2)
        self.assertEqual(df.schema.fields[0].name, "name")
        self.assertIsInstance(df.schema.fields[0].dataType, StringType)
        self.assertEqual(df.schema.fields[1].name, "age")
        self.assertIsInstance(df.schema.fields[1].dataType, IntegerType)
        self.assertEqual(df.count(), 0)
        self.assertEqual(df.collect(), [])

    def test_empty_schema(self):
        schema = StructType()
        for data in [[], [Row()]]:
            with self.subTest(data=data):
                sdf = self.spark.createDataFrame(data, schema)
                assertDataFrameEqual(sdf, data)

    @unittest.skipIf(
        not have_pandas or not have_pyarrow,
        pandas_requirement_message or pyarrow_requirement_message,
    )
    def test_from_pandas_dataframe_with_zero_columns(self):
        """SPARK-55600: Test that row count is preserved when creating DataFrame from
        pandas with 0 columns but with explicit schema in classic Spark."""
        import pandas as pd

        # Create a pandas DataFrame with 5 rows but 0 columns
        pdf = pd.DataFrame(index=range(5))
        schema = StructType([])

        # Test with Arrow optimization enabled
        with self.sql_conf(
            {
                "spark.sql.execution.arrow.pyspark.enabled": True,
                "spark.sql.execution.arrow.pyspark.fallback.enabled": False,
            }
        ):
            df = self.spark.createDataFrame(pdf, schema=schema)
            self.assertEqual(df.schema, schema)
            self.assertEqual(df.count(), 5)
            self.assertEqual(len(df.collect()), 5)

        # Test with Arrow optimization disabled
        with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}):
            df = self.spark.createDataFrame(pdf, schema=schema)
            self.assertEqual(df.schema, schema)
            self.assertEqual(df.count(), 5)
            self.assertEqual(len(df.collect()), 5)


class DataFrameCreationTests(
    DataFrameCreationTestsMixin,
    ReusedSQLTestCase,
):
    pass


if __name__ == "__main__":
    from pyspark.testing import main

    main()