from __future__ import annotations

from dataclasses import dataclass

import pandas as pd


BENCHMARK_CATALOG: dict[str, dict[str, object]] = {
    "SWE-Bench Verified": {
        "category": "Coding",
        "capabilities": ["repository-repair", "software-engineering"],
    },
    "SWE-Bench Pro -- Ansible": {
        "category": "Coding",
        "capabilities": ["repository-repair", "software-engineering", "ansible"],
    },
    "RH SWE-Bench": {
        "category": "Coding",
        "capabilities": ["repository-repair", "software-engineering"],
    },
    "Terminal Bench 2.0": {
        "category": "Generalist",
        "capabilities": ["shell", "tool-use"],
    },
    "Shellbench": {
        "category": "Generalist",
        "capabilities": ["shell", "tool-use"],
    },
}

DEFAULT_BENCHMARK_CATEGORY = "Other"


@dataclass(frozen=True)
class MetricSpec:
    column: str
    label: str
    higher_is_better: bool
    positive_only: bool = False


RANKING_METRICS: dict[str, MetricSpec] = {
    "Score": MetricSpec("Score (%)", "Score (%)", True),
    "Total tokens": MetricSpec("Total Tokens Per Task", "Total tokens per task", False, True),
    "Input tokens": MetricSpec("Input Tokens Per Task", "Input tokens per task", False, True),
    "Output tokens": MetricSpec("Output Tokens Per Task", "Output tokens per task", False, True),
    "Cache tokens": MetricSpec("Cache Tokens Per Task", "Cache tokens per task", False, True),
    "Cost": MetricSpec("Cost Per Task", "Cost per task (USD)", False, True),
    "Response time": MetricSpec("Total Time Per Task", "Total time per task (seconds)", False, True),
    "Agent time": MetricSpec("Agent Time Per Task", "Agent time per task (seconds)", False, True),
    "Reliability": MetricSpec("Execution Error Rate (%)", "Execution error rate (%)", False),
    "Tokens per successful task": MetricSpec("Tokens Per Successful Task", "Tokens per successful task", False, True),
    "Cost per successful task": MetricSpec("Cost Per Successful Task", "Cost per successful task (USD)", False, True),
    "Time per successful task": MetricSpec("Time Per Successful Task", "Time per successful task (seconds)", False, True),
}


TRADEOFF_METRICS: dict[str, MetricSpec] = {
    "Score": RANKING_METRICS["Score"],
    "Cost per task": RANKING_METRICS["Cost"],
    "Total tokens per task": RANKING_METRICS["Total tokens"],
    "Total time per task": RANKING_METRICS["Response time"],
    "Agent time per task": RANKING_METRICS["Agent time"],
    "Execution error rate": RANKING_METRICS["Reliability"],
}


MATRIX_METRICS: dict[str, MetricSpec] = {
    "Score": RANKING_METRICS["Score"],
    "Within-benchmark percentile": MetricSpec("Within-Benchmark Percentile", "Within-benchmark percentile", True),
    "Within-benchmark rank": MetricSpec("Within-Benchmark Rank", "Within-benchmark rank", False),
    "Total tokens": RANKING_METRICS["Total tokens"],
    "Cost": RANKING_METRICS["Cost"],
    "Total time": RANKING_METRICS["Response time"],
    "Agent time": RANKING_METRICS["Agent time"],
    "Execution error rate": RANKING_METRICS["Reliability"],
}


def benchmark_metadata(name: str) -> dict[str, object]:
    metadata = BENCHMARK_CATALOG.get(name)
    if metadata is not None:
        return metadata
    return {"category": DEFAULT_BENCHMARK_CATEGORY, "capabilities": []}


def benchmark_category(name: str) -> str:
    return str(benchmark_metadata(name)["category"])


def benchmarks_for_category(dataframe: pd.DataFrame, category: str) -> list[str]:
    if dataframe is None or dataframe.empty or "Benchmark" not in dataframe:
        return []
    names = sorted(str(value) for value in dataframe["Benchmark"].dropna().unique())
    return [name for name in names if benchmark_category(name) == category]


def enrich_analysis_df(dataframe: pd.DataFrame) -> pd.DataFrame:
    """Add PR2 taxonomy, reliability, per-success and normalized performance fields."""
    if dataframe is None:
        return pd.DataFrame()
    df = dataframe.copy()
    if df.empty:
        for column in (
            "Benchmark Category",
            "Execution Error Rate (%)",
            "Tokens Per Successful Task",
            "Cost Per Successful Task",
            "Time Per Successful Task",
            "Total Benchmark Cost",
            "Within-Benchmark Rank",
            "Within-Benchmark Percentile",
        ):
            if column not in df:
                df[column] = pd.Series(dtype="float64" if column != "Benchmark Category" else "object")
        return df

    df["Benchmark Category"] = df["Benchmark"].map(lambda value: benchmark_category(str(value)))

    tasks = pd.to_numeric(df.get("Tasks"), errors="coerce")
    errors = pd.to_numeric(df.get("Errors"), errors="coerce")
    valid_tasks = tasks.notna() & (tasks > 0)
    df["Execution Error Rate (%)"] = (errors / tasks * 100).where(valid_tasks & errors.notna())

    score_fraction = pd.to_numeric(df.get("Score"), errors="coerce")
    successful = score_fraction.notna() & (score_fraction > 0)
    for source, target in (
        ("Total Tokens Per Task", "Tokens Per Successful Task"),
        ("Cost Per Task", "Cost Per Successful Task"),
        ("Total Time Per Task", "Time Per Successful Task"),
    ):
        values = pd.to_numeric(df.get(source), errors="coerce")
        df[target] = (values / score_fraction).where(successful & values.notna() & (values > 0))

    cost = pd.to_numeric(df.get("Cost Per Task"), errors="coerce")
    df["Total Benchmark Cost"] = (cost * tasks).where(cost.notna() & (cost > 0) & valid_tasks)

    scores = pd.to_numeric(df.get("Score (%)"), errors="coerce")
    df["Within-Benchmark Rank"] = scores.groupby(df["Benchmark"]).rank(method="min", ascending=False)

    def percentile(group: pd.Series) -> pd.Series:
        valid = group.dropna()
        out = pd.Series(index=group.index, dtype=float)
        if valid.empty:
            return out
        ranks = valid.rank(method="average", ascending=False)
        if len(valid) == 1:
            out.loc[valid.index] = 100.0
        else:
            out.loc[valid.index] = 100.0 * (len(valid) - ranks) / (len(valid) - 1)
        return out

    df["Within-Benchmark Percentile"] = scores.groupby(df["Benchmark"], group_keys=False).apply(percentile)
    return df


def filter_category(dataframe: pd.DataFrame, category: str) -> pd.DataFrame:
    df = enrich_analysis_df(dataframe)
    return df[df["Benchmark Category"] == category].copy()


def ranking_df(
    dataframe: pd.DataFrame,
    metric: str,
    benchmark: str | None = None,
    category: str | None = None,
) -> pd.DataFrame:
    df = enrich_analysis_df(dataframe)
    spec = RANKING_METRICS[metric]
    if category:
        df = df[df["Benchmark Category"] == category]
    if benchmark and benchmark != "All benchmarks":
        df = df[df["Benchmark"] == benchmark]
    elif benchmark == "All benchmarks":
        if metric != "Score":
            # Resource units can be compared across benchmarks, but rows remain per benchmark;
            # do not silently aggregate them.
            pass
        else:
            return cross_benchmark_ranking_df(df)

    values = pd.to_numeric(df[spec.column], errors="coerce")
    valid = values.notna()
    if spec.positive_only:
        valid &= values > 0
    df = df.loc[valid].copy()
    df[spec.column] = values.loc[valid]
    columns = [
        "Model", "Harness", "Benchmark", "Benchmark Category", spec.column,
        "Score (%)", "Execution Error Rate (%)"
    ]
    columns = list(dict.fromkeys(column for column in columns if column in df))
    return df.sort_values(
        [spec.column, "Model", "Harness"],
        ascending=[not spec.higher_is_better, True, True],
        kind="mergesort",
    )[columns].reset_index(drop=True)


def cross_benchmark_ranking_df(
    dataframe: pd.DataFrame,
    minimum_coverage: float | int = 0.5,
) -> pd.DataFrame:
    """Aggregate within-benchmark percentiles without averaging incompatible raw scores."""
    df = enrich_analysis_df(dataframe)
    if df.empty:
        return pd.DataFrame(columns=[
            "Model", "Harness", "Normalized Performance", "Benchmarks Covered",
            "Eligible Benchmarks", "Coverage (%)",
        ])
    eligible = int(df["Benchmark"].nunique())
    grouped = (
        df.dropna(subset=["Within-Benchmark Percentile"])
        .groupby(["Model", "Harness"], as_index=False)
        .agg(
            **{
                "Normalized Performance": ("Within-Benchmark Percentile", "mean"),
                "Benchmarks Covered": ("Benchmark", "nunique"),
            }
        )
    )
    grouped["Eligible Benchmarks"] = eligible
    grouped["Coverage (%)"] = grouped["Benchmarks Covered"] / eligible * 100 if eligible else 0.0

    if isinstance(minimum_coverage, float) and minimum_coverage <= 1:
        threshold_pct = minimum_coverage * 100
        grouped = grouped[grouped["Coverage (%)"] >= threshold_pct]
    else:
        grouped = grouped[grouped["Benchmarks Covered"] >= int(minimum_coverage)]
    return grouped.sort_values(
        ["Normalized Performance", "Benchmarks Covered", "Model", "Harness"],
        ascending=[False, False, True, True],
        kind="mergesort",
    ).reset_index(drop=True)


def matrix_df(
    dataframe: pd.DataFrame,
    metric: str,
    category: str | None = None,
    include_incomplete: bool = True,
    sort_by: str = "Normalized performance",
) -> pd.DataFrame:
    df = enrich_analysis_df(dataframe)
    if category:
        df = df[df["Benchmark Category"] == category]
    if df.empty:
        return pd.DataFrame()

    df["Agent"] = df["Model"].astype(str) + " / " + df["Harness"].astype(str)
    benchmarks = sorted(df["Benchmark"].dropna().unique())
    agents = sorted(df["Agent"].dropna().unique())

    if metric == "Coverage":
        available = df.assign(_coverage=1).pivot_table(
            index="Agent", columns="Benchmark", values="_coverage", aggfunc="max"
        )
        matrix = available.reindex(index=agents, columns=benchmarks)
    else:
        spec = MATRIX_METRICS[metric]
        values = pd.to_numeric(df[spec.column], errors="coerce")
        work = df.assign(_value=values)
        matrix = work.pivot_table(index="Agent", columns="Benchmark", values="_value", aggfunc="mean")
        matrix = matrix.reindex(index=agents, columns=benchmarks)

    if not include_incomplete:
        matrix = matrix.dropna(axis=0, how="any")

    if matrix.empty:
        return matrix

    if sort_by in {"Coverage", "Coverage (high to low)", "Coverage (low to high)"}:
        ascending = sort_by == "Coverage (low to high)"
        coverage = matrix.notna().sum(axis=1)
        order = coverage.sort_values(ascending=ascending, kind="mergesort").index
        matrix = matrix.loc[order]
    elif sort_by in {"Normalized performance", "Normalized performance (high to low)", "Normalized performance (low to high)"}:
        perf = cross_benchmark_ranking_df(df, minimum_coverage=0)
        perf["Agent"] = perf["Model"] + " / " + perf["Harness"]
        perf = perf.sort_values(
            ["Normalized Performance", "Agent"],
            ascending=[sort_by == "Normalized performance (low to high)", True],
            kind="mergesort",
        )
        order = [agent for agent in perf["Agent"] if agent in matrix.index]
        order += [agent for agent in matrix.index if agent not in order]
        matrix = matrix.loc[order]
    elif sort_by == "Alphabetical (Z–A)":
        matrix = matrix.loc[sorted(matrix.index, reverse=True)]
    elif sort_by in {"Alphabetical", "Stable"}:
        matrix = matrix.loc[sorted(matrix.index)]
    return matrix


def coverage_summary(dataframe: pd.DataFrame) -> dict[str, float | int]:
    df = enrich_analysis_df(dataframe)
    total = len(df)
    def pct(column: str, positive: bool = False) -> float:
        if total == 0:
            return 0.0
        values = pd.to_numeric(df[column], errors="coerce")
        mask = values.notna()
        if positive:
            mask &= values > 0
        return float(mask.mean() * 100)
    return {
        "results": total,
        "models": int(df["Model"].nunique()) if total else 0,
        "harnesses": int(df["Harness"].nunique()) if total else 0,
        "benchmarks": int(df["Benchmark"].nunique()) if total else 0,
        "token_coverage_pct": pct("Total Tokens Per Task", True),
        "cost_coverage_pct": pct("Cost Per Task", True),
        "time_coverage_pct": pct("Total Time Per Task", True),
    }