"""
SparkSession + SparkContext — entry point for singlespark.

Usage mirrors real PySpark:
    spark = SparkSession.builder.appName("MyApp").getOrCreate()
    df = spark.createDataFrame([("Alice", 30)], ["name", "age"])
    spark.sql("SELECT * FROM people").show()
    spark.stop()
"""
from __future__ import annotations

import threading
from typing import Any

import polars as pl

from singlespark.dataframe import DataFrame, _list_to_polars
from singlespark.hdfs import FakeHDFS
from singlespark.sql.catalog import Catalog


class SparkSession:
    """
    Fake SparkSession — mirrors pyspark.sql.SparkSession.
    Backed by Polars for DataFrames and DuckDB for SQL.
    """

    _active: "SparkSession | None" = None
    _lock = threading.Lock()

    # ------------------------------------------------------------------
    # Builder
    # ------------------------------------------------------------------

    class Builder:
        def __init__(self):
            self._app_name: str = "singlespark-app"
            self._config: dict[str, Any] = {}
            self._master: str = "local[*]"

        def appName(self, name: str) -> "SparkSession.Builder":
            self._app_name = name
            return self

        def master(self, url: str) -> "SparkSession.Builder":
            self._master = url  # accepted but ignored
            return self

        def config(self, key: str, value: Any = None, conf=None) -> "SparkSession.Builder":
            if conf is not None:
                self._config.update(conf._config)
            else:
                self._config[key] = value
            return self

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

        def getOrCreate(self) -> "SparkSession":
            with SparkSession._lock:
                if SparkSession._active is None:
                    SparkSession._active = SparkSession(self._app_name, dict(self._config))
                return SparkSession._active

        def create(self) -> "SparkSession":
            """Always create a new session (even if one already exists)."""
            with SparkSession._lock:
                session = SparkSession(self._app_name, dict(self._config))
                SparkSession._active = session
                return session

    builder: "SparkSession.Builder" = Builder()

    # ------------------------------------------------------------------
    # Constructor
    # ------------------------------------------------------------------

    def __init__(self, app_name: str = "singlespark-app", config: dict | None = None):
        self._app_name = app_name
        self._config: dict[str, Any] = config or {}
        self.catalog = Catalog()
        self._hdfs = FakeHDFS(self._config.get("spark.singlespark.hdfs.root"))
        self._sc: "SparkContext | None" = None

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def version(self) -> str:
        return "3.5.0-singlespark"

    @property
    def sparkContext(self) -> "SparkContext":
        if self._sc is None:
            self._sc = SparkContext(self)
        return self._sc

    @property
    def read(self) -> "DataFrameReader":
        from singlespark.reader import DataFrameReader
        return DataFrameReader(self)

    @property
    def readStream(self):
        raise NotImplementedError("Streaming is not supported in singlespark")

    @property
    def conf(self) -> "SparkConf":
        from singlespark.conf import SparkConf
        return SparkConf(self._config)

    @property
    def udf(self):
        return _UDFRegistration(self)

    # ------------------------------------------------------------------
    # Core methods
    # ------------------------------------------------------------------

    def sql(self, query: str, **kwargs) -> DataFrame:
        """Execute a SQL query using DuckDB and return a DataFrame."""
        result = self.catalog.sql(query)
        return DataFrame(result, self)

    def table(self, name: str) -> DataFrame:
        """Return a DataFrame from a registered temp view."""
        return self.sql(f"SELECT * FROM {name}")

    def createDataFrame(
        self,
        data,
        schema=None,
        samplingRatio: float | None = None,
        verifySchema: bool = True,
    ) -> DataFrame:
        """
        Create a DataFrame from various data sources.

        Parameters
        ----------
        data :
            - list of tuples/dicts/Rows
            - polars.DataFrame
            - pandas.DataFrame
        schema :
            - StructType
            - list of column name strings
            - None (infer)
        """
        if isinstance(data, pl.DataFrame):
            df = data
        elif hasattr(data, "to_frame"):
            # pandas DataFrame or Series
            df = pl.from_pandas(data)
        elif isinstance(data, list):
            df = _list_to_polars(data, schema)
            # schema already applied in _list_to_polars
            return DataFrame(df, self)
        else:
            raise TypeError(
                f"createDataFrame does not support data of type {type(data).__name__}.\n"
                "Accepted types: list, polars.DataFrame, pandas.DataFrame"
            )

        # Apply column rename if schema is a list of names
        if schema is not None:
            from singlespark.sql.types import StructType
            if isinstance(schema, (list, tuple)) and all(isinstance(s, str) for s in schema):
                df = df.rename(dict(zip(df.columns, schema)))
            elif isinstance(schema, StructType):
                from singlespark.sql.types import spark_type_to_polars
                casts = [
                    pl.col(f.name).cast(spark_type_to_polars(f.dataType))
                    for f in schema.fields
                    if f.name in df.columns
                ]
                if casts:
                    df = df.with_columns(casts)

        return DataFrame(df, self)

    def range(
        self,
        start: int,
        end: int | None = None,
        step: int = 1,
        numPartitions: int | None = None,
    ) -> DataFrame:
        """Create a DataFrame with a single 'id' column from a range."""
        if end is None:
            start, end = 0, start
        ids = list(range(start, end, step))
        return DataFrame(pl.DataFrame({"id": ids}), self)

    # ------------------------------------------------------------------
    # Session lifecycle
    # ------------------------------------------------------------------

    def stop(self) -> None:
        """Stop this SparkSession and release resources."""
        with SparkSession._lock:
            if SparkSession._active is self:
                SparkSession._active = None
        self.catalog.close()

    def newSession(self) -> "SparkSession":
        """Create a new SparkSession sharing the same config."""
        return SparkSession(self._app_name, dict(self._config))

    @classmethod
    def getActiveSession(cls) -> "SparkSession | None":
        return cls._active

    @classmethod
    def active(cls) -> "SparkSession":
        session = cls._active
        if session is None:
            raise RuntimeError("No active SparkSession. Call SparkSession.builder.getOrCreate() first.")
        return session

    # ------------------------------------------------------------------
    # Repr
    # ------------------------------------------------------------------

    def __repr__(self):
        return f"SparkSession(app='{self._app_name}', version='{self.version}')"

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.stop()


# ---------------------------------------------------------------------------
# SparkContext stub
# ---------------------------------------------------------------------------

class SparkContext:
    """
    Minimal SparkContext stub.
    Many PySpark programs access sc = spark.sparkContext.
    """

    def __init__(self, session: SparkSession):
        self._session = session

    @property
    def appName(self) -> str:
        return self._session._app_name

    @property
    def master(self) -> str:
        return "local[*]"

    @property
    def version(self) -> str:
        return self._session.version

    def textFile(self, path: str, minPartitions: int = 1) -> list[str]:
        """Read a text file and return lines as a list (fake RDD)."""
        from pathlib import Path
        resolved = str(self._session._hdfs.resolve(path))
        return Path(resolved).read_text(encoding="utf-8").splitlines()

    def parallelize(self, data, numSlices: int | None = None) -> list:
        """Return data as-is (fake RDD = plain list)."""
        return list(data)

    def stop(self) -> None:
        self._session.stop()

    def getConf(self):
        return self._session.conf

    def __repr__(self):
        return f"SparkContext(master='{self.master}', appName='{self.appName}')"


# ---------------------------------------------------------------------------
# UDF registration stub
# ---------------------------------------------------------------------------

class _UDFRegistration:
    def __init__(self, session: SparkSession):
        self._session = session

    def register(self, name: str, f, returnType=None):
        """Register a Python function as a SQL UDF in DuckDB."""
        from singlespark.sql.types import DataType, spark_type_to_polars
        # DuckDB supports Python UDFs; wrap and register
        try:
            self._session.catalog._conn.create_function(name, f)
        except Exception:
            pass  # may already exist
        return f

    def __call__(self, *args, **kwargs):
        from singlespark.sql.functions import udf
        return udf(*args, **kwargs)
