from __future__ import annotations

from collections.abc import Iterable, Sequence
from typing import Any, Literal

import pandas as pd

BalanceStrategy = Literal["none", "downsample", "cap"]


def canon_label(x: Any) -> str:
    """
    Canonicalize labels for balancing purposes only:
    - 'benign' and 'normal' -> 'benign'
    - 'malicious' -> 'malicious'
    - otherwise, returns lowercase
    """
    s = str(x).strip().lower()
    if s in {"benign", "normal"}:
        return "benign"
    if s == "malicious":
        return "malicious"
    return s


def _ensure_labels(df: pd.DataFrame, labels: Sequence[str] | None) -> list[str]:
    """
    Return the canonical target labels to consider and validate presence.
    If labels is None, returns all canonical labels present in df['label'].
    """
    canon_series = df["label"].map(canon_label)
    present = canon_series.value_counts().sort_index().index.tolist()
    if labels is None:
        return present

    requested = [canon_label(x) for x in labels]
    missing = [x for x in requested if x not in present]
    if missing:
        raise ValueError(
            f"Requested labels not present in data (after normalization): {missing}"
        )
    return requested


def _filter_to_labels(df: pd.DataFrame, labels: Iterable[str]) -> pd.DataFrame:
    """
    Keep only rows whose canonical label is in 'labels'.
    Preserves original df['label'] values.
    """
    canon = df["label"].map(canon_label)
    mask = canon.isin(set(labels))
    return df.loc[mask].copy()


def _stratified_take_total(
    df: pd.DataFrame,
    canon_series: pd.Series,
    *,
    target_total: int,
    order: Sequence[str],
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Stratified sampling to an exact 'target_total' rows without upsampling.
    Distributes rows across classes in 'order' as evenly as possible:

      base = target_total // k
      remainder distributed in deterministic label order, respecting each class capacity
    """
    k = len(order)
    if target_total <= 0 or k == 0:
        return df.iloc[0:0].copy()

    counts = canon_series.value_counts()
    per_class = {lbl: min(counts.get(lbl, 0), target_total // k) for lbl in order}
    current = sum(per_class.values())
    remainder = max(0, target_total - current)

    for lbl in order:
        if remainder <= 0:
            break
        cap = counts.get(lbl, 0) - per_class[lbl]
        if cap > 0:
            take = min(cap, remainder)
            per_class[lbl] += take
            remainder -= take

    parts = []
    for lbl in order:
        n = per_class[lbl]
        if n <= 0:
            continue
        g = df[canon_series == lbl]
        if not g.empty:
            parts.append(g.sample(n=n, random_state=random_state, replace=False))

    if not parts:
        return df.iloc[0:0].copy()

    return pd.concat(parts).sample(frac=1.0, random_state=random_state).reset_index(drop=True)


def balance_downsample(
    df: pd.DataFrame,
    labels: Sequence[str] | None = None,
    *,
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Classic downsampling: reduce each class to the size of the smallest class.
    """
    if "label" not in df.columns or df.empty:
        return df

    selected = _ensure_labels(df, labels)
    work = _filter_to_labels(df, selected)
    if work.empty:
        return work

    canon = work["label"].map(canon_label)
    counts = canon.value_counts()
    if len(counts) <= 1:
        return work

    n_min = int(counts.min())
    parts = []
    for lbl in selected:
        g = work[canon == lbl]
        take = min(n_min, len(g))
        if take > 0:
            parts.append(g.sample(n=take, random_state=random_state, replace=False))

    if not parts:
        return work.iloc[0:0].copy()

    return pd.concat(parts).sample(frac=1.0, random_state=random_state).reset_index(drop=True)


def balance_cap(
    df: pd.DataFrame,
    labels: Sequence[str] | None = None,
    *,
    per_class_cap: int,
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Per-class cap (no upsampling): limit each class to at most 'per_class_cap' rows.
    """
    if per_class_cap is None or per_class_cap <= 0:
        raise ValueError("per_class_cap must be a positive integer")

    if "label" not in df.columns or df.empty:
        return df

    selected = _ensure_labels(df, labels)
    work = _filter_to_labels(df, selected)
    if work.empty:
        return work

    canon = work["label"].map(canon_label)
    parts = []
    for lbl in selected:
        g = work[canon == lbl]
        take = min(len(g), int(per_class_cap))
        if take > 0:
            parts.append(g.sample(n=take, random_state=random_state, replace=False))

    if not parts:
        return work.iloc[0:0].copy()

    return pd.concat(parts).sample(frac=1.0, random_state=random_state).reset_index(drop=True)


def stratified_cut(
    df: pd.DataFrame,
    *,
    labels: Sequence[str] | None = None,
    total: int,
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Deterministic, stratified *cut* to 'total' rows (no upsampling).
    Useful AFTER balancing to keep the final sample balanced when total is even
    (for two classes, ensures 50/50 if data allows).
    """
    if total is None or total <= 0:
        return df.iloc[0:0].copy()

    selected = _ensure_labels(df, labels)
    work = _filter_to_labels(df, selected)
    if work.empty:
        return work

    canon = work["label"].map(canon_label)
    total = min(total, len(work))  # never exceed available rows

    return _stratified_take_total(
        work,
        canon,
        target_total=total,
        order=list(selected),
        random_state=random_state,
    )


def balance_dataframe(
    df: pd.DataFrame,
    *,
    strategy: BalanceStrategy = "none",
    labels: Sequence[str] | None = None,
    per_class_cap: int | None = None,
    random_state: int = 42,
) -> pd.DataFrame:
    """
    Unified entry point for balancing:
      - strategy='none'       -> no balancing
      - strategy='downsample' -> reduce to the smallest class
      - strategy='cap'        -> cap each class to per_class_cap
    """
    if strategy == "none" or "label" not in df.columns or df.empty:
        return df

    if strategy == "downsample":
        return balance_downsample(df, labels, random_state=random_state)

    if strategy == "cap":
        if per_class_cap is None:
            raise ValueError("per_class_cap is required for strategy='cap'")
        return balance_cap(df, labels, per_class_cap=per_class_cap, random_state=random_state)

    raise ValueError(f"Unknown balance strategy: {strategy}")
