from __future__ import annotations

import csv

from evals.export_policy_dataset import main as export_policy_dataset_main
from llmdebug import _json as json


def test_export_policy_dataset_writes_flat_csv_and_dedupes(tmp_path) -> None:
    events_path = tmp_path / "events.jsonl"
    output_csv = tmp_path / "dataset.csv"
    events = [
        {
            "schema_version": 1,
            "run_id": "run_a",
            "case_id": "case_1",
            "condition": "adaptive_llm_discretion",
            "attempt": 1,
            "note": "no_patch_proposed",
            "decision_id": "d1",
            "decision_index": 1,
            "mode": "on",
            "algorithm": "linucb",
            "selected_action": "invoke_helper",
            "applied_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.8,
            "decision_strategy": "linucb_greedy",
            "scores": {"skip_helper": 0.2, "invoke_helper": 0.5},
            "feature_values": {"bias": 1.0, "attempt_frac": 0.5},
            "reward": 0.2,
            "reward_success": 0.0,
            "reward_token_penalty": 0.5,
            "reward_latency_penalty": 0.3,
            "success": False,
            "tokens_total": 321,
            "attempt_wall_sec": 4.2,
            "override_reason": "",
            "model_updated": True,
            "update_reason": "on_policy_update",
            "created_at_utc": "2026-02-26T00:00:00Z",
        },
        {
            "schema_version": 1,
            "run_id": "run_a",
            "case_id": "case_1",
            "condition": "adaptive_llm_discretion",
            "attempt": 1,
            "note": "no_patch_proposed",
            "decision_id": "d1",
            "decision_index": 1,
            "mode": "on",
            "algorithm": "linucb",
            "selected_action": "invoke_helper",
            "applied_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.8,
            "decision_strategy": "linucb_greedy",
            "scores": {"skip_helper": 0.2, "invoke_helper": 0.5},
            "feature_values": {"bias": 1.0, "attempt_frac": 0.5},
            "reward": 0.9,
            "reward_success": 1.0,
            "reward_token_penalty": 0.05,
            "reward_latency_penalty": 0.05,
            "success": True,
            "tokens_total": 100,
            "attempt_wall_sec": 1.0,
            "override_reason": "",
            "model_updated": True,
            "update_reason": "on_policy_update",
            "created_at_utc": "2026-02-26T00:00:01Z",
        },
        {
            "schema_version": 1,
            "run_id": "run_a",
            "case_id": "case_2",
            "condition": "adaptive_llm_discretion",
            "attempt": 1,
            "note": "tests_still_failing",
            "decision_id": "d2",
            "decision_index": 2,
            "mode": "shadow",
            "algorithm": "linucb",
            "selected_action": "skip_helper",
            "applied_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.0,
            "decision_strategy": "epsilon_uniform",
            "scores": {"skip_helper": 0.4, "invoke_helper": 0.3},
            "feature_values": {"bias": 1.0, "attempt_frac": 1.0},
            "reward": -0.2,
            "reward_success": 0.0,
            "reward_token_penalty": 0.1,
            "reward_latency_penalty": 0.1,
            "success": False,
            "tokens_total": 50,
            "attempt_wall_sec": 2.0,
            "override_reason": "shadow_mode_baseline",
            "model_updated": False,
            "update_reason": "shadow_no_update",
            "created_at_utc": "2026-02-26T00:00:02Z",
        },
    ]
    events_path.write_text(
        "\n".join(json.dumps(event, ensure_ascii=False) for event in events) + "\n",
        encoding="utf-8",
    )

    exit_code = export_policy_dataset_main(
        [
            "--events",
            str(events_path),
            "--output-csv",
            str(output_csv),
            "--min-propensity",
            "0.1",
        ]
    )

    assert exit_code == 0
    with output_csv.open("r", encoding="utf-8", newline="") as handle:
        rows = list(csv.DictReader(handle))

    assert len(rows) == 2
    by_decision = {row["decision_id"]: row for row in rows}
    assert by_decision["d1"]["reward"] == "0.9"
    assert by_decision["d1"]["feature_bias"] == "1.0"
    assert by_decision["d1"]["feature_attempt_frac"] == "0.5"
    assert by_decision["d1"]["importance_weight"] == "1.25"
    assert by_decision["d2"]["propensity_clipped"] == "0.1"
    assert by_decision["d2"]["importance_weight"] == "10.0"


def test_export_policy_dataset_strict_input_fails_fast(tmp_path) -> None:
    events_path = tmp_path / "events_invalid.jsonl"
    events_path.write_text('{"run_id":"r1"}\n{invalid-json\n', encoding="utf-8")

    exit_code = export_policy_dataset_main(
        [
            "--events",
            str(events_path),
            "--strict-input",
        ]
    )

    assert exit_code == 2
