from __future__ import annotations

import json

import pytest

from evals.analyze_results import (
    _derive_case_condition_rows,
    _load_and_group_records,
    _summarize_bucket,
)
from evals.analyze_results import (
    main as analyze_main,
)


def test_analyze_grouping_and_uplift(tmp_path) -> None:
    records = [
        {
            "run_id": "run_a",
            "case_id": "case_one",
            "condition": "traceback_only",
            "category": "hidden_state",
            "difficulty": "easy",
            "snapshot_dependency": "required",
            "attempt": 1,
            "success": False,
            "before_duration_sec": 1.0,
            "after_duration_sec": 0.5,
            "cumulative_runtime_sec": 1.6,
            "patch_meta": {"usage": {"total_tokens": 10}},
        },
        {
            "run_id": "run_a",
            "case_id": "case_one",
            "condition": "traceback_only",
            "category": "hidden_state",
            "difficulty": "easy",
            "snapshot_dependency": "required",
            "attempt": 2,
            "success": True,
            "before_duration_sec": 0.5,
            "after_duration_sec": 0.2,
            "cumulative_runtime_sec": 2.1,
            "patch_meta": {"usage": {"total_tokens": 20}},
        },
        {
            "run_id": "run_a",
            "case_id": "case_one",
            "condition": "with_snapshot",
            "category": "hidden_state",
            "difficulty": "easy",
            "snapshot_dependency": "required",
            "attempt": 1,
            "success": True,
            "before_duration_sec": 0.9,
            "after_duration_sec": 0.2,
            "cumulative_runtime_sec": 1.4,
            "patch_meta": {"usage": {"total_tokens": 15}},
        },
        {
            "run_id": "run_a",
            "case_id": "case_two",
            "condition": "traceback_only",
            "category": "parse_value",
            "difficulty": "medium",
            "snapshot_dependency": "helpful",
            "attempt": 1,
            "success": False,
            "before_duration_sec": 1.1,
            "cumulative_runtime_sec": 1.3,
        },
        {
            "run_id": "run_a",
            "case_id": "case_two",
            "condition": "with_snapshot",
            "category": "parse_value",
            "difficulty": "medium",
            "snapshot_dependency": "helpful",
            "attempt": 1,
            "success": False,
            "before_duration_sec": 1.0,
            "cumulative_runtime_sec": 1.5,
            "patch_meta": {"usage": {"total_tokens": 5}},
        },
        {
            "run_id": "run_b",
            "case_id": "case_one",
            "condition": "traceback_only",
            "category": "hidden_state",
            "difficulty": "easy",
            "snapshot_dependency": "required",
            "attempt": 1,
            "success": False,
            "before_duration_sec": 1.3,
            "cumulative_runtime_sec": 1.8,
        },
        {
            "run_id": "run_b",
            "case_id": "case_one",
            "condition": "with_snapshot",
            "category": "hidden_state",
            "difficulty": "easy",
            "snapshot_dependency": "required",
            "attempt": 1,
            "success": True,
            "before_duration_sec": 1.2,
            "after_duration_sec": 0.1,
            "cumulative_runtime_sec": 1.6,
            "patch_meta": {"usage": {"total_tokens": 12}},
        },
    ]
    jsonl = tmp_path / "results.jsonl"
    jsonl.write_text("\n".join(json.dumps(record) for record in records) + "\n", encoding="utf-8")

    grouped, input_errors = _load_and_group_records([str(jsonl)])
    assert len(grouped) == 6
    assert input_errors == {"files_unreadable": 0, "json_lines_invalid": 0, "records_skipped": 0}

    rows = _derive_case_condition_rows(grouped)
    assert len(rows) == 6
    row_index = {(row["run_id"], row["case_id"], row["condition"]): row for row in rows}
    assert row_index[("run_a", "case_one", "traceback_only")]["runtime_sec"] == pytest.approx(2.1)

    overall = _summarize_bucket(rows)
    assert overall["traceback_only"]["cases"] == 3
    assert overall["traceback_only"]["successes"] == 1
    assert overall["traceback_only"]["success_rate"] == pytest.approx(1 / 3)
    assert overall["traceback_only"]["median_attempts_to_success"] == pytest.approx(2.0)
    assert overall["with_snapshot"]["cases"] == 3
    assert overall["with_snapshot"]["successes"] == 2
    assert overall["with_snapshot"]["success_rate"] == pytest.approx(2 / 3)
    assert overall["uplift_success_rate"] == pytest.approx(1 / 3)

    by_category = {"hidden_state": [], "parse_value": []}
    for row in rows:
        by_category[row["category"]].append(row)
    hidden_state_summary = _summarize_bucket(by_category["hidden_state"])
    assert hidden_state_summary["with_snapshot"]["success_rate"] == pytest.approx(1.0)
    assert hidden_state_summary["traceback_only"]["success_rate"] == pytest.approx(0.5)
    assert hidden_state_summary["uplift_success_rate"] == pytest.approx(0.5)

    by_snapshot_dependency = {"required": [], "helpful": []}
    for row in rows:
        by_snapshot_dependency[row["snapshot_dependency"]].append(row)
    required_summary = _summarize_bucket(by_snapshot_dependency["required"])
    assert required_summary["traceback_only"]["cases"] == 2
    assert required_summary["with_snapshot"]["cases"] == 2
    assert required_summary["uplift_success_rate"] == pytest.approx(0.5)


def test_analyze_grouping_tolerates_malformed_jsonl(tmp_path, capsys) -> None:
    jsonl = tmp_path / "results_mixed.jsonl"
    jsonl.write_text(
        (
            '{"run_id":"run","case_id":"case","condition":"traceback_only","attempt":1,"success":false}\n'
            "this-is-not-json\n"
            '{"run_id":"run","case_id":"case","condition":"with_snapshot","attempt":1,"success":true}\n'
            '{"run_id":"run","condition":"with_snapshot","attempt":2}\n'
        ),
        encoding="utf-8",
    )

    grouped, input_errors = _load_and_group_records([str(jsonl)])
    assert len(grouped) == 2
    assert input_errors["json_lines_invalid"] == 1
    assert input_errors["records_skipped"] == 1

    exit_code = analyze_main([str(jsonl)])
    assert exit_code == 0
    output = capsys.readouterr().out
    assert "Input diagnostics:" in output
    assert "json_lines_invalid=1" in output
    assert "records_skipped=1" in output


def test_analyze_grouping_strict_input_fails_on_malformed_jsonl(tmp_path, capsys) -> None:
    jsonl = tmp_path / "results_bad.jsonl"
    jsonl.write_text(
        '{"run_id":"run","case_id":"case","condition":"traceback_only","attempt":1,"success":false}\n'
        "not-json\n",
        encoding="utf-8",
    )

    exit_code = analyze_main(["--strict-input", str(jsonl)])
    assert exit_code == 2
    output = capsys.readouterr().out
    assert "Input error:" in output
