"""
DataFrameWriter — mirrors pyspark.sql.DataFrameWriter.
Writes to local paths and fake HDFS paths.
"""
from __future__ import annotations

import shutil
from pathlib import Path
from typing import TYPE_CHECKING

import polars as pl

if TYPE_CHECKING:
    from singlespark.dataframe import DataFrame


class DataFrameWriter:
    """
    Mirrors pyspark.sql.DataFrameWriter.
    Chain option/mode/format calls then call a terminal write method.
    """

    def __init__(self, df: "DataFrame"):
        self._df = df
        self._fmt: str = "parquet"
        self._mode: str = "error"
        self._options: dict = {}
        self._partition_cols: list[str] = []
        self._sort_cols: list[str] = []

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

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

    def mode(self, saveMode: str) -> "DataFrameWriter":
        self._mode = saveMode.lower()
        return self

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

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

    def partitionBy(self, *cols) -> "DataFrameWriter":
        flat = list(cols[0]) if len(cols) == 1 and isinstance(cols[0], (list, tuple)) else list(cols)
        self._partition_cols = flat
        return self

    def sortBy(self, *cols) -> "DataFrameWriter":
        flat = list(cols[0]) if len(cols) == 1 and isinstance(cols[0], (list, tuple)) else list(cols)
        self._sort_cols = flat
        return self

    bucketBy = sortBy  # stub

    # ------------------------------------------------------------------
    # Terminal write methods
    # ------------------------------------------------------------------

    def save(self, path: str | None = None) -> None:
        writer = getattr(self, self._fmt, None)
        if writer is None:
            raise ValueError(f"Unsupported format: {self._fmt!r}")
        if path is not None:
            writer(path)
        else:
            raise ValueError("path is required for save()")

    def parquet(self, path: str) -> None:
        resolved = self._resolve(path)
        if not self._check_mode(resolved):
            return
        resolved.mkdir(parents=True, exist_ok=True)
        if self._partition_cols:
            self._write_partitioned(resolved, "parquet")
        else:
            self._df._df.write_parquet(str(resolved / "part-00000.snappy.parquet"))

    def csv(self, path: str) -> None:
        resolved = self._resolve(path)
        if not self._check_mode(resolved):
            return
        resolved.mkdir(parents=True, exist_ok=True)
        sep = self._options.get("sep", self._options.get("delimiter", ","))
        header = self._options.get("header", True)
        if isinstance(header, str):
            header = header.lower() == "true"
        self._df._df.write_csv(str(resolved / "part-00000.csv"), separator=sep, include_header=header)

    def json(self, path: str) -> None:
        resolved = self._resolve(path)
        if not self._check_mode(resolved):
            return
        resolved.mkdir(parents=True, exist_ok=True)
        self._df._df.write_ndjson(str(resolved / "part-00000.json"))

    def text(self, path: str) -> None:
        """Write a single-column DataFrame as a plain text file."""
        resolved = self._resolve(path)
        if not self._check_mode(resolved):
            return
        resolved.mkdir(parents=True, exist_ok=True)
        col = self._df._df.columns[0]
        lines = "\n".join(str(v) for v in self._df._df[col].to_list())
        (resolved / "part-00000.txt").write_text(lines, encoding="utf-8")

    def saveAsTable(self, name: str) -> None:
        """Save as a named temp view (in-memory table)."""
        self._df.createOrReplaceTempView(name)

    # ------------------------------------------------------------------
    # Partitioned write helper
    # ------------------------------------------------------------------

    def _write_partitioned(self, base: Path, fmt: str) -> None:
        """Write partitioned output (Hive-style: col=value/part-XXXXX.parquet)."""
        partition_cols = self._partition_cols
        groups = self._df._df.partition_by(partition_cols, as_dict=True)
        for key_tuple, part_df in groups.items():
            # Build Hive partition path
            if not isinstance(key_tuple, tuple):
                key_tuple = (key_tuple,)
            subdir = Path(*[f"{col}={val}" for col, val in zip(partition_cols, key_tuple)])
            out_dir = base / subdir
            out_dir.mkdir(parents=True, exist_ok=True)
            # Drop partition columns from data (Spark behaviour)
            data_df = part_df.drop(partition_cols)
            if fmt == "parquet":
                data_df.write_parquet(str(out_dir / "part-00000.snappy.parquet"))
            elif fmt == "csv":
                data_df.write_csv(str(out_dir / "part-00000.csv"))
            elif fmt == "json":
                data_df.write_ndjson(str(out_dir / "part-00000.json"))

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _resolve(self, path: str) -> Path:
        return Path(str(self._df._session._hdfs.resolve(path)))

    def _check_mode(self, resolved: Path) -> bool:
        """
        Handle save mode. Returns True if writing should proceed.
        """
        if not resolved.exists():
            return True

        mode = self._mode
        if mode in ("error", "errorifexists", "default"):
            raise FileExistsError(
                f"Path already exists: {resolved}\n"
                "Use .mode('overwrite') or .mode('append') to handle existing data."
            )
        if mode == "ignore":
            return False  # skip write silently
        if mode == "overwrite":
            if resolved.is_dir():
                shutil.rmtree(resolved)
            else:
                resolved.unlink()
            return True
        if mode == "append":
            return True  # will write new part file alongside existing

        raise ValueError(f"Unknown save mode: {mode!r}")
