# fmatch/core/blocking/guardrails.py
from dataclasses import dataclass
from typing import Iterable, Optional, Tuple, List
import pandas as pd
import logging
import time
from fmatch.core.sampling import AdaptiveSampler  # reuse your module

log = logging.getLogger(__name__)


@dataclass
class BlockGuardrails:
    max_block_rows: int = 50_000
    max_pairs: int = 50_000_000
    emergency_sample: int = 1000
    split_timeout_s: float = 5.0


def _is_safe_block(n: int, guard: BlockGuardrails) -> bool:
    pairs = (n * (n - 1)) // 2
    return (n <= guard.max_block_rows) and (pairs <= guard.max_pairs)


def _split_with_timeout(
    df: pd.DataFrame, idx: pd.Index, secondary: Tuple[str, int], guard: BlockGuardrails
) -> List[pd.Index]:
    start = time.time()
    col, key_len = secondary
    if col not in df.columns:
        return []
    keys = df.loc[idx, col].astype(str).str.lower().str[:key_len]
    shards = []
    for _, grp in df.loc[idx].groupby(keys, sort=False):
        s = grp.index
        if _is_safe_block(len(s), guard):
            shards.append(s)
        else:
            shards.append(s)  # still too big; caller will sample these
        if time.time() - start > guard.split_timeout_s:
            raise TimeoutError("block split timeout")
    return shards


def _adaptive_sample(df: pd.DataFrame, idx: pd.Index, target: int) -> pd.Index:
    sampler = AdaptiveSampler()
    # Sample indices until stable (use proportion metric)
    series = pd.Series(idx)
    sample, _, _, _ = sampler.sample_until_stable(
        series, lambda xs: min(1.0, len(xs) / max(1, len(series)))
    )
    return pd.Index(sample[:target])


def enforce_block_limits(
    df: pd.DataFrame,
    row_idx: Iterable[int],
    guard: BlockGuardrails,
    secondary: Optional[Tuple[str, int]] = None,
) -> List[pd.Index]:
    try:
        if df is None or len(df) == 0:
            log.warning("Guardrails: empty dataframe")
            return []
        idx = pd.Index(row_idx)
        n = len(idx)
        if _is_safe_block(n, guard):
            return [idx]
        # Try split
        if secondary:
            try:
                shards = _split_with_timeout(df, idx, secondary, guard)
                safe, big = [], []
                for s in shards:
                    (safe if _is_safe_block(len(s), guard) else big).append(s)
                if big:
                    log.warning(f"{len(big)} shard(s) still exceed limits; sampling.")
                    safe.extend(
                        [_adaptive_sample(df, s, guard.max_block_rows) for s in big]
                    )
                if safe:
                    return safe
            except TimeoutError:
                log.warning(f"Shard timeout for block size {n}")
            except Exception as e:
                log.error(f"Shard failed: {e}", exc_info=True)
        # Adaptive sample whole block
        return [_adaptive_sample(df, idx, guard.max_block_rows)]
    except Exception as e:
        log.error(f"Guardrail enforcement failed: {e} — emergency sample.")
        sample_size = min(guard.emergency_sample, len(row_idx))
        return [df.sample(n=sample_size, random_state=42).index]
