import pytest

pd = pytest.importorskip("pandas")

from dot_metrics import ComputedValue, MetricSet
from dot_metrics.explore import _to_dataframe


def _make_batch_tuple_keys():
    ms = MetricSet()
    ms.add("score", lambda x: x["s"])
    ms.add_constraint("errors", lambda x: x["e"], threshold=1)
    return ms.compute_batch({
        ("gpt4", "en"): {"s": 0.9, "e": 0},
        ("gpt4", "fr"): {"s": 0.7, "e": 2},
    })


def _make_batch_string_keys():
    ms = MetricSet()
    ms.add("score", lambda x: x)
    return ms.compute_batch({"a": 0.3, "b": 0.7})


def _make_eval_result():
    ms = MetricSet()
    ms.add("score", lambda x: x)
    return ms.compute(0.5)


def test_tuple_keys_columns():
    df = _to_dataframe(_make_batch_tuple_keys())
    assert "key[0]" in df.columns
    assert "key[1]" in df.columns
    assert "key" not in df.columns
    assert len(df) == 2


def test_string_keys_column():
    df = _to_dataframe(_make_batch_string_keys())
    assert "key" in df.columns
    assert "key[0]" not in df.columns
    assert list(df["key"]) == ["a", "b"]


def test_eval_result_wrapping():
    df = _to_dataframe(_make_eval_result())
    assert len(df) == 1
    assert "key" in df.columns
    assert df["key"].iloc[0] == "_single"


def test_metric_and_constraint_columns():
    df = _to_dataframe(_make_batch_tuple_keys())
    assert "metric:score" in df.columns
    assert "constraint:errors" in df.columns
    assert "constraint_passed:errors" in df.columns
    assert "constraints_ok" in df.columns


def test_constraint_passed_values():
    df = _to_dataframe(_make_batch_tuple_keys())
    # ("gpt4", "en") has errors=0 <= threshold=1, so passed
    row_en = df[df["key[1]"] == "en"].iloc[0]
    assert row_en["constraint_passed:errors"] == True
    # ("gpt4", "fr") has errors=2 > threshold=1, so failed
    row_fr = df[df["key[1]"] == "fr"].iloc[0]
    assert row_fr["constraint_passed:errors"] == False


def test_constraints_ok_values():
    df = _to_dataframe(_make_batch_tuple_keys())
    row_en = df[df["key[1]"] == "en"].iloc[0]
    assert row_en["constraints_ok"] == True
    row_fr = df[df["key[1]"] == "fr"].iloc[0]
    assert row_fr["constraints_ok"] == False


def test_no_constraints_no_constraints_ok():
    df = _to_dataframe(_make_batch_string_keys())
    assert "constraints_ok" not in df.columns


def test_debug_columns_present():
    ms = MetricSet()
    ms.add("info", lambda x: ComputedValue(value=x, debug={"detail": "yes"}))
    batch = ms.compute_batch({"a": 0.5})
    df = _to_dataframe(batch)
    assert "debug:info" in df.columns
    assert '"detail"' in df["debug:info"].iloc[0]


def test_debug_columns_absent_when_no_debug():
    df = _to_dataframe(_make_batch_string_keys())
    debug_cols = [c for c in df.columns if c.startswith("debug:")]
    assert debug_cols == []


def test_key_names_tuple_keys():
    ms = MetricSet()
    ms.add("score", lambda x: x["s"])
    ms.add_constraint("errors", lambda x: x["e"], threshold=1)
    batch = ms.compute_batch({
        ("gpt4", "en"): {"s": 0.9, "e": 0},
        ("gpt4", "fr"): {"s": 0.7, "e": 2},
    }, key_names=["model", "lang"])
    df = _to_dataframe(batch)
    assert "model" in df.columns
    assert "lang" in df.columns
    assert "key[0]" not in df.columns


def test_key_names_string_keys():
    ms = MetricSet()
    ms.add("score", lambda x: x)
    batch = ms.compute_batch({"a": 0.3, "b": 0.7}, key_names=["run"])
    df = _to_dataframe(batch)
    assert "run" in df.columns
    assert "key" not in df.columns
