import pandas as pd
import pytest

from flownarrator.utils import balancing as bal
from flownarrator.utils.balancing import (
    _stratified_take_total,
    balance_cap,
    balance_dataframe,
    balance_downsample,
    canon_label,
    stratified_cut,
)

# -------------------------------
# Helpers
# -------------------------------

SEED = 7

def make_df(protocol="TCP", labels=None):
    labels = labels or []
    return pd.DataFrame({"protocol": [protocol] * len(labels), "label": labels})

def counts_lower(s):
    return s.str.lower().value_counts()

# -------------------------------
# _stratified_take_total (private)
# -------------------------------

def test__stratified_take_total_k_zero_returns_empty():
    df = make_df(labels=["a", "b"])
    canon = df["label"].map(canon_label)
    out = _stratified_take_total(df, canon, target_total=2, order=[], random_state=SEED)
    assert out.empty

def test__stratified_take_total_base_zero_with_remainder_and_zero_capacity():
    # df contiene solo 'A'; total=1; order fissa -> prende 'a'
    df = make_df(labels=["A"])
    canon = df["label"].map(canon_label)
    out = _stratified_take_total(df, canon, target_total=1, order=["a", "b"], random_state=SEED)
    assert len(out) == 1 and out["label"].str.lower().iloc[0] == "a"

def test__stratified_take_total_df_empty_no_parts():
    df = make_df(labels=[])
    canon = df["label"].map(canon_label)
    out = _stratified_take_total(df, canon, target_total=3, order=["a", "b"], random_state=SEED)
    assert out.empty

# -------------------------------
# Canonicalizzazione
# -------------------------------

@pytest.mark.parametrize(
    "inp,exp",
    [
        ("Benign", "benign"),
        ("NORMAL", "benign"),
        ("MALICIOUS", "malicious"),
        ("AnythingElse", "anythingelse"),
        (42, "42"),
        ("Normal", "benign"),
        ("Malicious", "malicious"),
        ("Foo", "foo"),
        (5, "5"),
    ],
)
def test_canon_label(inp, exp):
    assert canon_label(inp) == exp
    assert bal.canon_label(inp) == exp  # path tramite alias modulo

# -------------------------------
# balance_cap
# -------------------------------

@pytest.mark.parametrize("cap", [0, -3])
def test_balance_cap_invalid_cap_raises(cap):
    df = make_df(labels=["x"])
    with pytest.raises(ValueError):
        balance_cap(df, per_class_cap=cap, random_state=SEED)

def test_balance_cap_empty_df_returns_empty():
    df = pd.DataFrame(columns=["protocol", "label"])
    out = balance_cap(df, per_class_cap=2, random_state=SEED)
    assert out.empty

def test_balance_cap_labels_missing_raises():
    df = make_df(labels=["x", "y"])
    with pytest.raises(ValueError):
        balance_cap(df, labels=["x", "zzz"], per_class_cap=1, random_state=SEED)

def test_balance_cap_subset_and_capacity():
    # x=3, y=1, cap=2 -> x:2, y:1
    df = make_df(labels=["x", "x", "x", "y"])
    out = balance_cap(df, labels=["x", "y"], per_class_cap=2, random_state=SEED)
    c = counts_lower(out["label"])
    assert c.get("x", 0) == 2 and c.get("y", 0) == 1

# -------------------------------
# balance_downsample
# -------------------------------

def test_balance_downsample_labels_empty_returns_empty():
    df = make_df(labels=["a", "a", "b"])
    out = balance_downsample(df, labels=[], random_state=SEED)
    assert out.empty

def test_balance_downsample_ensure_labels_missing_raises():
    df = make_df(labels=["a", "a", "b", "b"])
    with pytest.raises(ValueError):
        balance_downsample(df, labels=["a", "c"], random_state=SEED)

def test_balance_downsample_single_class_noop():
    df = make_df(labels=["a", "a", "a"])
    out = balance_downsample(df, random_state=SEED)
    assert len(out) == 3

def test_balance_downsample_multi_class_min_rule():
    # a:3, b:1, c:2 -> min=1 => 1 per classe
    df = make_df(labels=["a", "a", "a", "b", "c", "c"])
    out = balance_downsample(df, random_state=SEED)
    c = counts_lower(out["label"])
    assert c.get("a", 0) == 1 and c.get("b", 0) == 1 and c.get("c", 0) == 1

# -------------------------------
# stratified_cut
# -------------------------------

@pytest.mark.parametrize("total", [0, -3])
def test_stratified_cut_total_le_zero_returns_empty(total):
    df = make_df(labels=["a", "a", "b", "b"])
    assert stratified_cut(df, total=total, random_state=SEED).empty

def test_stratified_cut_total_exceeds_len_clamped():
    df = make_df(labels=["a", "a", "b"])
    out = stratified_cut(df, total=10, random_state=SEED)
    assert len(out) == 3

def test_stratified_cut_remainder_distribution_and_capacity():
    # A:3, B:1, total=3 -> base=1 ciascuna, remainder=1 -> va su A (cap>0)
    df = make_df(labels=["A", "A", "A", "B"])
    out = stratified_cut(df, total=3, random_state=SEED)
    c = counts_lower(out["label"])
    assert c.get("a", 0) == 2 and c.get("b", 0) == 1

def test_stratified_cut_labels_empty_list_results_empty():
    df = make_df(labels=["x", "x", "y"])
    out = stratified_cut(df, labels=[], total=2, random_state=SEED)
    assert out.empty

def test_stratified_cut_even_and_missing_labels_error_paths():
    df = make_df(labels=["a", "a", "b", "b", "a", "a"])  # a:4, b:2
    out = stratified_cut(df, total=4, random_state=SEED)
    c = counts_lower(out["label"])
    assert c.get("a", 0) == 2 and c.get("b", 0) == 2
    with pytest.raises(ValueError):
        stratified_cut(df, labels=["a", "z"], total=2)

# -------------------------------
# balance_dataframe (dispatcher)
# -------------------------------

def test_balance_dataframe_none_passthrough():
    df = make_df(labels=["a", "b"])
    out = balance_dataframe(df, strategy="none", random_state=SEED)
    assert out.equals(df)

def test_balance_dataframe_downsample_and_cap_paths_and_errors():
    df = make_df(labels=["x", "y"])
    # path "downsample" not raise
    balance_dataframe(df, strategy="downsample", random_state=SEED)
    # "cap" senza per_class_cap raise
    with pytest.raises(ValueError):
        balance_dataframe(df, strategy="cap", per_class_cap=None, random_state=SEED)

@pytest.mark.parametrize("bad", ["whatever", "__unknown__", "???"])
def test_balance_dataframe_unknown_strategy_raises(bad):
    df = make_df(labels=["a"])
    with pytest.raises(ValueError):
        balance_dataframe(df, strategy=bad, random_state=SEED)


def test_downsample_and_cap_unit_alias_module():
    df = make_df(labels=["A", "A", "A", "A", "B", "B"])
    out = bal.balance_downsample(df, random_state=SEED)
    c = out["label"].value_counts()
    assert c["A"] == c["B"]
    out2 = bal.balance_cap(df, per_class_cap=1, random_state=SEED)
    assert all(v == 1 for v in out2["label"].value_counts().values)
