from __future__ import annotations

"""
Generic DataFrame utility functions.
"""

from typing import Any
import pandas as pd
import dask.dataframe as dd
import numpy as np
from sibi_flux.dask_cluster import safe_compute


def safe_where(series: dd.Series, condition: dd.Series, other: Any) -> dd.Series:
    """Ensure condition is boolean before calling .where()."""

    def _coerce_bool(part: pd.Series) -> pd.Series:
        if pd.api.types.is_bool_dtype(part):
            return part.fillna(False)
        if pd.api.types.is_numeric_dtype(part):
            return part.fillna(0).astype(bool)
        return pd.to_numeric(part, errors="coerce").fillna(0).astype(bool)

    cond = condition.map_partitions(_coerce_bool, meta=("cond", "bool"))
    return series.where(cond, other)


def unique_list(
    series: pd.Series | dd.Series, dask_client: Any | None = None
) -> list[Any]:
    """
    Compute a list of unique non-null values from a Series.
    Uses Dask's native .unique() for efficiency and correctness.
    """
    if isinstance(series, pd.Series):
        return pd.unique(series.dropna()).tolist()

    # Dask path: use built-in unique()
    try:
        unique_vals = series.dropna().unique()
        result = safe_compute(unique_vals, dask_client=dask_client)

        if isinstance(result, pd.Series):
            return result.dropna().tolist()
        elif isinstance(result, np.ndarray):
            return result.tolist()
        elif isinstance(result, (list, tuple)):
            return [x for x in result if pd.notna(x)]
        elif pd.notna(result):
            return [result]
        else:
            return []
    except Exception:
        return []
