from __future__ import annotations

import hashlib
from typing import TYPE_CHECKING, Any, Iterable, Final

import pandas as pd
from sibi_flux.logger import Logger

if TYPE_CHECKING:
    import dask.dataframe as dd

# Standardizing on the same UTC definition used conceptually in DateUtils
UTC_DATETIME_DTYPE: Final[str] = "datetime64[ns, UTC]"


class DataUtils:
    """
    Helpers for transforming columns, safe emptiness checks, datetime coercion,
    and joining lookup data for Pandas or Dask DataFrames.
    """

    def __init__(self, logger: Logger | None = None, **kwargs: Any) -> None:
        self.logger = logger or Logger.default_logger(
            logger_name=self.__class__.__name__
        )
        self.debug: bool = bool(kwargs.get("debug", False))

    # ---------- numeric / boolean transforms ----------

    @staticmethod
    def _transform_column_pandas(
        series: pd.Series, fill_value: Any, dtype: type
    ) -> pd.Series:
        """Vectorized pandas transformation."""
        # Use downcast to minimize memory if possible, otherwise standard cast
        return pd.to_numeric(series, errors="coerce").fillna(fill_value).astype(dtype)

    def transform_numeric_columns(
        self,
        df: pd.DataFrame | dd.DataFrame,
        columns: list[str],
        fill_value: Any = 0,
        dtype: type = int,
    ) -> pd.DataFrame | dd.DataFrame:
        """
        Convert selected columns to numeric → fillna → cast dtype.
        Convert selected columns to numeric → fillna → cast dtype.
        Works for Pandas and Dask (partition-wise).
        Uses df.assign() for Dask to avoid graph layer bloat.
        """
        if not columns:
            self.logger.warning("No columns specified for transform_numeric_columns.")
            return df

        # Filter for existing columns to avoid KeyErrors
        cols = [c for c in columns if c in df.columns]
        if not cols:
            self.logger.warning("None of the requested columns exist in the DataFrame.")
            return df

        if isinstance(df, pd.DataFrame):
            # Optimize: Modify in place or use assign for chainability?
            # In-place is more memory efficient for large frames.
            for col in cols:
                df[col] = self._transform_column_pandas(df[col], fill_value, dtype)
            return df

        # Dask path: batch transformations into a single assign()
        import dask.dataframe as dd
        
        transforms = {
            col: df[col].map_partitions(
                self._transform_column_pandas,
                fill_value,
                dtype,
                meta=(col, dtype),
            )
            for col in cols
        }
        return df.assign(**transforms)

    def transform_boolean_columns(
        self,
        df: pd.DataFrame | dd.DataFrame,
        columns: list[str],
        fill_value: Any = 0,
    ) -> pd.DataFrame | dd.DataFrame:
        """Convenience wrapper: cast to boolean via numeric→fillna→astype(bool)."""
        return self.transform_numeric_columns(
            df, columns, fill_value=fill_value, dtype=bool
        )

    # ---------- PII transforms ----------

    @staticmethod
    def _hash_series(series: pd.Series) -> pd.Series:
        """
        Hash a pandas Series using SHA256.
        Strings are encoded to utf-8 before hashing.
        Non-string values are converted to string first.
        """
        def hash_val(val: Any) -> str | None:
            if pd.isna(val) or val == "":
                return None
            val_str = str(val).encode("utf-8")
            return hashlib.sha256(val_str).hexdigest()

        # map is often faster than apply for simple string ops on object columns
        return series.map(hash_val)

    def transform_pii_columns(
        self,
        df: pd.DataFrame | dd.DataFrame,
        columns: list[str],
    ) -> pd.DataFrame | dd.DataFrame:
        """
        Hashes specified columns using SHA256 to obfuscate PII.
        """
        if not columns:
            return df

        cols = [c for c in columns if c in df.columns]
        if not cols:
            self.logger.warning(
                "None of the requested PII columns exist in the DataFrame."
            )
            return df

        if isinstance(df, pd.DataFrame):
            for col in cols:
                df[col] = self._hash_series(df[col])
            return df

        # Dask path
        import dask.dataframe as dd
        transforms = {
            col: df[col].map_partitions(self._hash_series, meta=(col, object))
            for col in cols
        }
        return df.assign(**transforms)

    # ---------- lookup merge ----------

    def merge_lookup_data(
        self,
        classname: Any,
        df: pd.DataFrame | dd.DataFrame,
        *,
        max_unique_ids: int = 50000,
        **kwargs: Any,
    ) -> pd.DataFrame | dd.DataFrame:
        """
        Merge lookup data for ids present in `source_col`.
        Optimized to avoid full lookup table loads by filtering on `source_col` IDs.
        
        :param max_unique_ids: Safeguard against OOM when computing unique IDs.
        """
        if self.is_dataframe_empty(df):
            self.logger.debug(
                "merge_lookup_data: input DataFrame empty — nothing to merge."
            )
            return df

        # 1. Parameter extraction
        try:
            source_col = kwargs.pop("source_col")
            lookup_col = kwargs.pop("lookup_col")
            # These are popped but used in the load_kwargs construction implicitly or explicitly
            lookup_description_col = kwargs.pop("lookup_description_col")
            source_description_alias = kwargs.pop("source_description_alias")

            fillna_alias = bool(kwargs.pop("fillna_source_description_alias", False))
            fieldnames = kwargs.pop("fieldnames", (lookup_col, lookup_description_col))
            column_names = kwargs.pop(
                "column_names", ["temp_join_col", source_description_alias]
            )
        except KeyError as e:
            raise ValueError(f"merge_lookup_data missing required parameter: {e}")

        if source_col not in df.columns:
            self.logger.debug(
                f"merge_lookup_data: '{source_col}' not found — skipping."
            )
            return df

        # 2. Collect unique IDs for filtering (Dimension table optimization)
        import dask.dataframe as dd
        try:
            ids_series = df[source_col].dropna()
            if isinstance(df, dd.DataFrame):
                # Compute unique IDs. Warning: If high cardinality, this is expensive.
                # Assuming lookup tables are typically low cardinality (Categories, SKUs, etc.)
                ids = ids_series.unique().compute()
            else:
                ids = ids_series.unique()

            if len(ids) > max_unique_ids:
                self.logger.warning(
                    f"merge_lookup_data: high cardinality ({len(ids)}) detected for '{source_col}'. "
                    f"Limiting to first {max_unique_ids} for safety."
                )
                ids = ids[:max_unique_ids]

            # Convert to list for DB lookup
            ids_list = sorted(ids.tolist() if not isinstance(ids, list) else ids)
        except Exception as e:
            self.logger.error(
                f"merge_lookup_data: failed extracting ids from '{source_col}': {e}"
            )
            return df

        if not ids_list:
            return df

        # 3. Load only relevant lookup data
        load_kwargs = {
            **kwargs,
            "fieldnames": fieldnames,
            "column_names": column_names,
            f"{lookup_col}__in": ids_list,
        }

        lookup_instance = classname(debug=self.debug, logger=self.logger)
        result = lookup_instance.load(**load_kwargs)

        # Handle Dask result from lookup (materialize it, as it should be small filtered data)
        if isinstance(result, dd.DataFrame):
            result = result.compute()

        if not isinstance(result, pd.DataFrame):
            raise TypeError("merge_lookup_data: lookup 'load' must return a DataFrame.")

        if result.empty:
            return df

        # 4. Merge
        temp_join_col = (
            "temp_join_col" if "temp_join_col" in column_names else lookup_col
        )

        # Dask supports merging a Dask DF (left) with a Pandas DF (right)
        merged = df.merge(
            result, how="left", left_on=source_col, right_on=temp_join_col
        )

        # 5. Cleanup
        if fillna_alias and source_description_alias in merged.columns:
            # Vectorized fillna works on both Dask/Pandas
            merged[source_description_alias] = merged[source_description_alias].fillna(
                ""
            )

        merged = merged.drop(columns="temp_join_col", errors="ignore")
        return merged

    # ---------- emptiness & datetime ----------

    def is_dataframe_empty(self, df: pd.DataFrame | dd.DataFrame) -> bool:
        """
        Safe emptiness check using dask_cluster.
        """
        import dask.dataframe as dd
        if isinstance(df, dd.DataFrame):
            from sibi_flux.dask_cluster import dask_is_empty

            # We don't have easy access to client here, so we rely on global/singleton
            return dask_is_empty(df, logger=self.logger)

        if isinstance(df, pd.DataFrame):
            return df.empty

        self.logger.error(
            "is_dataframe_empty: input must be a pandas or dask DataFrame."
        )
        return False

    @staticmethod
    def convert_to_datetime_dask(
        df: dd.DataFrame, date_fields: Iterable[str]
    ) -> dd.DataFrame:
        """
        Convert specified columns to UTC-aware datetime64[ns, UTC].
        Aligns with DateUtils philosophy: UTC everywhere.
        """

        def _to_pandas_datetime(part: pd.DataFrame) -> pd.DataFrame:
            part = part.copy()
            for col in date_fields:
                if col not in part.columns:
                    continue

                # FIXED: Added format="mixed" to handle columns containing
                # both naive strings and aware (Z-suffix) strings safely.
                part[col] = pd.to_datetime(
                    part[col], errors="coerce", utc=True, format="mixed"
                )
            return part

        meta = df._meta.copy()
        for col in date_fields:
            if col in meta.columns:
                meta[col] = pd.Series([], dtype=UTC_DATETIME_DTYPE)

        return df.map_partitions(_to_pandas_datetime, meta=meta)

    @staticmethod
    def enforce_pandas_dtypes(df: dd.DataFrame) -> dd.DataFrame:
        """
        Convert PyArrow / Dask-specific dtypes to standard Pandas equivalents
        to ensure compatibility with downstream tools.
        """

        def _convert_partition(part: pd.DataFrame) -> pd.DataFrame:
            part = part.copy()
            for col in part.columns:
                dtype_name = str(part[col].dtype)

                # 1. Handle PyArrow Timestamps -> Standard UTC Datetime
                if "timestamp" in dtype_name and "pyarrow" in dtype_name:
                    part[col] = pd.to_datetime(part[col], errors="coerce", utc=True)

                # 2. Handle Strings (Clean up 'None' strings from CSV/Arrow issues)
                elif "string" in dtype_name:
                    # Optimized: Convert to object/string, but avoid the costly replace("None")
                    # unless you explicitly know your data source writes "None" as a string.
                    # If dealing with true PyArrow strings, .astype(object) handles nulls correctly.
                    part[col] = part[col].astype(object)

                    # Only do the slow replace if absolutely necessary:
                    # mask = part[col] == "None"
                    # if mask.any():
                    #     part.loc[mask, col] = None

                # 3. Handle Booleans
                elif dtype_name == "bool":
                    part[col] = part[col].astype("boolean")  # Nullable boolean
            return part

        # Update metadata to match the transformation
        meta = df._meta.copy()
        for col in meta.columns:
            dtype_name = str(meta[col].dtype)
            if "timestamp" in dtype_name and "pyarrow" in dtype_name:
                meta[col] = pd.Series([], dtype=UTC_DATETIME_DTYPE)
            elif "string" in dtype_name:
                meta[col] = pd.Series([], dtype="object")
            elif dtype_name == "bool":
                meta[col] = pd.Series([], dtype="boolean")

        return df.map_partitions(_convert_partition, meta=meta)

    # ---------- filtering ----------

    def exclude_from_dataframe(
        self,
        df: pd.DataFrame | dd.DataFrame,
        conditions: list[tuple[str, str, Any]],
    ) -> pd.DataFrame | dd.DataFrame:
        """
        Exclude rows based on a list of conditions (col, op, val).
        Supported operators: '==', '!=', '>', '<', '>=', '<=', 'in', 'not in'.
        """
        if self.is_dataframe_empty(df):
            return df

        import operator

        ops: Final[dict[str, Any]] = {
            "==": operator.eq,
            "!=": operator.ne,
            ">": operator.gt,
            "<": operator.lt,
            ">=": operator.ge,
            "<=": operator.le,
            # "in" and "not in" need special handling if not using df.query/eval
        }

        for col, op_str, val in conditions:
            if col not in df.columns:
                self.logger.warning(
                    f"exclude_from_dataframe: column '{col}' not found. Skipping."
                )
                continue

            # Invert logic: we want to EXCLUDE rows where condition is TRUE.
            # So we KEEP rows where condition is FALSE.
            # query logic: mask = (col op val) -> df = df[~mask]

            # Simple operators
            if op_str in ops:
                op_func = ops[op_str]
                mask = op_func(df[col], val)
                df = df[~mask]

            elif op_str == "in":
                mask = df[col].isin(val)
                df = df[~mask]

            elif op_str == "not in":
                mask = ~df[col].isin(val)
                df = df[~mask]

            else:
                self.logger.warning(
                    f"exclude_from_dataframe: unsupported operator '{op_str}'. Skipping."
                )

        return df
