import uuid
from collections.abc import Iterable, Sized
from typing import TYPE_CHECKING, Any, NoReturn, Union

import duckdb

if TYPE_CHECKING:
    from pandas.core.frame import DataFrame as PandasDataFrame

    from .catalog import Catalog


from ..conf import SparkConf
from ..context import SparkContext
from ..errors import PySparkTypeError
from ..exception import ContributionsAcceptedError
from .conf import RuntimeConfig
from .dataframe import DataFrame
from .readwriter import DataFrameReader
from .streaming import DataStreamReader
from .types import StructType
from .udf import UDFRegistration

# In spark:
# SparkSession holds a SparkContext
# SparkContext gets created from SparkConf
# At this level the check is made to determine whether the instance already exists and just needs
# to be retrieved or it needs to be created.

# For us this is done inside of `duckdb.connect`, based on the passed in path + configuration
# SparkContext can be compared to our Connection class, and SparkConf to our ClientContext class


# data is a List of rows
# every value in each row needs to be turned into a Value
def _combine_data_and_schema(data: Iterable[Any], schema: StructType) -> list[duckdb.Value]:
    from duckdb import Value

    new_data = []
    for row in data:
        new_row = [Value(x, dtype.duckdb_type) for x, dtype in zip(row, [y.dataType for y in schema], strict=False)]
        new_data.append(new_row)
    return new_data


class SparkSession:  # noqa: D101
    def __init__(self, context: SparkContext) -> None:  # noqa: D107
        self.conn = context.connection
        self._context = context
        self._conf = RuntimeConfig(self.conn)

    def _create_dataframe(self, data: Union[Iterable[Any], "PandasDataFrame"]) -> DataFrame:
        try:
            import pandas

            has_pandas = True
        except ImportError:
            has_pandas = False
        if has_pandas and isinstance(data, pandas.DataFrame):
            unique_name = f"pyspark_pandas_df_{uuid.uuid1()}"
            self.conn.register(unique_name, data)
            return DataFrame(self.conn.sql(f'select * from "{unique_name}"'), self)

        def verify_tuple_integrity(tuples: list[tuple]) -> None:
            if len(tuples) <= 1:
                return
            expected_length = len(tuples[0])
            for i, item in enumerate(tuples[1:]):
                actual_length = len(item)
                if expected_length == actual_length:
                    continue
                raise PySparkTypeError(
                    error_class="LENGTH_SHOULD_BE_THE_SAME",
                    message_parameters={
                        "arg1": f"data{i}",
                        "arg2": f"data{i + 1}",
                        "arg1_length": str(expected_length),
                        "arg2_length": str(actual_length),
                    },
                )

        if not isinstance(data, list):
            data = list(data)
        verify_tuple_integrity(data)

        def construct_query(tuples: Iterable) -> str:
            def construct_values_list(row: Sized, start_param_idx: int) -> str:
                parameter_count = len(row)
                parameters = [f"${x + start_param_idx}" for x in range(parameter_count)]
                parameters = "(" + ", ".join(parameters) + ")"
                return parameters

            row_size = len(tuples[0])
            values_list = [construct_values_list(x, 1 + (i * row_size)) for i, x in enumerate(tuples)]
            values_list = ", ".join(values_list)

            query = f"""
                select * from (values {values_list})
            """
            return query

        query = construct_query(data)

        def construct_parameters(tuples: Iterable) -> list[list]:
            parameters = []
            for row in tuples:
                parameters.extend(list(row))
            return parameters

        parameters = construct_parameters(data)

        rel = self.conn.sql(query, params=parameters)
        return DataFrame(rel, self)

    def _createDataFrameFromPandas(
        self, data: "PandasDataFrame", types: list[str] | None, names: list[str] | None
    ) -> DataFrame:
        df = self._create_dataframe(data)

        # Cast to types
        if types:
            df = df._cast_types(*types)
        # Alias to names
        if names:
            df = df.toDF(*names)
        return df

    def createDataFrame(  # noqa: D102
        self,
        data: Union["PandasDataFrame", Iterable[Any]],
        schema: StructType | list[str] | None = None,
        samplingRatio: float | None = None,
        verifySchema: bool = True,
    ) -> DataFrame:
        if samplingRatio:
            raise NotImplementedError
        if not verifySchema:
            raise NotImplementedError
        types = None
        names = None

        if isinstance(data, DataFrame):
            raise PySparkTypeError(
                error_class="SHOULD_NOT_DATAFRAME",
                message_parameters={"arg_name": "data"},
            )

        if schema:
            if isinstance(schema, StructType):
                types, names = schema.extract_types_and_names()
            else:
                names = schema

        try:
            import pandas

            has_pandas = True
        except ImportError:
            has_pandas = False
        # Falsey check on pandas dataframe is not defined, so first check if it's not a pandas dataframe
        # Then check if 'data' is None or []
        if has_pandas and isinstance(data, pandas.DataFrame):
            return self._createDataFrameFromPandas(data, types, names)

        # Finally check if a schema was provided
        is_empty = False
        if not data and names:
            # Create NULLs for every type in our dataframe
            is_empty = True
            data = [tuple(None for _ in names)]

        if schema and isinstance(schema, StructType):
            # Transform the data into Values to combine the data+schema
            data = _combine_data_and_schema(data, schema)

        df = self._create_dataframe(data)
        if is_empty:
            rel = df.relation
            # Add impossible where clause
            rel = rel.filter("1=0")
            df = DataFrame(rel, self)

        # Cast to types
        if types:
            df = df._cast_types(*types)
        # Alias to names
        if names:
            df = df.toDF(*names)
        return df

    def newSession(self) -> "SparkSession":  # noqa: D102
        return SparkSession(self._context)

    def range(  # noqa: D102
        self,
        start: int,
        end: int | None = None,
        step: int = 1,
        numPartitions: int | None = None,
    ) -> "DataFrame":
        if numPartitions:
            raise ContributionsAcceptedError

        if end is None:
            end = start
            start = 0

        return DataFrame(self.conn.table_function("range", parameters=[start, end, step]), self)

    def sql(self, sqlQuery: str, **kwargs: Any) -> DataFrame:  # noqa: D102, ANN401
        if kwargs:
            raise NotImplementedError
        relation = self.conn.sql(sqlQuery)
        return DataFrame(relation, self)

    def stop(self) -> None:  # noqa: D102
        self._context.stop()

    def table(self, tableName: str) -> DataFrame:  # noqa: D102
        relation = self.conn.table(tableName)
        return DataFrame(relation, self)

    def getActiveSession(self) -> "SparkSession":  # noqa: D102
        return self

    @property
    def catalog(self) -> "Catalog":  # noqa: D102
        if not hasattr(self, "_catalog"):
            from duckdb.experimental.spark.sql.catalog import Catalog

            self._catalog = Catalog(self)
        return self._catalog

    @property
    def conf(self) -> RuntimeConfig:  # noqa: D102
        return self._conf

    @property
    def read(self) -> DataFrameReader:  # noqa: D102
        return DataFrameReader(self)

    @property
    def readStream(self) -> DataStreamReader:  # noqa: D102
        return DataStreamReader(self)

    @property
    def sparkContext(self) -> SparkContext:  # noqa: D102
        return self._context

    @property
    def streams(self) -> NoReturn:  # noqa: D102
        raise ContributionsAcceptedError

    @property
    def udf(self) -> UDFRegistration:  # noqa: D102
        return UDFRegistration(self)

    @property
    def version(self) -> str:  # noqa: D102
        return "1.0.0"

    class Builder:  # noqa: D106
        def __init__(self) -> None:  # noqa: D107
            pass

        def master(self, name: str) -> "SparkSession.Builder":  # noqa: D102
            # no-op
            return self

        def appName(self, name: str) -> "SparkSession.Builder":  # noqa: D102
            # no-op
            return self

        def remote(self, url: str) -> "SparkSession.Builder":  # noqa: D102
            # no-op
            return self

        def getOrCreate(self) -> "SparkSession":  # noqa: D102
            context = SparkContext("__ignored__")
            return SparkSession(context)

        def config(  # noqa: D102
            self,
            key: str | None = None,
            value: Any | None = None,  # noqa: ANN401
            conf: SparkConf | None = None,
        ) -> "SparkSession.Builder":
            return self

        def enableHiveSupport(self) -> "SparkSession.Builder":  # noqa: D102
            # no-op
            return self

    builder = Builder()


__all__ = ["SparkSession"]
