"""
DataFrameReader — mirrors pyspark.sql.DataFrameReader.
Reads from local paths and fake HDFS paths.
"""
from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

import polars as pl

if TYPE_CHECKING:
    from singlespark.session import SparkSession
    from singlespark.dataframe import DataFrame
    from singlespark.sql.types import StructType


class DataFrameReader:
    def __init__(self, session: "SparkSession"):
        self._session = session
        self._fmt: str | None = None
        self._options: dict = {}
        self._schema_hint: "StructType | None" = None

    # ------------------------------------------------------------------
    # Builder methods
    # ------------------------------------------------------------------

    def format(self, fmt: str) -> "DataFrameReader":
        self._fmt = fmt.lower()
        return self

    def option(self, key: str, value) -> "DataFrameReader":
        self._options[key] = value
        return self

    def options(self, **kwargs) -> "DataFrameReader":
        self._options.update(kwargs)
        return self

    def schema(self, schema) -> "DataFrameReader":
        self._schema_hint = schema
        return self

    # ------------------------------------------------------------------
    # load() — dispatches to format-specific reader
    # ------------------------------------------------------------------

    def load(self, path: str | list[str] | None = None, format: str | None = None,
             schema=None, **options) -> "DataFrame":
        fmt = format or self._fmt or "parquet"
        if schema:
            self._schema_hint = schema
        self._options.update(options)
        reader = getattr(self, fmt, None)
        if reader is None:
            raise ValueError(f"Unsupported format: {fmt!r}")
        if path is not None:
            return reader(path)
        raise ValueError("path is required for load()")

    # ------------------------------------------------------------------
    # Format-specific readers
    # ------------------------------------------------------------------

    def parquet(self, *paths: str | list[str]) -> "DataFrame":
        from singlespark.dataframe import DataFrame
        flat_paths = _normalise_paths(paths)
        resolved = self._session._hdfs.expand_paths(flat_paths)
        if not resolved:
            raise FileNotFoundError(f"No files found for paths: {flat_paths}")
        dfs = [pl.read_parquet(p) for p in resolved]
        result = pl.concat(dfs, how="diagonal") if len(dfs) > 1 else dfs[0]
        result = self._apply_schema(result)
        return DataFrame(result, self._session)

    def csv(self, path: str | list[str], **kwargs) -> "DataFrame":
        from singlespark.dataframe import DataFrame
        paths = _normalise_paths((path,))
        resolved = self._session._hdfs.expand_paths(paths)

        # Merge caller kwargs + stored options
        opts = {**self._options, **kwargs}
        has_header = opts.pop("header", opts.pop("has_header", True))
        if isinstance(has_header, str):
            has_header = has_header.lower() == "true"
        sep = opts.pop("sep", opts.pop("delimiter", opts.pop("separator", ",")))
        infer_schema = opts.pop("inferSchema", opts.pop("infer_schema", True))
        if isinstance(infer_schema, str):
            infer_schema = infer_schema.lower() == "true"

        polars_opts = {
            "has_header": has_header,
            "separator": sep,
            "infer_schema": infer_schema,
        }

        dfs = [pl.read_csv(p, **polars_opts) for p in resolved]
        result = pl.concat(dfs, how="diagonal") if len(dfs) > 1 else dfs[0]
        result = self._apply_schema(result)
        return DataFrame(result, self._session)

    def json(self, path: str | list[str], **kwargs) -> "DataFrame":
        from singlespark.dataframe import DataFrame
        paths = _normalise_paths((path,))
        resolved = self._session._hdfs.expand_paths(paths)
        dfs = [pl.read_ndjson(p) for p in resolved]
        result = pl.concat(dfs, how="diagonal") if len(dfs) > 1 else dfs[0]
        result = self._apply_schema(result)
        return DataFrame(result, self._session)

    def text(self, path: str | list[str]) -> "DataFrame":
        from singlespark.dataframe import DataFrame
        paths = _normalise_paths((path,))
        resolved = self._session._hdfs.expand_paths(paths)
        lines: list[str] = []
        for p in resolved:
            lines.extend(Path(p).read_text(encoding="utf-8").splitlines())
        return DataFrame(pl.DataFrame({"value": lines}), self._session)

    def orc(self, path: str | list[str]) -> "DataFrame":
        raise NotImplementedError("ORC format is not supported in singlespark (no Polars ORC reader)")

    def table(self, name: str) -> "DataFrame":
        return self._session.sql(f"SELECT * FROM {name}")

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _apply_schema(self, df: pl.DataFrame) -> pl.DataFrame:
        """Cast columns to schema types if a schema hint was provided."""
        if self._schema_hint is None:
            return df
        from singlespark.sql.types import StructType, spark_type_to_polars
        if not isinstance(self._schema_hint, StructType):
            return df
        casts = []
        for field in self._schema_hint.fields:
            if field.name in df.columns:
                casts.append(pl.col(field.name).cast(spark_type_to_polars(field.dataType)))
        return df.with_columns(casts) if casts else df


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _normalise_paths(paths) -> list[str]:
    """Flatten nested path args to a list of strings."""
    result: list[str] = []
    for p in paths:
        if isinstance(p, (list, tuple)):
            result.extend([str(x) for x in p])
        elif p is not None:
            result.append(str(p))
    return result
