"""
Core functionality for resilient Dask operations.
- Dry Run: Graph complexity inspection and logging to OpenObserve.
- Resilience: Auto-healing via persistent client registry.
- Invariants: Strict prohibition of local fallback for Dask DataFrames.
"""

from __future__ import annotations

import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
    TypeVar,
)

import dask
import dask.dataframe as dd
import pandas as pd

try:
    from dask.distributed import Client, Future
    from dask.distributed import wait as dask_wait

    HAS_DISTRIBUTED = True
except ImportError:
    Client = object
    Future = object

    def dask_wait(*args: Any, **kwargs: Any) -> None:
        pass

    HAS_DISTRIBUTED = False

if TYPE_CHECKING:
    from sibi_flux.logger import Logger as SibiLogger


# Project-specific imports
from .client_manager import get_persistent_client
from .exceptions import RECOVERABLE_COMMS
from .utils import _to_int_safe

T = TypeVar("T")

# ---------------------------------------------------------------------------
# Late-Binding & Helpers
# ---------------------------------------------------------------------------


def _get_log() -> SibiLogger | logging.Logger:
    """Late-binds the Logger to prevent circular imports during init."""
    try:
        from sibi_flux.logger import Logger

        return Logger.default_logger(logger_name="dask_cluster.core")
    except ImportError:
        return logging.getLogger("dask_cluster.core")


def _is_dask_dataframe_like(obj: Any) -> bool:
    """Checks if object is a Dask collection relying on distributed state."""
    return isinstance(obj, (dd.DataFrame, dd.Series)) or hasattr(obj, "_meta")


def _get_active_client(
    provided_client: Optional[Client], logger: SibiLogger | logging.Logger | None = None
) -> Optional[Client]:
    """Retrieves a healthy client, healing the persistent one if necessary."""
    if provided_client and getattr(provided_client, "status", "") == "running":
        return provided_client
    try:
        # get_persistent_client handles internal healing/watchdog logic
        return get_persistent_client(logger=logger)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Inspection & Dry Run
# ---------------------------------------------------------------------------


def get_graph_metrics(obj: Any) -> dict[str, Any]:
    """Extract complexity metrics from a Dask object for observability."""
    try:
        # If it's a list (e.g., for safe_gather), check if any item is dask-backed
        if isinstance(obj, list):
            obj = obj[0] if obj and hasattr(obj[0], "__dask_graph__") else None

        if obj is None or not hasattr(obj, "__dask_graph__"):
            return {"is_dask": False}

        graph = obj.__dask_graph__()
        return {
            "type": type(obj).__name__,
            "is_dask": True,
            "task_count": len(graph),
            "n_partitions": getattr(obj, "npartitions", "N/A"),
            "layers": (
                len(getattr(graph, "layers", []))
                if hasattr(graph, "layers")
                else "unknown"
            ),
        }
    except Exception as e:
        return {"error_extracting_metrics": str(e)}


# ---------------------------------------------------------------------------
# Resilient Execution Engine
# ---------------------------------------------------------------------------


def _execute_with_resilience(
    op: Callable[..., T],
    obj: Any,
    dask_client: Optional[Client],
    logger: SibiLogger | logging.Logger | None = None,
    dry_run: bool = False,
    **kwargs: Any,
) -> Optional[T]:
    """
    Orchestrates Dask operations with dry-run logging and a single-retry
    healing mechanism for communication failures.
    """
    log = logger or _get_log()

    # 1. Observability: Log graph complexity to OpenObserve
    metrics = get_graph_metrics(obj)
    if metrics.get("is_dask"):
        log.info(
            "Dask Graph Inspection",
            extra={"graph_metrics": metrics, "dry_run": dry_run},
        )
        if dry_run:
            log.info("Dry Run: Execution skipped.")
            return None

    # 2. Execution with Auto-Healing
    active_client = _get_active_client(dask_client, logger=log)
    try:
        return op(obj, active_client, **kwargs)
    except RECOVERABLE_COMMS as e:
        log.warning(f"Dask comm failure ({type(e).__name__}). Healing and retrying.")

        # Trigger explicit heal via singleton refresh
        active_client = get_persistent_client(logger=log)

        if active_client:
            log.info("Client healed. Resubmitting task.")
            return op(obj, active_client, **kwargs)

        # Guard: Never fall back to local compute for DataFrames (Memory Safety)
        if _is_dask_dataframe_like(obj):
            raise RuntimeError(
                "Distributed client lost and cannot be healed. "
                "Local fallback forbidden for DataFrames."
            ) from e

        log.warning("Falling back to local threaded compute (safe for non-DataFrame).")
        return obj.compute(scheduler="threads")


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


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


def safe_compute(
    obj: Any,
    dask_client: Optional[Client] = None,
    logger: SibiLogger | logging.Logger | None = None,
    dry_run: bool = False,
) -> Any:
    """Compute with auto-healing and optional dry-run complexity logging."""
    return _execute_with_resilience(
        _compute_impl, obj, dask_client, logger, dry_run=dry_run
    )


def safe_persist(
    obj: Any,
    dask_client: Optional[Client] = None,
    logger: SibiLogger | logging.Logger | None = None,
) -> Any:
    """Persist a collection to distributed memory with auto-healing."""

    def _persist_op(o: Any, c: Optional[Client]) -> Any:
        return c.persist(o) if c else o.persist()

    return _execute_with_resilience(_persist_op, obj, dask_client, logger)


def safe_gather(
    objs: list[Any],
    dask_client: Optional[Client] = None,
    logger: SibiLogger | logging.Logger | None = None,
) -> list[Any]:
    """Gather multiple futures or collections into local memory."""
    if not objs:
        return []

    def _gather_op(items: list[Any], client: Optional[Client]) -> list[Any]:
        if client:
            return client.gather(client.compute(items))  # type: ignore
        return list(dask.compute(*items, scheduler="threads"))

    return _execute_with_resilience(_gather_op, objs, dask_client, logger)  # type: ignore


def safe_wait(
    obj: Any,
    dask_client: Optional[Client] = None,
    timeout: Optional[float] = None,
    logger: SibiLogger | logging.Logger | None = None,
) -> Any:
    """Wait for completion. Safe from local-fallback for DataFrames."""
    log = logger or _get_log()
    client = _get_active_client(dask_client, logger=log)
    try:
        if client:
            dask_wait(obj, timeout=timeout)
        elif not _is_dask_dataframe_like(obj) and hasattr(obj, "compute"):
            obj.compute(scheduler="threads")
        return obj
    except Exception as e:
        log.warning(f"safe_wait: {type(e).__name__}: {e}")
        return obj


def safe_dry_run(
    obj: Any, logger: SibiLogger | logging.Logger | None = None
) -> dict[str, Any]:
    """Utility to log and return graph metrics without execution."""
    metrics = get_graph_metrics(obj)
    log = logger or _get_log()
    log.info("Manual Dask Dry Run", extra={"graph_metrics": metrics})
    return metrics


# ---------------------------------------------------------------------------
# Heuristic Emptiness Checks
# ---------------------------------------------------------------------------


def dask_is_probably_empty(ddf: dd.DataFrame) -> bool:
    """Metadata check (zero partitions)."""
    return int(getattr(ddf, "npartitions", 0)) == 0


def dask_is_empty_truthful(
    ddf: dd.DataFrame,
    dask_client: Optional[Client] = None,
    logger: SibiLogger | logging.Logger | None = None,
) -> bool:
    """Expensive but accurate full-table count check."""
    total = safe_compute(
        ddf.map_partitions(len).sum(),
        dask_client=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: SibiLogger | logging.Logger | None = None,
) -> bool:
    """
    Multi-stage check:
    1. Metadata
    2. Parallel sampling of first K partitions
    3. Truthful sum (fallback)
    """
    if dask_is_probably_empty(ddf):
        return True

    n_parts = int(getattr(ddf, "npartitions", 0))
    k = min(max(sample, 1), n_parts)
    try:
        parts = [
            ddf.get_partition(i).map_partitions(len)
            for i in range(k)
        ]
        probes = safe_gather(parts, dask_client, logger=logger)

        if any(_to_int_safe(n) > 0 for n in probes):
            return False

        if k == n_parts:
            return True
    except Exception as e:
        _get_log().warning(f"dask_is_empty probe failed: {e}")
        return False

    return dask_is_empty_truthful(ddf, dask_client=dask_client, logger=logger)


# ---------------------------------------------------------------------------
# Data Extraction
# ---------------------------------------------------------------------------


class UniqueValuesExtractor:
    """Resilient unique value extraction from Dask columns."""

    def __init__(
        self,
        dask_client: Optional[Client] = None,
        logger: SibiLogger | logging.Logger | None = None,
    ) -> None:
        self.dask_client = dask_client
        self.logger = logger

    async def extract_unique_values(
        self, df: dd.DataFrame, *columns: str, limit: int = 100_000
    ) -> dict[str, list[Any]]:
        async def _extract(col: str) -> tuple[str, list[Any]]:
            # Optimization: drop duplicates on the distributed collection first
            unique_dd = df[col].dropna().drop_duplicates()

            # Fetch only the stats head to avoid OOM
            # npartitions=-1 forces a logical 'head' across partitions if needed,
            # but usually for unique values we want to be careful.
            # Using compute=True on head triggers the fetch.

            # We run this in a thread because head(compute=True) is blocking
            res = await asyncio.to_thread(
                lambda: unique_dd.head(limit, npartitions=-1, compute=True)
            )

            if len(res) >= limit:
                if self.logger:
                    self.logger.warning(
                        f"Unique value extraction for column '{col}' truncated at {limit} items. "
                        "High cardinality detected.",
                        extra={"column": col, "limit": limit},
                    )

            return col, res.tolist()

        results = await asyncio.gather(*(_extract(c) for c in columns))
        return dict(results)
