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

import math
from collections.abc import Mapping, MutableMapping
from typing import (
    TYPE_CHECKING,
    Any,
    Iterable,
    Sequence,
)

import dask.dataframe as dd
import pandas as pd

from sibi_flux.core import SecureResource
from .types import DdfLike

if TYPE_CHECKING:
    from sibi_flux.logger import Logger


class DfMerger(SecureResource):
    """
    Generic Dask/Pandas DataFrame merger with integrated join-key normalization.

    Responsibilities:
      - Normalize incoming objects to Dask DataFrames.
      - Normalize join-key dtypes (left/right) to 'string' when mismatched.
      - Apply a sequence of left merges according to a merge plan.
      - Dynamic partitioning for in-memory data to optimize merge parallelism.
    """

    # Partitioning heuristics
    ROWS_PER_PARTITION = 100_000

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.dask_client = kwargs.get("dask_client")
        # Track dtype normalization changes for logging/inspection
        self._dtype_changes: MutableMapping[str, MutableMapping[str, str]] = {}

    def _normalize_merge_dtypes(
        self,
        df: dd.DataFrame,
        results: Mapping[str, dd.DataFrame | None],
        merge_plan: Iterable[tuple[str, Sequence[str], Sequence[str], Sequence[str]]],
        persist: bool = False,
    ) -> tuple[dd.DataFrame, dict[str, dd.DataFrame | None]]:
        """
        Normalize join key dtypes using safe promotion rules.
        """
        if not results:
            return df, results

        base_casts: MutableMapping[str, str] = {}

        for key, left_on, right_on, _ in merge_plan:
            other = results.get(key)
            if other is None:
                continue

            for lcol, rcol in zip(left_on, right_on):
                if lcol in df.columns and rcol in other.columns:
                    l_dtype = df._meta[lcol].dtype
                    r_dtype = other._meta[rcol].dtype

                    if l_dtype == r_dtype:
                        continue

                    # Mismatch Detected - Decide on Target Type
                    target_type = None

                    is_l_num = pd.api.types.is_numeric_dtype(l_dtype)
                    is_r_num = pd.api.types.is_numeric_dtype(r_dtype)

                    if is_l_num and is_r_num:
                        # Both numeric (e.g. Int vs Float) -> Promote to Float
                        target_type = "float64"
                    else:
                        # At least one is non-numeric -> Fallback to String
                        target_type = "string"

                    # If Left needs casting
                    if str(l_dtype) != target_type:
                        base_casts[lcol] = target_type
                        if self.logger:
                            self.logger.debug(
                                f"Merge Type Fix: '{lcol}' ({l_dtype}) vs '{rcol}' ({r_dtype}) -> Casting '{lcol}' to {target_type}"
                            )

        # Apply casts to Base DataFrame
        if base_casts:
            self._dtype_changes.clear()
            for col, dtype in base_casts.items():
                self._dtype_changes[col] = {str(df._meta[col].dtype): dtype}

            df = df.astype(base_casts)

        if persist:
            df = df.persist()
            results = {
                k: v.persist() if isinstance(v, dd.DataFrame) else v
                for k, v in results.items()
            }

        return df, dict(results)

    def log_normalization_changes(self) -> None:
        """Log dtype alignment actions."""
        if not self.logger or not self._dtype_changes:
            return

        for col, mapping in self._dtype_changes.items():
            for src, tgt in mapping.items():
                self.logger.info(
                    f"Merge Normalization: Cast Base Column '{col}' from {src} to {tgt}"
                )

    def _to_ddf(self, obj: Any) -> dd.DataFrame | None:
        """Normalize arbitrary objects into a Dask DataFrame with dynamic partitioning."""
        if obj is None:
            return None
        if isinstance(obj, dd.DataFrame):
            return obj

        # Conversion to Pandas first for length detection
        if isinstance(obj, pd.DataFrame):
            pdf = obj
        elif isinstance(obj, dict):
            pdf = pd.DataFrame([obj])
        elif isinstance(obj, list) and obj and all(isinstance(r, dict) for r in obj):
            pdf = pd.DataFrame(obj)
        else:
            raise TypeError(f"Unexpected type for merge: {type(obj)}")

        # Heuristic: 1 partition per ROWS_PER_PARTITION rows, min 1
        n_parts = max(1, math.ceil(len(pdf) / self.ROWS_PER_PARTITION))
        if n_parts > 1 and self.logger:
            self.logger.debug(
                f"Dynamic Partitioning: Splitting in-memory data ({len(pdf)} rows) into {n_parts} partitions."
            )

        return dd.from_pandas(pdf, npartitions=n_parts)

    @staticmethod
    def _drop_cols(df: dd.DataFrame, cols_to_drop: Sequence[str]) -> dd.DataFrame:
        """Drop columns that actually exist in the frame."""
        to_drop = [c for c in cols_to_drop if c in df.columns]
        return df.drop(columns=to_drop) if to_drop else df

    def _merge_left(
        self,
        left: dd.DataFrame,
        right_like: dd.DataFrame | None,
        *,
        left_on: Sequence[str],
        right_on: Sequence[str],
    ) -> dd.DataFrame:
        """
        Perform a left merge. Assumes join-key dtypes already normalized.
        """
        if right_like is None:
            return left

        return left.merge(
            right_like,
            how="left",
            left_on=list(left_on),
            right_on=list(right_on),
        )

    def apply_merges(
        self,
        base_df: DdfLike,
        results: Mapping[str, DdfLike | None],
        merge_plan: Iterable[tuple[str, Sequence[str], Sequence[str], Sequence[str]]],
        *,
        normalize_persist: bool = False,
    ) -> dd.DataFrame:
        """
        Apply a sequence of left merges defined by `merge_plan`.
        """
        # Ensure we work with Dask for the base
        if isinstance(base_df, dd.DataFrame):
            df = base_df
        else:
            df_to_ddf = self._to_ddf(base_df)
            if df_to_ddf is None:
                 raise ValueError("Base DataFrame cannot be None")
            df = df_to_ddf

        # Ensure all auxiliary results are dd.DataFrame or None
        dd_results: MutableMapping[str, dd.DataFrame | None] = {}
        for key, val in results.items():
            dd_results[key] = self._to_ddf(val) if val is not None else None

        # Normalize join-key dtypes across base and related frames
        df, norm_results = self._normalize_merge_dtypes(
            df, dd_results, merge_plan, persist=normalize_persist
        )
        self.log_normalization_changes()

        # Apply merges
        out = df
        start_partitions = out.npartitions
        
        for i, (key, left_on, right_on, drop_cols) in enumerate(merge_plan):
            other = norm_results.get(key)
            if other is None:
                continue
            out = self._merge_left(out, other, left_on=left_on, right_on=right_on)
            out = self._drop_cols(out, drop_cols)

            # Resilience: Break graph lineage after multiple joins or partition growth
            # We repartition if the number of partitions has grown significantly (> 2x)
            # or every 5 joins to keep the graph shallow.
            if i > 0 and (i % 5 == 0 or out.npartitions > start_partitions * 2):
                if self.logger:
                    self.logger.info(
                        f"Refactoring merge graph at step {i} (partitions: {out.npartitions})"
                    )
                try:
                    # Consolidate partitions if they became too many
                    out = out.repartition(partition_size="100MB")
                    if normalize_persist:
                        out = out.persist()
                    start_partitions = out.npartitions
                except Exception as e:
                    if self.logger:
                        self.logger.warning(
                            f"Failed to repartition merge intermediate: {e}"
                        )

        return out
