from __future__ import annotations

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


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=0.7,
        reward_success=0.0,
        reward_token_penalty=0.2,
        reward_latency_penalty=0.1,
        success=False,
        tokens_total=123,
        attempt_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 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=-0.1,
        reward_success=0.0,
        reward_token_penalty=0.1,
        reward_latency_penalty=0.0,
        success=False,
        tokens_total=10,
        attempt_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_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_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(
            {
                "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_SKIP_HELPER: 1,
                    ACTION_INVOKE_HELPER: 2,
                },
                "A": {
                    ACTION_SKIP_HELPER: [[3.0]],
                    ACTION_INVOKE_HELPER: [[4.0]],
                },
                "b": {
                    ACTION_SKIP_HELPER: [0.7],
                    ACTION_INVOKE_HELPER: [0.9],
                },
            },
            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 (ACTION_SKIP_HELPER, ACTION_INVOKE_HELPER):
        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
