"""W20E5 — dispatch-side image_call_budget gate + worker install (RED).

Two surfaces under test:

1. ``evaluate_image_call_budget_gate`` — pure decision used at dispatch
   time. When the W20 shot-aware image path is active
   (``settings.background_render_reference_mode == "shot_aware_plan"``)
   and ``category in {"image", "all"}``, it must fail closed unless the
   caller explicitly approved image generation AND provided a positive
   cap. Otherwise it returns the budget that the worker should install
   (or None for non-image categories / opt-out).

2. ``run_steps_batch`` — worker. When given a budget, it must install it
   for the duration of the run and uninstall it afterwards so the budget
   does not leak across jobs.
"""
from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from app.core.errors import AppError
from app.core.image_call_budget import (
    ImageCallBudget,
    get_current_budget,
    uninstall_budget,
)


@pytest.fixture(autouse=True)
def _isolate_budget():
    uninstall_budget()
    yield
    uninstall_budget()


# ─────────────────────────────────────────────────────────────────────────────
# evaluate_image_call_budget_gate — pure decision
# ─────────────────────────────────────────────────────────────────────────────


def _patch_render_mode(monkeypatch, value: str) -> None:
    from app.core.config import settings as _s
    monkeypatch.setattr(_s, "background_render_reference_mode", value, raising=False)


def test_gate_returns_none_for_analysis_category(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    out = evaluate_image_call_budget_gate(
        category="analysis", image_call_cap=None, approve_image_generation=False,
    )
    assert out is None


def test_gate_returns_none_for_legacy_image_without_cap(monkeypatch):
    """category=image on legacy mode without explicit cap → unchanged, no budget."""
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "legacy")
    out = evaluate_image_call_budget_gate(
        category="image", image_call_cap=None, approve_image_generation=False,
    )
    assert out is None


def test_gate_returns_budget_when_caller_opts_in_on_legacy(monkeypatch):
    """If caller passes a cap on legacy mode, install it (advisory opt-in)."""
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "legacy")
    out = evaluate_image_call_budget_gate(
        category="image", image_call_cap=3, approve_image_generation=False,
    )
    assert isinstance(out, ImageCallBudget)
    assert out.snapshot()["cap"] == 3


def test_gate_fails_closed_when_shot_aware_image_without_approval(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    with pytest.raises(AppError) as ei:
        evaluate_image_call_budget_gate(
            category="image",
            image_call_cap=5,
            approve_image_generation=False,
        )
    assert ei.value.status_code == 400
    assert ei.value.code == "step.image_generation_not_approved"


def test_gate_fails_closed_when_shot_aware_image_without_cap(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    with pytest.raises(AppError) as ei:
        evaluate_image_call_budget_gate(
            category="image",
            image_call_cap=None,
            approve_image_generation=True,
        )
    assert ei.value.status_code == 400
    assert ei.value.code == "step.image_call_cap_required"


def test_gate_fails_closed_when_shot_aware_image_with_zero_cap(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    with pytest.raises(AppError) as ei:
        evaluate_image_call_budget_gate(
            category="image",
            image_call_cap=0,
            approve_image_generation=True,
        )
    assert ei.value.code == "step.image_call_cap_required"


def test_gate_returns_budget_when_shot_aware_image_fully_approved(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    out = evaluate_image_call_budget_gate(
        category="image",
        image_call_cap=7,
        approve_image_generation=True,
    )
    assert isinstance(out, ImageCallBudget)
    assert out.snapshot()["cap"] == 7


def test_gate_applies_to_category_all_under_shot_aware_plan(monkeypatch):
    from app.services.analysis_dispatch_service import evaluate_image_call_budget_gate
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    with pytest.raises(AppError):
        evaluate_image_call_budget_gate(
            category="all", image_call_cap=None, approve_image_generation=False,
        )


# ─────────────────────────────────────────────────────────────────────────────
# run_steps_batch — budget install/uninstall lifecycle
# ─────────────────────────────────────────────────────────────────────────────


def test_run_steps_batch_installs_and_uninstalls_budget(monkeypatch):
    """Worker must install the budget at entry and clear it at exit."""
    from app.services import analysis_dispatch_service as ads

    seen_during_run: dict = {}

    class _FakeStepRunner:
        def run(self, mode):
            seen_during_run["budget"] = get_current_budget()
            return {"status": "done"}

    monkeypatch.setattr(
        ads, "get_step_runner",
        lambda *_a, **_kw: _FakeStepRunner(),
    )
    monkeypatch.setattr(ads, "orchestrate_full_sync", lambda *a, **kw: None)
    monkeypatch.setattr(
        ads, "SessionLocal",
        lambda: MagicMock(close=lambda: None, rollback=lambda: None),
    )
    # step_manifest lookup — yield empty dict (no pre-sync).
    monkeypatch.setattr(
        "app.core.step_manifest.get_manifest_dict",
        lambda _sid: {},
    )
    # _recover_episode_status_on_failure — no-op for this test path.
    monkeypatch.setattr(
        ads, "_recover_episode_status_on_failure",
        lambda *a, **kw: None,
    )

    budget = ImageCallBudget(cap=3)
    assert get_current_budget() is None

    ads.run_steps_batch(
        project_id="p1",
        episode_id="e1",
        step_ids=["fake_step"],
        run_mode="resume",
        project_config={},
        opik_context={},
        budget=budget,
    )

    # During the run, the budget was visible.
    assert seen_during_run["budget"] is budget
    # After return, it has been cleared.
    assert get_current_budget() is None


def test_run_steps_batch_without_budget_leaves_global_state_clean(monkeypatch):
    from app.services import analysis_dispatch_service as ads

    class _FakeStepRunner:
        def run(self, mode):
            return {"status": "done"}

    monkeypatch.setattr(ads, "get_step_runner", lambda *_a, **_kw: _FakeStepRunner())
    monkeypatch.setattr(ads, "orchestrate_full_sync", lambda *a, **kw: None)
    monkeypatch.setattr(
        ads, "SessionLocal",
        lambda: MagicMock(close=lambda: None, rollback=lambda: None),
    )
    monkeypatch.setattr(
        "app.core.step_manifest.get_manifest_dict", lambda _sid: {},
    )
    monkeypatch.setattr(
        ads, "_recover_episode_status_on_failure", lambda *a, **kw: None,
    )

    assert get_current_budget() is None
    ads.run_steps_batch(
        project_id="p1", episode_id="e1", step_ids=["x"],
        run_mode="resume", project_config={}, opik_context={},
    )
    assert get_current_budget() is None


def test_run_steps_batch_uninstalls_budget_even_on_failure(monkeypatch):
    """Even if the worker crashes mid-batch, the budget must be cleared."""
    from app.services import analysis_dispatch_service as ads

    class _CrashRunner:
        def run(self, mode):
            raise RuntimeError("boom")

    monkeypatch.setattr(ads, "get_step_runner", lambda *_a, **_kw: _CrashRunner())
    monkeypatch.setattr(ads, "orchestrate_full_sync", lambda *a, **kw: None)
    monkeypatch.setattr(
        ads, "SessionLocal",
        lambda: MagicMock(close=lambda: None, rollback=lambda: None),
    )
    monkeypatch.setattr(
        "app.core.step_manifest.get_manifest_dict", lambda _sid: {},
    )
    monkeypatch.setattr(
        ads, "_recover_episode_status_on_failure", lambda *a, **kw: None,
    )

    budget = ImageCallBudget(cap=1)
    ads.run_steps_batch(
        project_id="p1", episode_id="e1", step_ids=["crashy"],
        run_mode="resume", project_config={}, opik_context={},
        budget=budget,
    )
    assert get_current_budget() is None


# ─────────────────────────────────────────────────────────────────────────────
# dispatch_category_run — end-to-end gate wiring (no-network)
# ─────────────────────────────────────────────────────────────────────────────


def _stub_dispatch_dependencies(monkeypatch):
    """Stub everything dispatch touches so we can assert routing only."""
    from app.services import analysis_dispatch_service as ads
    monkeypatch.setattr(
        "app.core.task_registry.is_task_running", lambda _k: False,
    )
    monkeypatch.setattr(
        ads, "preflight_analysis_start", lambda *_a, **_kw: None,
    )
    monkeypatch.setattr(
        ads, "load_project_llm_config", lambda *_a, **_kw: {},
    )
    monkeypatch.setattr(
        ads, "select_steps_for_category", lambda *_a, **_kw: ["fake_step"],
    )
    monkeypatch.setattr(
        ads, "build_opik_context", lambda *_a, **_kw: {},
    )
    monkeypatch.setattr(
        ads, "rollback_episode_status", lambda *a, **kw: None,
    )


def test_dispatch_image_under_shot_aware_without_approval_raises_before_submit(
    monkeypatch,
):
    from app.services import analysis_dispatch_service as ads
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    _stub_dispatch_dependencies(monkeypatch)

    submitted = {"called": False}

    def _spy_submit(**_kw):
        submitted["called"] = True
        return True

    monkeypatch.setattr(ads, "submit_background_job", _spy_submit)

    db = MagicMock()
    with pytest.raises(AppError) as ei:
        ads.dispatch_category_run(
            project_id="p1", episode_id="e1",
            category="image", mode="resume", db=db,
            image_call_cap=None, approve_image_generation=False,
        )
    assert ei.value.code == "step.image_generation_not_approved"
    assert submitted["called"] is False


def test_dispatch_image_under_shot_aware_with_approval_submits_budget(monkeypatch):
    from app.services import analysis_dispatch_service as ads
    _patch_render_mode(monkeypatch, "shot_aware_plan")
    _stub_dispatch_dependencies(monkeypatch)

    captured: dict = {}

    def _spy_submit(*, job_key, target, args, description):
        captured["args"] = args
        captured["target"] = target
        captured["job_key"] = job_key
        return True

    monkeypatch.setattr(ads, "submit_background_job", _spy_submit)

    db = MagicMock()
    out = ads.dispatch_category_run(
        project_id="p1", episode_id="e1",
        category="image", mode="resume", db=db,
        image_call_cap=4, approve_image_generation=True,
    )
    assert out["ok"] is True
    # Budget must be in the args tuple passed to run_steps_batch.
    assert any(isinstance(a, ImageCallBudget) for a in captured["args"]), (
        "submit_background_job args must carry the ImageCallBudget instance"
    )
    budget_arg = next(a for a in captured["args"] if isinstance(a, ImageCallBudget))
    assert budget_arg.snapshot()["cap"] == 4


def test_dispatch_image_under_legacy_no_cap_passes_through_with_no_budget(monkeypatch):
    from app.services import analysis_dispatch_service as ads
    _patch_render_mode(monkeypatch, "legacy")
    _stub_dispatch_dependencies(monkeypatch)

    captured: dict = {}

    def _spy_submit(*, job_key, target, args, description):
        captured["args"] = args
        return True

    monkeypatch.setattr(ads, "submit_background_job", _spy_submit)

    db = MagicMock()
    ads.dispatch_category_run(
        project_id="p1", episode_id="e1",
        category="image", mode="resume", db=db,
    )
    # No ImageCallBudget instance threaded through.
    assert not any(isinstance(a, ImageCallBudget) for a in captured["args"])