"""Tests for exporting policy events to a flat CSV dataset."""

from __future__ import annotations

import csv

from evals.export_policy_dataset import (
    POLICY_EVENT_SCHEMA_VERSION,
)
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": POLICY_EVENT_SCHEMA_VERSION,
            "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",
            "behavior_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.8,
            "behavior_propensity": 0.8,
            "decision_strategy": "linucb_greedy",
            "feature_scaling_profile": "v1",
            "scores": {"skip_helper": 0.2, "invoke_helper": 0.5},
            "feature_values": {"bias": 1.0, "attempt_frac": 0.5},
            "reward_scope": "helper_local_v1",
            "reward": 0.2,
            "reward_success_hint": 0.0,
            "reward_progress_hint": 0.1,
            "reward_token_penalty": 0.5,
            "reward_latency_penalty": 0.3,
            "reward_stagnation_penalty": 0.0,
            "reward_helper_materialization_bonus": 0.05,
            "success": False,
            "local_llm_call_count": 1,
            "local_tokens_prompt": 200,
            "local_tokens_completion": 121,
            "local_tokens_total": 321,
            "local_wall_sec": 4.2,
            "override_reason": "",
            "model_updated": True,
            "update_reason": "on_policy_update",
            "created_at_utc": "2026-02-26T00:00:00Z",
        },
        {
            "schema_version": POLICY_EVENT_SCHEMA_VERSION,
            "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",
            "behavior_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.8,
            "behavior_propensity": 0.8,
            "decision_strategy": "linucb_greedy",
            "feature_scaling_profile": "v1",
            "scores": {"skip_helper": 0.2, "invoke_helper": 0.5},
            "feature_values": {"bias": 1.0, "attempt_frac": 0.5},
            "reward_scope": "helper_local_v1",
            "reward": 0.9,
            "reward_success_hint": 0.25,
            "reward_progress_hint": 0.0,
            "reward_token_penalty": 0.05,
            "reward_latency_penalty": 0.05,
            "reward_stagnation_penalty": 0.1,
            "reward_helper_materialization_bonus": 0.0,
            "success": True,
            "local_llm_call_count": 1,
            "local_tokens_prompt": 70,
            "local_tokens_completion": 30,
            "local_tokens_total": 100,
            "local_wall_sec": 1.0,
            "override_reason": "",
            "model_updated": True,
            "update_reason": "on_policy_update",
            "created_at_utc": "2026-02-26T00:00:01Z",
        },
        {
            "schema_version": POLICY_EVENT_SCHEMA_VERSION,
            "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",
            "behavior_action": "invoke_helper",
            "default_action": "invoke_helper",
            "propensity": 0.0,
            "behavior_propensity": 1.0,
            "decision_strategy": "epsilon_uniform",
            "feature_scaling_profile": "v1",
            "scores": {"skip_helper": 0.4, "invoke_helper": 0.3},
            "feature_values": {"bias": 1.0, "attempt_frac": 1.0},
            "reward_scope": "helper_local_v1",
            "reward": -0.2,
            "reward_success_hint": 0.0,
            "reward_progress_hint": 0.0,
            "reward_token_penalty": 0.1,
            "reward_latency_penalty": 0.1,
            "reward_stagnation_penalty": 0.1,
            "reward_helper_materialization_bonus": 0.0,
            "success": False,
            "local_llm_call_count": 1,
            "local_tokens_prompt": 30,
            "local_tokens_completion": 20,
            "local_tokens_total": 50,
            "local_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"]["reward_scope"] == "helper_local_v1"
    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["d1"]["behavior_importance_weight"] == "1.25"
    assert by_decision["d1"]["feature_scaling_profile"] == "v1"
    assert by_decision["d1"]["reward_success_hint"] == "0.25"
    assert by_decision["d1"]["local_tokens_total"] == "100"
    assert by_decision["d1"]["local_wall_sec"] == "1.0"
    assert "tokens_total" not in by_decision["d1"]
    assert "attempt_wall_sec" not in by_decision["d1"]
    assert by_decision["d2"]["propensity_clipped"] == "0.1"
    assert by_decision["d2"]["importance_weight"] == "10.0"
    assert by_decision["d2"]["behavior_propensity"] == "1.0"
    assert by_decision["d2"]["behavior_importance_weight"] == "1.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


def test_export_policy_dataset_supports_budget_policy_events(tmp_path) -> None:
    events_path = tmp_path / "budget_events.jsonl"
    output_csv = tmp_path / "budget_dataset.csv"
    events = [
        {
            "schema_version": POLICY_EVENT_SCHEMA_VERSION,
            "run_id": "run_budget",
            "case_id": "case_budget_1",
            "condition": "adaptive_llm_discretion",
            "attempt": 1,
            "note": "no_patch_proposed",
            "decision_id": "b1",
            "decision_index": 1,
            "mode": "on",
            "algorithm": "linucb",
            "selected_action": "increase_budget",
            "applied_action": "increase_budget",
            "behavior_action": "increase_budget",
            "default_action": "keep_budget",
            "selected_multiplier": 1.5,
            "applied_multiplier": 1.5,
            "default_multiplier": 1.0,
            "tier_index_before": 0,
            "tier_index_after": 1,
            "propensity": 0.6,
            "behavior_propensity": 0.6,
            "decision_strategy": "linucb_greedy",
            "feature_scaling_profile": "v1",
            "scores": {"keep_budget": 0.1, "increase_budget": 0.4},
            "feature_values": {"bias": 1.0, "attempt_frac": 0.5},
            "reward_scope": "budget_local_v1",
            "reward": -0.3,
            "reward_success_hint": 0.0,
            "reward_progress_hint": 0.15,
            "reward_token_penalty": 0.2,
            "reward_latency_penalty": 0.1,
            "reward_stagnation_penalty": 0.0,
            "reward_budget_exceeded_penalty": 0.5,
            "request_budget_exceeded": True,
            "success": False,
            "local_llm_call_count": 2,
            "local_tokens_prompt": 80,
            "local_tokens_completion": 40,
            "local_tokens_total": 120,
            "local_wall_sec": 3.0,
            "override_reason": "",
            "model_updated": True,
            "update_reason": "on_policy_update",
            "created_at_utc": "2026-03-04T00:00:00Z",
        }
    ]
    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),
        ]
    )

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

    assert len(rows) == 1
    row = rows[0]
    assert row["score_keep_budget"] == "0.1"
    assert row["score_increase_budget"] == "0.4"
    assert row["request_budget_exceeded"] == "1"
    assert row["reward_budget_exceeded_penalty"] == "0.5"
    assert row["applied_multiplier"] == "1.5"
    assert row["behavior_propensity"] == "0.6"
    assert row["feature_scaling_profile"] == "v1"
    assert row["reward_progress_hint"] == "0.15"
    assert row["local_tokens_total"] == "120"


def test_export_policy_dataset_rejects_legacy_event_schema(tmp_path) -> None:
    events_path = tmp_path / "legacy_events.jsonl"
    output_csv = tmp_path / "legacy_dataset.csv"
    events = [
        {
            "schema_version": 1,
            "run_id": "run_legacy",
            "case_id": "case_legacy_1",
            "condition": "adaptive_llm_discretion",
            "attempt": 1,
            "decision_id": "legacy_1",
            "selected_action": "skip_helper",
            "applied_action": "invoke_helper",
            "propensity": 0.25,
            "reward": 0.1,
            "scores": {"skip_helper": 0.4, "invoke_helper": 0.6},
            "feature_values": {"bias": 1.0},
            "created_at_utc": "2026-03-05T00:00:00Z",
        }
    ]
    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),
        ]
    )

    assert exit_code == 1
    assert not output_csv.exists()
