#!/usr/bin/env python3
"""Aggregate Agent evaluation JSONL and compare paired variants."""

from __future__ import annotations

import argparse
import json
import math
import statistics
from collections import defaultdict
from pathlib import Path
from typing import Any


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="JSONL result records.")
    parser.add_argument("--baseline", help="Baseline variant_id for paired comparison.")
    parser.add_argument("--variant", help="Candidate variant_id for paired comparison.")
    parser.add_argument("--format", choices=("json", "markdown"), default="json")
    parser.add_argument("--output", help="Optional output file.")
    return parser.parse_args()


def load_records(path: Path) -> list[dict[str, Any]]:
    records = []
    for line_number, raw_line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        if not raw_line.strip():
            continue
        value = json.loads(raw_line)
        if not isinstance(value, dict):
            raise ValueError(f"line {line_number}: record must be an object")
        for field in ("task_id", "variant_id", "model_id", "run_id", "status"):
            if not str(value.get(field, "")).strip():
                raise ValueError(f"line {line_number}: missing {field}")
        records.append(value)
    if not records:
        raise ValueError("input contains no result records")
    return records


def metric(record: dict[str, Any], name: str) -> float:
    value = (record.get("metrics") or {}).get(name, 0)
    return float(value or 0)


def is_infrastructure_error(record: dict[str, Any]) -> bool:
    return record.get("status") in {"provider_error", "infrastructure_error"}


def is_success(record: dict[str, Any]) -> bool:
    return (
        not is_infrastructure_error(record)
        and record.get("status") == "completed"
        and bool((record.get("validator") or {}).get("passed", False))
    )


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = (len(ordered) - 1) * fraction
    lower = int(index)
    upper = min(lower + 1, len(ordered) - 1)
    weight = index - lower
    return round(ordered[lower] * (1 - weight) + ordered[upper] * weight, 3)


def rate(numerator: int, denominator: int) -> float | None:
    return round(numerator / denominator, 4) if denominator else None


def wilson_interval(successes: int, total: int, z: float = 1.959963984540054) -> list[float] | None:
    if total == 0:
        return None
    proportion = successes / total
    z_squared = z * z
    denominator = 1 + z_squared / total
    center = (proportion + z_squared / (2 * total)) / denominator
    margin = (
        z
        * math.sqrt(
            proportion * (1 - proportion) / total
            + z_squared / (4 * total * total)
        )
        / denominator
    )
    return [round(max(0.0, center - margin), 4), round(min(1.0, center + margin), 4)]


def mcnemar_exact_p(gains: int, regressions: int) -> float:
    discordant = gains + regressions
    if discordant == 0:
        return 1.0
    tail = sum(math.comb(discordant, value) for value in range(min(gains, regressions) + 1))
    return round(min(1.0, 2 * tail / (2**discordant)), 6)


def mean(values: list[float]) -> float | None:
    return round(statistics.fmean(values), 3) if values else None


def aggregate(records: list[dict[str, Any]]) -> dict[str, Any]:
    eligible = [record for record in records if not is_infrastructure_error(record)]
    successes = [record for record in eligible if is_success(record)]
    false_completions = [
        record
        for record in eligible
        if bool(record.get("agent_claimed_complete", False)) and not is_success(record)
    ]
    safety_violations = sum(len(record.get("safety_violations") or []) for record in eligible)
    failed_tools = sum(int(metric(record, "failed_tool_calls")) for record in eligible)
    recovered_failures = sum(int(metric(record, "recovered_failures")) for record in eligible)
    unresolved_failures = sum(int(metric(record, "unresolved_tool_failures")) for record in eligible)
    elapsed = [metric(record, "elapsed_seconds") for record in eligible]
    return {
        "records": len(records),
        "eligible_records": len(eligible),
        "infrastructure_errors": len(records) - len(eligible),
        "successes": len(successes),
        "success_rate": rate(len(successes), len(eligible)),
        "success_rate_ci95": wilson_interval(len(successes), len(eligible)),
        "false_completions": len(false_completions),
        "false_completion_rate": rate(len(false_completions), len(eligible)),
        "false_completion_rate_ci95": wilson_interval(len(false_completions), len(eligible)),
        "safety_violations": safety_violations,
        "failed_tool_calls": failed_tools,
        "recovered_failures": recovered_failures,
        "failure_recovery_rate": rate(recovered_failures, failed_tools),
        "unresolved_tool_failures": unresolved_failures,
        "elapsed_seconds_p50": percentile(elapsed, 0.50),
        "elapsed_seconds_p95": percentile(elapsed, 0.95),
        "mean_total_tokens": mean([metric(record, "total_tokens") for record in eligible]),
        "mean_tool_calls": mean([metric(record, "tool_calls") for record in eligible]),
    }


def pair_key(record: dict[str, Any]) -> tuple[str, str, str, int]:
    return (
        str(record["task_id"]),
        str(record.get("task_version") or "1"),
        str(record["model_id"]),
        int(record.get("repetition") or 0),
    )


def records_by_pair(
    records: list[dict[str, Any]],
    variant_id: str,
) -> dict[tuple[str, str, str, int], dict[str, Any]]:
    indexed: dict[tuple[str, str, str, int], dict[str, Any]] = {}
    for record in records:
        if record["variant_id"] != variant_id or is_infrastructure_error(record):
            continue
        key = pair_key(record)
        if key in indexed:
            raise ValueError(f"duplicate paired record for {variant_id}: {key}")
        indexed[key] = record
    return indexed


def compare(
    records: list[dict[str, Any]],
    baseline_id: str,
    variant_id: str,
) -> dict[str, Any]:
    baseline = records_by_pair(records, baseline_id)
    variant = records_by_pair(records, variant_id)
    keys = sorted(set(baseline) & set(variant))
    counts = {"both_pass": 0, "gain": 0, "regression": 0, "both_fail": 0}
    for key in keys:
        baseline_pass = is_success(baseline[key])
        variant_pass = is_success(variant[key])
        if baseline_pass and variant_pass:
            counts["both_pass"] += 1
        elif not baseline_pass and variant_pass:
            counts["gain"] += 1
        elif baseline_pass and not variant_pass:
            counts["regression"] += 1
        else:
            counts["both_fail"] += 1

    baseline_summary = aggregate([baseline[key] for key in keys])
    variant_summary = aggregate([variant[key] for key in keys])
    blockers = []
    if variant_summary["safety_violations"] > baseline_summary["safety_violations"]:
        blockers.append("安全违规增加")
    if variant_summary["false_completions"] > baseline_summary["false_completions"]:
        blockers.append("假完成增加")

    recommendation = "continue"
    if blockers:
        recommendation = "reject"
    elif counts["gain"] > counts["regression"]:
        recommendation = "promotion-candidate"
    elif (
        counts["gain"] == counts["regression"]
        and variant_summary["success_rate"] == baseline_summary["success_rate"]
        and variant_summary["elapsed_seconds_p50"] is not None
        and baseline_summary["elapsed_seconds_p50"] is not None
        and variant_summary["elapsed_seconds_p50"] <= baseline_summary["elapsed_seconds_p50"] * 0.85
    ):
        recommendation = "efficiency-candidate"

    return {
        "baseline": baseline_id,
        "variant": variant_id,
        "paired_records": len(keys),
        "unpaired_baseline_records": len(set(baseline) - set(variant)),
        "unpaired_variant_records": len(set(variant) - set(baseline)),
        **counts,
        "discordant_pairs": counts["gain"] + counts["regression"],
        "net_gain": counts["gain"] - counts["regression"],
        "mcnemar_exact_p_two_sided": mcnemar_exact_p(counts["gain"], counts["regression"]),
        "baseline_summary": baseline_summary,
        "variant_summary": variant_summary,
        "blockers": blockers,
        "advisory_recommendation": recommendation,
    }


def render_markdown(report: dict[str, Any]) -> str:
    lines = ["# Agent 评估汇总", ""]
    lines.append("| 变体 | 可判分运行 | 成功率 [95% CI] | 假完成率 [95% CI] | 安全违规 | p50 秒 | 平均 Token |")
    lines.append("| --- | ---: | ---: | ---: | ---: | ---: | ---: |")
    for variant, summary in sorted(report["variants"].items()):
        lines.append(
            f"| {variant} | {summary['eligible_records']} | "
            f"{summary['success_rate']} {summary['success_rate_ci95']} | "
            f"{summary['false_completion_rate']} {summary['false_completion_rate_ci95']} | "
            f"{summary['safety_violations']} | {summary['elapsed_seconds_p50']} | "
            f"{summary['mean_total_tokens']} |"
        )
    comparison = report.get("comparison")
    if comparison:
        lines.extend(
            [
                "",
                "## 配对对比",
                "",
                f"- 基线：`{comparison['baseline']}`",
                f"- 候选变体：`{comparison['variant']}`",
                f"- 配对运行数：{comparison['paired_records']}",
                f"- 未配对基线运行：{comparison['unpaired_baseline_records']}",
                f"- 未配对候选运行：{comparison['unpaired_variant_records']}",
                f"- 新增成功：{comparison['gain']}",
                f"- 回归失败：{comparison['regression']}",
                f"- 双方成功：{comparison['both_pass']}",
                f"- 双方失败：{comparison['both_fail']}",
                f"- 净增益：{comparison['net_gain']}",
                f"- 配对差异精确检验 p 值：{comparison['mcnemar_exact_p_two_sided']}",
                f"- 机械建议：`{comparison['advisory_recommendation']}`",
            ]
        )
        if comparison["blockers"]:
            lines.append(f"- 阻断项：{', '.join(comparison['blockers'])}")
    return "\n".join(lines) + "\n"


def main() -> int:
    args = parse_args()
    if bool(args.baseline) != bool(args.variant):
        raise SystemExit("--baseline and --variant must be provided together")
    records = load_records(Path(args.input))
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for record in records:
        grouped[str(record["variant_id"])].append(record)
    report: dict[str, Any] = {
        "input_records": len(records),
        "variants": {variant: aggregate(items) for variant, items in sorted(grouped.items())},
    }
    if args.baseline and args.variant:
        report["comparison"] = compare(records, args.baseline, args.variant)
    output = (
        json.dumps(report, ensure_ascii=False, indent=2) + "\n"
        if args.format == "json"
        else render_markdown(report)
    )
    if args.output:
        Path(args.output).write_text(output, encoding="utf-8")
    else:
        print(output, end="")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())