"""
Core functionality for resilient Dask operations.
"""

import asyncio
from typing import Any, Dict, List, Optional, Tuple

import dask
import dask.dataframe as dd
import numpy as np
import pandas as pd
from dask.distributed import Client, Future

from .utils import _to_int_safe


# ---------------------------------------------------------------------------
# Safe compute / persist / gather with auto rebind and local fallback
# ---------------------------------------------------------------------------
def _rebind_client(logger=None) -> Optional[Client]:
    from .client_manager import get_persistent_client  # Avoid circular import

    try:
        return get_persistent_client(logger=logger)
    except Exception:
        return None


def _retry_with_rebind(op, *args, dask_client: Optional[Client], logger=None, **kwargs):
    from .client_manager import (
        RECOVERABLE_COMMS,
    )  # Import here to avoid circular import

    try:
        return op(*args, dask_client=dask_client, logger=logger, **kwargs)
    except RECOVERABLE_COMMS:
        if logger:
            logger.warning("Dask comm closed. Rebinding client and retrying once.")
        c2 = _rebind_client(logger)
        if c2:
            return op(*args, dask_client=c2, logger=logger, **kwargs)
        # Last resort: local compute path if applicable
        obj = args[0] if args else None
        if hasattr(obj, "compute"):
            return obj.compute(scheduler="threads")
        raise


def _compute_impl(obj: Any, dask_client: Optional[Client], **_) -> Any:
    if dask_client:
        res = dask_client.compute(obj)
        return res.result() if isinstance(res, Future) else res
    return obj.compute()


def _persist_impl(obj: Any, dask_client: Optional[Client], **_) -> Any:
    if dask_client:
        res = dask_client.persist(obj)
        return res.result() if isinstance(res, Future) else res
    return obj.persist()


def _gather_impl(objs: List[Any], dask_client: Optional[Client], **_) -> List[Any]:
    if dask_client:
        futs = [dask_client.compute(o) for o in objs]
        return dask_client.gather(futs)
    return list(dask.compute(*objs, scheduler="threads"))


def _safe_compute(obj: Any, dask_client: Optional[Client] = None, logger=None) -> Any:
    return _retry_with_rebind(
        _compute_impl, obj, dask_client=dask_client, logger=logger
    )


def _safe_persist(obj: Any, dask_client: Optional[Client] = None, logger=None) -> Any:
    return _retry_with_rebind(
        _persist_impl, obj, dask_client=dask_client, logger=logger
    )


def _safe_gather(
    objs: List[Any], dask_client: Optional[Client] = None, logger=None
) -> List[Any]:
    if not objs:
        return []

    from .client_manager import (
        RECOVERABLE_COMMS,
        _default_logger,
    )  # Import here to avoid circular import

    LOG = logger or _default_logger

    try:
        return _gather_impl(objs, dask_client)
    except RECOVERABLE_COMMS:
        LOG.warning("Dask comm closed during gather. Rebinding.")
        c2 = _rebind_client(LOG)
        if c2:
            return _gather_impl(objs, c2)
        return list(dask.compute(*objs, scheduler="threads"))
    except ValueError as e:
        if "Missing dependency" in str(e):
            LOG.warning("Detected orphaned Dask graph. Recomputing locally.")
            try:
                return [o.compute(scheduler="threads") for o in objs]
            except Exception as inner:
                LOG.error(f"Local recompute failed: {inner}")
                raise
        raise


def _safe_wait(
    obj: Any,
    dask_client: Optional[Client] = None,
    timeout: Optional[float] = None,
    logger=None,
) -> Any:
    from .client_manager import _default_logger  # Import here to avoid circular import

    LOG = logger or _default_logger

    if obj is None:
        return None
    try:
        if dask_client:
            from dask.distributed import get_client as get_dask_client
            from contextlib import suppress

            with suppress(Exception):
                dask_client.wait_for_workers(1, timeout=10)
            dask_client.wait(obj, timeout=timeout)
            return obj
        try:
            c = get_dask_client()
            c.wait(obj, timeout=timeout)
            return obj
        except ValueError:
            if hasattr(obj, "compute"):
                obj.compute(scheduler="threads")
            return obj
    except Exception as e:
        LOG.warning(f"_safe_wait: {type(e).__name__}: {e}")
        return obj


# ---------------------------------------------------------------------------
# Heuristic emptiness checks
# ---------------------------------------------------------------------------
def dask_is_probably_empty(ddf: dd.DataFrame) -> bool:
    """Check if a Dask DataFrame is probably empty based on metadata."""
    return getattr(ddf, "npartitions", 0) == 0 or len(ddf._meta.columns) == 0


def dask_is_empty_truthful(
    ddf: dd.DataFrame, dask_client: Optional[Client] = None, logger=None
) -> bool:
    """Check if a Dask DataFrame is actually empty by computing the total length."""
    total = _safe_compute(ddf.map_partitions(len).sum(), dask_client, logger=logger)
    return int(_to_int_safe(total)) == 0


def dask_is_empty(
    ddf: dd.DataFrame,
    *,
    sample: int = 4,
    dask_client: Optional[Client] = None,
    logger=None,
) -> bool:
    """
    Check if a Dask DataFrame is empty using sampling and fallback methods.

    Args:
        ddf: The Dask DataFrame to check
        sample: Number of partitions to sample initially
        dask_client: Optional Dask client to use
        logger: Optional logger instance

    Returns:
        True if the DataFrame is empty, False otherwise
    """
    if dask_is_probably_empty(ddf):
        return True
    k = min(max(sample, 1), ddf.npartitions)
    probes = _safe_gather(
        [ddf.get_partition(i).map_partitions(len) for i in range(k)],
        dask_client,
        logger=logger,
    )
    if any(_to_int_safe(n) > 0 for n in probes):
        return False
    if k == ddf.npartitions and all(_to_int_safe(n) == 0 for n in probes):
        return True
    return dask_is_empty_truthful(ddf, dask_client=dask_client, logger=logger)


# ---------------------------------------------------------------------------
# Unique value extractor
# ---------------------------------------------------------------------------
class UniqueValuesExtractor:
    """Extract unique values from Dask DataFrame columns with resilience."""

    def __init__(self, dask_client: Optional[Client] = None, logger=None):
        self.dask_client = dask_client
        self.logger = logger

    def _compute_to_list_sync(self, series) -> List[Any]:
        if hasattr(series, "compute"):
            if self.dask_client:
                result = self.dask_client.compute(series).result()
            else:
                result = series.compute()
        else:
            result = series
        if isinstance(result, (np.ndarray, pd.Series, list)):
            return pd.Series(result).dropna().unique().tolist()
        return [result]

    async def compute_to_list(self, series) -> List[Any]:
        return await asyncio.to_thread(self._compute_to_list_sync, series)

    async def extract_unique_values(self, df, *columns: str) -> Dict[str, List[Any]]:
        """Extract unique values from specified columns."""

        async def one(col: str):
            ser = df[col].dropna().unique()
            return col, await self.compute_to_list(ser)

        kv = await asyncio.gather(*(one(c) for c in columns))
        return dict(kv)
