"""Tests for the adaptive budget online policy controller and state round-trips."""

from __future__ import annotations

from evals.adaptive_budget_online_policy import (
    ACTION_INCREASE_BUDGET,
    ACTION_KEEP_BUDGET,
    STATE_SCHEMA_VERSION,
    AdaptiveBudgetPolicyController,
    BudgetPolicyConfig,
)
from llmdebug import _json as json


def test_budget_policy_on_updates_and_state_roundtrip(tmp_path) -> None:
    state_path = tmp_path / "budget_policy_state.json"
    events_path = tmp_path / "budget_policy_events.jsonl"
    config = BudgetPolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        alpha=0.2,
        epsilon=0.0,
        l2=1.0,
        warmup_decisions=0,
        save_every=1,
        state_path=state_path,
        events_path=events_path,
        random_seed=7,
    )
    controller = AdaptiveBudgetPolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_KEEP_BUDGET)
    event = controller.record_outcome(
        decision=decision,
        run_id="run_budget_online",
        case_id="case_a",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="tests_still_failing",
        reward_scope="budget_credit_v2",
        reward=0.6,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.5,
        reward_token_penalty=0.2,
        reward_latency_penalty=0.1,
        reward_stagnation_penalty=0.0,
        reward_budget_exceeded_penalty=0.0,
        request_budget_exceeded=False,
        success=False,
        local_llm_call_count=2,
        local_tokens_prompt=200,
        local_tokens_completion=121,
        local_tokens_total=321,
        local_wall_sec=2.4,
        selected_multiplier=1.5,
        applied_multiplier=1.5,
        default_multiplier=1.0,
        default_controls={"propose_budget_multiplier": 1.0},
        selected_controls={"propose_budget_multiplier": 1.5},
        applied_controls={"propose_budget_multiplier": 1.5},
        selected_controls_changed={"propose_budget": True},
        applied_controls_changed={"propose_budget": True},
        partial_override=False,
        partial_override_errors={},
        tier_index_before=0,
        tier_index_after=1,
        override_reason="",
    )
    state = controller.state
    controller.close()

    assert event.model_updated is True
    assert event.update_reason == "on_policy_update"
    assert event.reward_progress_score == 0.5
    assert event.applied_controls_changed["propose_budget"] is True
    assert state.decisions == 1
    assert state.updates == 1
    assert state.action_counts[ACTION_KEEP_BUDGET] == 1
    assert events_path.exists()
    assert len([line for line in events_path.read_text(encoding="utf-8").splitlines() if line]) == 1

    reloaded = AdaptiveBudgetPolicyController(config)
    reloaded_state = reloaded.state
    reloaded.close()
    assert reloaded.load_error is None
    assert reloaded_state.decisions == 1
    assert reloaded_state.updates == 1


def test_budget_policy_shadow_logs_but_does_not_update(tmp_path) -> None:
    state_path = tmp_path / "budget_policy_state_shadow.json"
    events_path = tmp_path / "budget_policy_events_shadow.jsonl"
    config = BudgetPolicyConfig.normalize(
        mode="shadow",
        algorithm="linucb",
        alpha=0.0,
        epsilon=0.0,
        l2=1.0,
        warmup_decisions=0,
        save_every=1,
        state_path=state_path,
        events_path=events_path,
        random_seed=0,
    )
    controller = AdaptiveBudgetPolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_KEEP_BUDGET)
    event = controller.record_outcome(
        decision=decision,
        run_id="run_budget_shadow",
        case_id="case_b",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="no_patch_proposed",
        reward_scope="budget_credit_v2",
        reward=-0.2,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.0,
        reward_token_penalty=0.1,
        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=1,
        local_tokens_prompt=8,
        local_tokens_completion=2,
        local_tokens_total=10,
        local_wall_sec=0.3,
        selected_multiplier=1.5,
        applied_multiplier=1.0,
        default_multiplier=1.0,
        default_controls={"propose_budget_multiplier": 1.0},
        selected_controls={"propose_budget_multiplier": 1.5},
        applied_controls={"propose_budget_multiplier": 1.0},
        selected_controls_changed={"propose_budget": True},
        applied_controls_changed={"propose_budget": False},
        partial_override=False,
        partial_override_errors={},
        tier_index_before=0,
        tier_index_after=0,
        override_reason="shadow_mode_baseline",
    )
    state = controller.state
    controller.close()

    assert decision.applied_action == ACTION_KEEP_BUDGET
    assert event.model_updated is False
    assert event.update_reason == "shadow_no_update"
    assert state.decisions == 1
    assert state.updates == 0
    assert len([line for line in events_path.read_text(encoding="utf-8").splitlines() if line]) == 1


def test_budget_policy_roundtrip_without_decisions(tmp_path) -> None:
    state_path = tmp_path / "budget_policy_state_empty.json"
    config = BudgetPolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "budget_policy_events_empty.jsonl",
    )
    controller = AdaptiveBudgetPolicyController(config)
    controller.close()

    assert state_path.exists()
    reloaded = AdaptiveBudgetPolicyController(config)
    state = reloaded.state
    reloaded.close()

    assert reloaded.load_error is None
    assert state.decisions == 0
    assert state.updates == 0
    assert state.feature_names == []


def test_budget_policy_reports_load_error_for_invalid_state(tmp_path) -> None:
    state_path = tmp_path / "budget_invalid_state.json"
    state_path.write_text(
        json.dumps({"schema_version": STATE_SCHEMA_VERSION, "algorithm": "linucb"}),
        encoding="utf-8",
    )
    config = BudgetPolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "events.jsonl",
    )
    controller = AdaptiveBudgetPolicyController(config)
    state = controller.state
    controller.close()

    assert controller.load_error is not None
    assert state.feature_names == []
    assert state.decisions == 0


def test_budget_policy_rejects_schema_v1_state(tmp_path) -> None:
    state_path = tmp_path / "budget_schema_v1_state.json"
    state_path.write_text(
        json.dumps({"schema_version": 1, "algorithm": "linucb"}),
        encoding="utf-8",
    )
    config = BudgetPolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "events.jsonl",
    )
    controller = AdaptiveBudgetPolicyController(config)
    state = controller.state
    controller.close()

    assert controller.load_error is not None
    assert "unsupported schema version" in str(controller.load_error)
    assert state.feature_names == []
    assert state.decisions == 0


def test_budget_policy_expands_loaded_state_for_new_feature_dimensions(tmp_path) -> None:
    state_path = tmp_path / "budget_policy_state_expand.json"
    events_path = tmp_path / "budget_policy_events_expand.jsonl"
    state_path.write_text(
        json.dumps(
            {
                "schema_version": STATE_SCHEMA_VERSION,
                "mode": "on",
                "algorithm": "linucb",
                "alpha": 0.2,
                "epsilon": 0.0,
                "l2": 1.5,
                "warmup_decisions": 0,
                "save_every": 1,
                "feature_names": ["bias"],
                "decisions": 3,
                "updates": 2,
                "action_counts": {
                    ACTION_KEEP_BUDGET: 1,
                    ACTION_INCREASE_BUDGET: 2,
                },
                "A": {
                    ACTION_KEEP_BUDGET: [[3.0]],
                    ACTION_INCREASE_BUDGET: [[4.0]],
                },
                "b": {
                    ACTION_KEEP_BUDGET: [0.7],
                    ACTION_INCREASE_BUDGET: [0.9],
                },
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    config = BudgetPolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        alpha=0.2,
        epsilon=0.0,
        l2=1.5,
        warmup_decisions=0,
        save_every=1,
        state_path=state_path,
        events_path=events_path,
        random_seed=0,
    )
    controller = AdaptiveBudgetPolicyController(config)
    controller.decide(
        feature_values={"bias": 1.0, "attempt_frac": 0.5},
        default_action=ACTION_KEEP_BUDGET,
    )
    state = controller.state
    controller.close()

    assert state.feature_names == ["attempt_frac", "bias"]
    for action in (ACTION_KEEP_BUDGET, ACTION_INCREASE_BUDGET):
        assert state.A[action].shape == (2, 2)
        assert state.b[action].shape == (2,)

    bias_idx = state.feature_names.index("bias")
    new_feature_idx = state.feature_names.index("attempt_frac")
    assert state.A[ACTION_KEEP_BUDGET][bias_idx, bias_idx] == 3.0
    assert state.A[ACTION_INCREASE_BUDGET][bias_idx, bias_idx] == 4.0
    assert state.b[ACTION_KEEP_BUDGET][bias_idx] == 0.7
    assert state.b[ACTION_INCREASE_BUDGET][bias_idx] == 0.9
    assert state.A[ACTION_KEEP_BUDGET][new_feature_idx, new_feature_idx] == 1.5
    assert state.A[ACTION_INCREASE_BUDGET][new_feature_idx, new_feature_idx] == 1.5
