# src/sibi_flux/df_enricher/async_enricher.py
from __future__ import annotations

import asyncio
from typing import (
    TYPE_CHECKING,
    Any,
    Sequence,
)

import pandas as pd
import dask.dataframe as dd

from sibi_flux.core import SecureResource
from sibi_flux.utils.dask_utils import sanitize_dask_df
from .types import DdfLike
from .specs import AttachmentSpec
from .merger import DfMerger

if TYPE_CHECKING:
    from sibi_flux.logger import Logger


class AsyncDfEnricher(SecureResource):
    """
    Generic async DataFrame enricher driven by AttachmentSpec definitions.

    - Uses a base DataFrame (Pandas or Dask).
    - Chooses applicable AttachmentSpecs based on required_cols.
    - Extracts unique values per spec.col_to_kwarg with safety guards.
    - Runs all attachment_fn calls concurrently.
    - Merges results using DfMerger.
    """

    MAX_UNIQUE_LIMIT = 500_000
    MAX_PARTITIONS_FOR_UNIQUE = 1_000  # Safety guard for massive dask columns
    MAX_CONCURRENT_COMPUTES = 3  # Limit parallel compute() calls to avoid OOM

    def __init__(
        self,
        base_df: DdfLike,
        specs: Sequence[AttachmentSpec],
        max_unique: int = MAX_UNIQUE_LIMIT,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.specs = list(specs)
        self.max_unique = max_unique
        self.df: dd.DataFrame = (
            base_df
            if isinstance(base_df, dd.DataFrame)
            else dd.from_pandas(base_df, npartitions=1)
        )

        if self.df is not None:
            self.df = sanitize_dask_df(self.df)
        
        # Semaphore for controlling heavy compute operations
        self._compute_semaphore = asyncio.Semaphore(self.MAX_CONCURRENT_COMPUTES)

    async def _get_unique_values(self, col_name: str) -> list[Any]:
        """
        Asynchronously retrieve unique values from a column with a safety limit.
        Includes a check for high partition counts to avoid excessive scheduler load.
        """

        async with self._compute_semaphore:
            def _fetch() -> list[Any]:
                series = self.df[col_name]
                if isinstance(series, pd.Series):
                    uniques = series.dropna().unique()
                else:
                    # Dask Series Safety Check
                    n_parts = series.npartitions
                    if n_parts > self.MAX_PARTITIONS_FOR_UNIQUE:
                        if self.logger:
                            self.logger.warning(
                                f"Column '{col_name}' has {n_parts} partitions. "
                                "Calculating uniques might be slow/heavy."
                            )

                    # Use dask.dataframe.unique() but compute safely
                    uniques_dask = series.dropna().unique()
                    
                    # Computation happens here
                    uniques = uniques_dask.compute()

                if len(uniques) > self.max_unique:
                    raise ValueError(
                        f"Column '{col_name}' has {len(uniques)} unique values, "
                        f"exceeding limit of {self.max_unique}. Enrichment aborted to prevent OOM."
                    )
                return list(uniques)

            return await asyncio.to_thread(_fetch)

    async def _resolve_spec_kwargs(
        self, spec: AttachmentSpec
    ) -> dict[str, Any] | None:
        """
        Resolve kwargs for a single spec concurrently.
        """
        available_cols = self.df.columns
        if not spec.is_applicable(available_cols):
            return None

        kwargs_for_fn: dict[str, Any] = {}

        tasks = []
        kw_names = []

        for col, kw in spec.col_to_kwarg.items():
            if col not in available_cols:
                continue
            kw_names.append(kw)
            tasks.append(self._get_unique_values(col))

        if not tasks:
            return None

        # Run all unique value extractions for this spec concurrently
        unique_lists = await asyncio.gather(*tasks)

        for kw, ulist in zip(kw_names, unique_lists):
            kwargs_for_fn[kw] = ulist

        return kwargs_for_fn

    async def _run_attachments(
        self,
        cols: Sequence[str],
    ) -> tuple[dict[str, DdfLike | None], list[AttachmentSpec]]:

        # 1. Prepare specs that apply
        available_cols = self.df.columns
        applicable_specs = [s for s in self.specs if s.is_applicable(available_cols)]
        if not applicable_specs:
            return {}, []

        # 2. Resolve kwargs for ALL specs in parallel
        resolve_tasks = [self._resolve_spec_kwargs(s) for s in applicable_specs]
        resolved_kwargs_list = await asyncio.gather(
            *resolve_tasks, return_exceptions=True
        )

        # 3. Launch attachment functions for successfully resolved kwargs
        attach_tasks: dict[str, asyncio.Task[DdfLike | None]] = {}
        valid_specs: list[AttachmentSpec] = []

        for spec, result in zip(applicable_specs, resolved_kwargs_list):
            if isinstance(result, Exception):
                if self.logger:
                    self.logger.warning(f"Spec setup failed for {spec.key}: {result}")
                continue
            if not result:  # None or empty dict
                continue

            # Launch the actual enrichment (API call / DB fetch)
            attach_tasks[spec.key] = asyncio.create_task(spec.attachment_fn(**result))
            valid_specs.append(spec)

        results: dict[str, DdfLike | None] = {}
        if attach_tasks:
            done = await asyncio.gather(*attach_tasks.values(), return_exceptions=True)
            for name, result in zip(attach_tasks.keys(), done):
                if isinstance(result, Exception):
                    if self.logger:
                        self.logger.warning("Attachment %s failed: %s", name, result)
                    results[name] = None
                else:
                    results[name] = result

        return results, valid_specs

    async def enrich(
        self,
        cols: Sequence[str],
        *,
        normalize_persist: bool = False,
        persist_result: bool = False,
    ) -> dd.DataFrame:
        if not cols:
            raise ValueError("Pass at least one column to drive enrichment.")

        results, used_specs = await self._run_attachments(cols)
        if not used_specs:
            return self.df.persist() if persist_result else self.df

        merge_plan: list[tuple[str, list[str], list[str], list[str]]] = [
            (spec.key, spec.left_on, spec.right_on, spec.drop_cols)
            for spec in used_specs
        ]

        merger = DfMerger(logger=self.logger, debug=self.debug)
        enriched = merger.apply_merges(
            base_df=self.df,
            results=results,
            merge_plan=merge_plan,
            normalize_persist=normalize_persist,
        )

        if persist_result:
            enriched = enriched.persist()

        self.df = enriched
        return enriched
