"""Tests for the adaptive online policy controller (actions, state persistence)."""

from __future__ import annotations

import pytest

from evals.adaptive_online_policy import (
    ACTION_FRAME_ONLY,
    ACTION_INVOKE_HELPER,
    ACTION_SKIP_HELPER,
    ONLINE_POLICY_ACTIONS,
    STATE_SCHEMA_VERSION,
    AdaptiveOnlinePolicyController,
    OnlinePolicyConfig,
)
from llmdebug import _json as json


def _policy_state_payload(
    *,
    preferred_action: str | None = None,
    preferred_value: float = 2.0,
    action_counts: dict[str, int] | None = None,
    diagonal_by_action: dict[str, float] | None = None,
    b_by_action: dict[str, float] | None = None,
) -> dict[str, object]:
    counts = dict.fromkeys(ONLINE_POLICY_ACTIONS, 0)
    counts.update(action_counts or {})
    diagonals = dict.fromkeys(ONLINE_POLICY_ACTIONS, 1.0)
    diagonals.update(diagonal_by_action or {})
    b_values = dict.fromkeys(ONLINE_POLICY_ACTIONS, 0.0)
    b_values.update(b_by_action or {})
    if preferred_action in ONLINE_POLICY_ACTIONS:
        b_values[str(preferred_action)] = float(preferred_value)
    return {
        "schema_version": STATE_SCHEMA_VERSION,
        "mode": "on",
        "algorithm": "linucb",
        "alpha": 0.0,
        "epsilon": 0.0,
        "l2": 1.0,
        "warmup_decisions": 0,
        "save_every": 1,
        "feature_names": ["bias"],
        "decisions": 0,
        "updates": 0,
        "action_counts": counts,
        "A": {action: [[float(diagonals[action])]] for action in ONLINE_POLICY_ACTIONS},
        "b": {action: [float(b_values[action])] for action in ONLINE_POLICY_ACTIONS},
    }


def test_online_policy_on_updates_and_state_roundtrip(tmp_path) -> None:
    state_path = tmp_path / "policy_state.json"
    events_path = tmp_path / "policy_events.jsonl"
    config = OnlinePolicyConfig.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 = AdaptiveOnlinePolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_INVOKE_HELPER)
    event = controller.record_outcome(
        decision=decision,
        run_id="run_online",
        case_id="case_a",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="tests_still_failing",
        reward_scope="helper_credit_v2",
        reward=0.7,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.4,
        reward_token_penalty=0.2,
        reward_latency_penalty=0.1,
        reward_stagnation_penalty=0.0,
        reward_helper_materialization_bonus=0.0,
        success=False,
        local_llm_call_count=1,
        local_tokens_prompt=100,
        local_tokens_completion=23,
        local_tokens_total=123,
        local_wall_sec=1.2,
        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.4
    assert state.decisions == 1
    assert state.updates == 1
    assert state.action_counts[ACTION_INVOKE_HELPER] == 1
    assert events_path.exists()
    assert len([line for line in events_path.read_text(encoding="utf-8").splitlines() if line]) == 1

    reloaded = AdaptiveOnlinePolicyController(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_online_policy_shadow_logs_but_does_not_update(tmp_path) -> None:
    state_path = tmp_path / "policy_state_shadow.json"
    events_path = tmp_path / "policy_events_shadow.jsonl"
    config = OnlinePolicyConfig.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 = AdaptiveOnlinePolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_SKIP_HELPER)
    event = controller.record_outcome(
        decision=decision,
        run_id="run_shadow",
        case_id="case_b",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="no_patch_proposed",
        reward_scope="helper_credit_v2",
        reward=-0.1,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.0,
        reward_token_penalty=0.1,
        reward_latency_penalty=0.0,
        reward_stagnation_penalty=0.0,
        reward_helper_materialization_bonus=0.0,
        success=False,
        local_llm_call_count=0,
        local_tokens_prompt=0,
        local_tokens_completion=0,
        local_tokens_total=10,
        local_wall_sec=0.2,
        override_reason="shadow_mode_baseline",
    )
    state = controller.state
    controller.close()

    assert decision.applied_action == ACTION_SKIP_HELPER
    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_online_policy_empty_state_roundtrips_without_load_error(tmp_path) -> None:
    state_path = tmp_path / "policy_state_empty.json"
    events_path = tmp_path / "policy_events_empty.jsonl"
    config = OnlinePolicyConfig.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=0,
    )
    controller = AdaptiveOnlinePolicyController(config)
    controller.close()

    reloaded = AdaptiveOnlinePolicyController(config)
    state = reloaded.state
    reloaded.close()

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


def test_online_policy_record_outcome_uses_executed_behavior_action(tmp_path) -> None:
    state_path = tmp_path / "policy_state_behavior_override.json"
    events_path = tmp_path / "policy_events_behavior_override.jsonl"
    state_path.write_text(
        json.dumps(
            _policy_state_payload(
                preferred_action=ACTION_SKIP_HELPER,
                preferred_value=2.0,
            ),
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.normalize(
        mode="on",
        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 = AdaptiveOnlinePolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_INVOKE_HELPER)

    assert decision.applied_action == ACTION_SKIP_HELPER

    event = controller.record_outcome(
        decision=decision,
        run_id="run_behavior_override",
        case_id="case_override",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="stagnation_override_force_helper",
        reward_scope="helper_credit_v2",
        reward=1.0,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.6,
        reward_token_penalty=0.0,
        reward_latency_penalty=0.0,
        reward_stagnation_penalty=0.0,
        reward_helper_materialization_bonus=0.0,
        success=False,
        local_llm_call_count=1,
        local_tokens_prompt=4,
        local_tokens_completion=6,
        local_tokens_total=10,
        local_wall_sec=0.1,
        override_reason="constraint_stagnation_override_force_helper",
        behavior_action=ACTION_INVOKE_HELPER,
        behavior_propensity=1.0,
    )
    state = controller.state
    controller.close()

    assert event.model_updated is True
    assert event.applied_action == ACTION_SKIP_HELPER
    assert event.behavior_action == ACTION_INVOKE_HELPER
    assert event.behavior_propensity == 1.0
    assert state.b[ACTION_SKIP_HELPER][0] == 2.0
    assert state.b[ACTION_INVOKE_HELPER][0] == 1.0


def test_online_policy_record_outcome_updates_non_binary_behavior_action(tmp_path) -> None:
    state_path = tmp_path / "policy_state_frame_only.json"
    events_path = tmp_path / "policy_events_frame_only.jsonl"
    state_path.write_text(
        json.dumps(
            _policy_state_payload(
                preferred_action=ACTION_FRAME_ONLY,
                preferred_value=2.0,
            ),
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.normalize(
        mode="on",
        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 = AdaptiveOnlinePolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_INVOKE_HELPER)

    assert decision.applied_action == ACTION_FRAME_ONLY

    event = controller.record_outcome(
        decision=decision,
        run_id="run_frame_only",
        case_id="case_frame_only",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="tests_still_failing",
        reward_scope="helper_credit_v2",
        reward=0.3,
        reward_success_hint=0.0,
        reward_progress_hint=0.0,
        reward_progress_score=0.2,
        reward_token_penalty=0.0,
        reward_latency_penalty=0.0,
        reward_stagnation_penalty=0.0,
        reward_helper_materialization_bonus=0.1,
        success=False,
        local_llm_call_count=1,
        local_tokens_prompt=5,
        local_tokens_completion=3,
        local_tokens_total=8,
        local_wall_sec=0.1,
        override_reason="",
        behavior_action=ACTION_FRAME_ONLY,
        behavior_propensity=1.0,
    )
    state = controller.state
    controller.close()

    assert event.behavior_action == ACTION_FRAME_ONLY
    assert event.model_updated is True
    assert state.b[ACTION_FRAME_ONLY][0] == pytest.approx(2.3)
    assert state.b[ACTION_INVOKE_HELPER][0] == 0.0


def test_online_policy_record_outcome_can_log_without_updating_model(tmp_path) -> None:
    state_path = tmp_path / "policy_state_no_update.json"
    events_path = tmp_path / "policy_events_no_update.jsonl"
    config = OnlinePolicyConfig.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=0,
    )
    controller = AdaptiveOnlinePolicyController(config)
    decision = controller.decide(feature_values={"bias": 1.0}, default_action=ACTION_INVOKE_HELPER)

    event = controller.record_outcome(
        decision=decision,
        run_id="run_no_update",
        case_id="case_no_update",
        condition="adaptive_llm_discretion",
        attempt=1,
        note="patcher_request_failed",
        reward_scope="helper_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_helper_materialization_bonus=0.0,
        success=False,
        local_llm_call_count=1,
        local_tokens_prompt=10,
        local_tokens_completion=2,
        local_tokens_total=12,
        local_wall_sec=0.2,
        override_reason="",
        allow_model_update=False,
        update_reason_override="helper_transport_no_update",
    )
    state = controller.state
    controller.close()

    assert event.model_updated is False
    assert event.update_reason == "helper_transport_no_update"
    assert event.reward_progress_score == 0.0
    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_online_policy_reports_load_error_for_invalid_state(tmp_path) -> None:
    state_path = tmp_path / "invalid_state.json"
    state_path.write_text(
        json.dumps({"schema_version": STATE_SCHEMA_VERSION, "algorithm": "linucb"}),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "events.jsonl",
    )
    controller = AdaptiveOnlinePolicyController(config)
    state = controller.state
    controller.close()
    assert controller.load_error is not None
    assert state.feature_names == []
    assert state.decisions == 0


def test_online_policy_rejects_schema_v1_state(tmp_path) -> None:
    state_path = tmp_path / "schema_v1_state.json"
    state_path.write_text(
        json.dumps({"schema_version": 1, "algorithm": "linucb"}),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "events.jsonl",
    )
    controller = AdaptiveOnlinePolicyController(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_online_policy_schema_v2_state_cold_starts_empty(tmp_path) -> None:
    state_path = tmp_path / "schema_v2_state.json"
    state_path.write_text(
        json.dumps(
            {
                **_policy_state_payload(preferred_action=ACTION_INVOKE_HELPER),
                "schema_version": 2,
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.normalize(
        mode="on",
        algorithm="linucb",
        state_path=state_path,
        events_path=tmp_path / "events_v2.jsonl",
    )
    controller = AdaptiveOnlinePolicyController(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_online_policy_expands_loaded_state_for_new_feature_dimensions(tmp_path) -> None:
    state_path = tmp_path / "policy_state_expand.json"
    events_path = tmp_path / "policy_events_expand.jsonl"
    state_path.write_text(
        json.dumps(
            {
                **_policy_state_payload(
                    action_counts={ACTION_SKIP_HELPER: 1, ACTION_INVOKE_HELPER: 2},
                    diagonal_by_action={
                        ACTION_SKIP_HELPER: 3.0,
                        ACTION_INVOKE_HELPER: 4.0,
                    },
                    b_by_action={
                        ACTION_SKIP_HELPER: 0.7,
                        ACTION_INVOKE_HELPER: 0.9,
                    },
                ),
                "alpha": 0.2,
                "l2": 1.5,
                "decisions": 3,
                "updates": 2,
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    config = OnlinePolicyConfig.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 = AdaptiveOnlinePolicyController(config)
    controller.decide(
        feature_values={"bias": 1.0, "attempt_frac": 0.5},
        default_action=ACTION_INVOKE_HELPER,
    )
    state = controller.state
    controller.close()

    assert state.feature_names == ["attempt_frac", "bias"]
    for action in ONLINE_POLICY_ACTIONS:
        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_SKIP_HELPER][bias_idx, bias_idx] == 3.0
    assert state.A[ACTION_INVOKE_HELPER][bias_idx, bias_idx] == 4.0
    assert state.b[ACTION_SKIP_HELPER][bias_idx] == 0.7
    assert state.b[ACTION_INVOKE_HELPER][bias_idx] == 0.9
    assert state.A[ACTION_SKIP_HELPER][new_feature_idx, new_feature_idx] == 1.5
    assert state.A[ACTION_INVOKE_HELPER][new_feature_idx, new_feature_idx] == 1.5
