from __future__ import annotations

import os
import warnings
from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
from functools import partial
from multiprocessing.pool import ThreadPool
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
)

import dask.config
import dask.dataframe as dd
import pandas as pd
import pyarrow as pa

from sibi_flux.core import SecureResource
from sibi_flux.dask_cluster.core import safe_persist
from ._write_gatekeeper import get_write_sem

try:
    from tqdm.dask import TqdmCallback
except ImportError:
    TqdmCallback = None

if TYPE_CHECKING:
    from dask.distributed import Client

warnings.filterwarnings(
    "ignore", message="Passing 'overwrite=True' to to_parquet is deprecated"
)


def _coerce_partition(
    pdf: pd.DataFrame,
    target: dict[str, pa.DataType],
    partition_cols: list[str] | None = None,
) -> pd.DataFrame:
    """
    Applies type conversions to a single pandas partition.
    """
    partition_cols = partition_cols or []

    for col, pa_type in target.items():
        if col not in pdf.columns:
            if pa.types.is_string(pa_type):
                pdf[col] = None
            elif pa.types.is_integer(pa_type):
                pdf[col] = pd.NA
            elif pa.types.is_floating(pa_type):
                pdf[col] = float("nan")
            elif pa.types.is_boolean(pa_type):
                pdf[col] = pd.NA
            elif pa.types.is_timestamp(pa_type):
                pdf[col] = pd.NaT
            else:
                pdf[col] = None

        series = pdf[col]
        current_dtype = series.dtype

        try:
            # --- 0. Partition Column Logic ---
            if col in partition_cols:
                if not pd.api.types.is_numeric_dtype(current_dtype):
                    try:
                        temp_series = pd.to_datetime(
                            series, errors="raise", utc=True, format="mixed"
                        )
                        pdf[col] = temp_series.dt.strftime("%Y-%m-%d").astype(
                            "string[pyarrow]"
                        )
                        continue
                    except (ValueError, TypeError):
                        pass

                if str(current_dtype) != "string[pyarrow]":
                    pdf[col] = series.astype("string[pyarrow]")
                continue

            # --- 1. Timestamp Coercion ---
            if pa.types.is_timestamp(pa_type):
                if not pd.api.types.is_datetime64_any_dtype(current_dtype):
                    series = pd.to_datetime(
                        series, errors="coerce", utc=True, format="mixed"
                    )

                if series.dt.tz is None:
                    series = series.dt.tz_localize("UTC")
                elif str(series.dt.tz) != "UTC":
                    series = series.dt.tz_convert("UTC")

                if str(series.dtype) != "timestamp[ns, tz=UTC][pyarrow]":
                    pdf[col] = series.astype("timestamp[ns, tz=UTC][pyarrow]")
                else:
                    pdf[col] = series
                continue

            # --- 2. Other Types ---
            current_dtype_str = str(current_dtype)

            if pa.types.is_string(pa_type):
                if "string" not in current_dtype_str:
                    pdf[col] = series.astype("string[pyarrow]")

            elif pa.types.is_boolean(pa_type):
                if "bool" not in current_dtype_str:
                    if pd.api.types.is_object_dtype(current_dtype) or pd.api.types.is_string_dtype(current_dtype):
                        series = series.astype(str).str.lower().isin(["true", "1", "yes"])
                    pdf[col] = series.astype("boolean[pyarrow]")

            elif pa.types.is_integer(pa_type):
                if "int" not in current_dtype_str:
                    pdf[col] = pd.to_numeric(series, errors="coerce").astype("int64[pyarrow]")

            elif pa.types.is_floating(pa_type):
                if "float" not in current_dtype_str and "double" not in current_dtype_str:
                    pdf[col] = pd.to_numeric(series, errors="coerce").astype("float64[pyarrow]")

        except (ValueError, TypeError):
            # Known conversion issues, skip and let pyarrow handle or fail at write time
            pass
        except Exception:
            # Catch-all for unexpected issues in partition processing
            pass

    return pdf


class ParquetSaver(SecureResource):
    """
    Production-grade Dask -> Parquet writer.
    """

    logger_extra: ClassVar[dict[str, str]] = {"sibi_flux_component": __name__}

    def __init__(
        self,
        df_result: dd.DataFrame,
        parquet_storage_path: str,
        *,
        repartition_size: str | None = "128MB",
        persist: bool = False,
        write_index: bool = False,
        write_metadata_file: bool = True,
        pyarrow_args: dict[str, Any] | None = None,
        writer_threads: int | str = "auto",
        arrow_cpu: int | None = None,
        partitions_per_round: int = 24,
        max_delete_workers: int = 8,
        write_gate_max: int = 2,
        write_gate_key: str | None = None,
        partition_on: list[str] | None = None,
        dask_client: Client | None = None,
        schema: pa.Schema | None = None,
        show_progress: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)

        if not isinstance(df_result, dd.DataFrame):
            raise TypeError("df_result must be a Dask DataFrame")
        if not self.fs:
            raise ValueError("File system (fs) must be provided to ParquetSaver.")

        self.df_result = df_result
        self.dask_client = dask_client
        self.parquet_storage_path = parquet_storage_path.rstrip("/")
        self.repartition_size = repartition_size
        self.persist = persist
        self.write_index = write_index
        self.write_metadata_file = write_metadata_file
        self.pyarrow_args = dict(pyarrow_args or {})

        if writer_threads == "auto" or writer_threads is None:
            self.writer_threads = self._calculate_optimal_threads()
            self.logger.debug(f"Auto-tuned writer_threads to: {self.writer_threads}")
        else:
            self.writer_threads = max(1, int(writer_threads))

        self.arrow_cpu = None if arrow_cpu is None else max(1, int(arrow_cpu))
        self.partitions_per_round = max(1, int(partitions_per_round))
        self.max_delete_workers = max(1, int(max_delete_workers))
        self.write_gate_max = max(1, int(write_gate_max))
        self.write_gate_key = (write_gate_key or self.parquet_storage_path).rstrip("/")
        self.partition_on = partition_on or []
        self.scheduler = kwargs.get("scheduler", None)
        self.force_local = kwargs.get("force_local", False)
        self.schema = schema
        self.show_progress = show_progress

        self.pyarrow_args.setdefault("compression", "zstd")

        self.protocol = "file"
        if "://" in self.parquet_storage_path:
            self.protocol = self.parquet_storage_path.split(":", 1)[0]

    def save_to_parquet(
        self, output_directory_name: str = "default_output", overwrite: bool = True
    ) -> str:
        # 1. Determine Execution Mode
        assert self.fs is not None, "FileSystem (fs) must be available"
        using_distributed = False
        if self.force_local:
            using_distributed = False
        elif self.dask_client is not None:
            using_distributed = True
        else:
            try:
                from dask.distributed import get_client
                get_client()
                using_distributed = True
            except (ImportError, ValueError, RuntimeError):
                pass

        # 2. Configure Context Factories
        if using_distributed:
            def get_config_context():
                return nullcontext()
            def get_pool_context():
                return nullcontext()
        else:
            def get_config_context():
                return dask.config.set({"dataframe.shuffle.method": "tasks"})

            def get_pool_context():
                if self.scheduler == "synchronous":
                    return dask.config.set(scheduler="synchronous")
                elif self.scheduler == "threads":
                    return dask.config.set(
                        pool=ThreadPool(self.writer_threads), scheduler="threads"
                    )
                elif not using_distributed and self.scheduler is None:
                    return dask.config.set(scheduler="single-threaded")
                return nullcontext()

        # 3. Path Resolution
        if self.partition_on:
            overwrite = False
            target_path = self.parquet_storage_path.rstrip("/")
        else:
            target_path = f"{self.parquet_storage_path}/{output_directory_name}".rstrip("/")

        sem = get_write_sem(self.write_gate_key, self.write_gate_max)

        # 4. Execution
        with get_config_context():
            with sem:
                if overwrite and self.fs.exists(target_path):
                    self._clear_directory_safely(target_path)
                self.fs.mkdirs(target_path, exist_ok=True)

                schema = self._define_schema()
                ddf = self._coerce_ddf_to_schema(self.df_result, schema)

                if self.repartition_size:
                    ddf = ddf.repartition(partition_size=self.repartition_size)

                if self.persist:
                    with get_pool_context():
                        if using_distributed:
                            ddf = safe_persist(ddf)
                        else:
                            ddf = ddf.persist()

                old_arrow_cpu = None
                if self.arrow_cpu:
                    old_arrow_cpu = pa.get_cpu_count()
                    pa.set_cpu_count(self.arrow_cpu)

                try:
                    params = {
                        "path": target_path,
                        "engine": "pyarrow",
                        "filesystem": self.fs,
                        "write_index": self.write_index,
                        "write_metadata_file": self.write_metadata_file,
                        "schema": schema,
                        **self.pyarrow_args,
                    }

                    if self.partition_on:
                        params["partition_on"] = self.partition_on

                    write_task = ddf.to_parquet(**params, compute=False)

                    compute_kwargs = {}
                    if not using_distributed:
                        if self.scheduler in ("synchronous", "threads", "single-threaded"):
                            compute_kwargs["scheduler"] = self.scheduler
                        elif self.scheduler is None:
                            compute_kwargs["scheduler"] = "single-threaded"

                    with get_pool_context():
                        ctx = TqdmCallback(desc=f"Saving Parquet to {output_directory_name}") if self.show_progress and TqdmCallback else nullcontext()
                        with ctx:
                            dask.compute(write_task, **compute_kwargs)
                finally:
                    if old_arrow_cpu is not None:
                        pa.set_cpu_count(old_arrow_cpu)

        self.logger.info(
            f"Parquet dataset written: {target_path}", extra=self.logger_extra
        )
        return target_path

    # ---------- Internals ----------

    def _calculate_optimal_threads(self) -> int:
        try:
            if "OMP_NUM_THREADS" in os.environ:
                available_cores = int(os.environ["OMP_NUM_THREADS"])
            elif hasattr(os, "sched_getaffinity"):
                available_cores = len(os.sched_getaffinity(0))
            else:
                available_cores = os.cpu_count() or 4
            recommended = available_cores * 3
            return max(4, min(recommended, 32))
        except Exception:
            return 8

    def _clear_directory_safely(self, directory: str) -> None:
        if self.protocol.startswith("s3"):
            entries = [p for p in self.fs.glob(f"{directory}/**") if p != directory]
            if not entries:
                return

            def _rm_one(p: str) -> None:
                try:
                    self.fs.rm_file(p)
                except Exception as e:
                    self.logger.warning(
                        f"Delete failed '{p}': {e}", extra=self.logger_extra
                    )

            with ThreadPoolExecutor(max_workers=self.max_delete_workers) as ex:
                list(ex.map(_rm_one, entries))
            try:
                self.fs.rm(directory, recursive=False)
            except Exception:
                pass
        else:
            self.fs.rm(directory, recursive=True)

    def _define_schema(self) -> pa.Schema:
        fields = []
        # Pre-calculate overrides from self.schema if provided
        overrides = {f.name: f.type for f in self.schema} if self.schema else {}

        for name, dtype in self.df_result.dtypes.items():
            if name in overrides:
                fields.append(pa.field(name, overrides[name]))
                continue

            if name in self.partition_on:
                fields.append(pa.field(name, pa.string()))
                continue

            dtype_str = str(dtype).lower()
            lower_name = name.lower()

            # --- PRIORITY 1: Heuristics ---
            if any(x in lower_name for x in ["_dt", "_date", "_at", "time"]):
                pa_type = pa.timestamp("ns", tz="UTC")
            elif lower_name.startswith(("is_", "has_", "flag_")):
                pa_type = pa.bool_()
            elif (
                any(x in lower_name for x in ["_id", "id"]) and "guid" not in lower_name
            ):
                pa_type = pa.int64()
            elif any(
                x in lower_name for x in ["amount", "price", "score", "percent", "rate"]
            ):
                pa_type = pa.float64()

            # --- PRIORITY 2: Dtypes ---
            elif "int" in dtype_str:
                pa_type = pa.int64()
            elif "float" in dtype_str or "double" in dtype_str:
                pa_type = pa.float64()
            elif "bool" in dtype_str:
                pa_type = pa.bool_()
            elif "datetime" in dtype_str or "timestamp" in dtype_str:
                pa_type = pa.timestamp("ns", tz="UTC")
            else:
                pa_type = pa.string()

            fields.append(pa.field(name, pa_type))

        return pa.schema(fields)

    def _coerce_ddf_to_schema(
        self, ddf: dd.DataFrame, schema: pa.Schema
    ) -> dd.DataFrame:
        target = {f.name: f.type for f in schema}

        meta_cols: dict[str, pd.Series] = {}
        for name, typ in target.items():
            if name in self.partition_on:
                meta_cols[name] = pd.Series([], dtype="string[pyarrow]")
            elif pa.types.is_timestamp(typ):
                meta_cols[name] = pd.Series([], dtype="timestamp[ns, tz=UTC][pyarrow]")
            elif pa.types.is_integer(typ):
                meta_cols[name] = pd.Series([], dtype="int64[pyarrow]")
            elif pa.types.is_floating(typ):
                meta_cols[name] = pd.Series([], dtype="float64[pyarrow]")
            elif pa.types.is_boolean(typ):
                meta_cols[name] = pd.Series([], dtype="boolean[pyarrow]")
            else:
                meta_cols[name] = pd.Series([], dtype="string[pyarrow]")

        new_meta = pd.DataFrame(meta_cols, index=ddf._meta.index)
        new_meta = new_meta[list(target.keys())]

        coerce_fn = partial(
            _coerce_partition, target=target, partition_cols=self.partition_on
        )
        return ddf.map_partitions(coerce_fn, meta=new_meta)
