"""W20E5 — location_floor_plan image_call_budget wiring (RED).

``generate_floor_plan_image`` calls ``openai_client.images.edit`` /
``openai_client.images.generate`` in a retry loop. Each attempt must reserve
one budget unit before the request; cap exhaustion raises
``ImageCallBudgetExceeded`` directly (not wrapped in RuntimeError).
"""
from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from app.core.image_call_budget import (
    ImageCallBudget,
    ImageCallBudgetExceeded,
    install_budget,
    uninstall_budget,
)
from app.modules.pipeline.location_floor_plan import generate_floor_plan_image


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


def _ok_response(b64: str = "aGVsbG8="):
    # The function decodes b64 and requires >=1024 bytes — use a long payload.
    import base64
    payload = base64.b64encode(b"x" * 2048).decode()
    resp = MagicMock()
    resp.data = [MagicMock(b64_json=payload)]
    return resp


def test_no_budget_installed_keeps_legacy_behaviour():
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    png = generate_floor_plan_image(prompt="x", openai_client=client, ref_paths=None)
    assert isinstance(png, bytes)
    assert client.images.generate.call_count == 1


def test_cap_zero_blocks_generate(monkeypatch):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    with pytest.raises(ImageCallBudgetExceeded):
        generate_floor_plan_image(prompt="x", openai_client=client, ref_paths=None)
    assert client.images.generate.call_count == 0


def test_cap_zero_blocks_edit(tmp_path):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.edit.return_value = _ok_response()
    ref = tmp_path / "r.png"
    ref.write_bytes(b"R" * 2048)  # > 1024 bytes
    with pytest.raises(ImageCallBudgetExceeded):
        generate_floor_plan_image(
            prompt="x", openai_client=client, ref_paths=[ref],
        )
    assert client.images.edit.call_count == 0


def test_cap_one_allows_one_call():
    budget = ImageCallBudget(cap=1)
    install_budget(budget)
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    png = generate_floor_plan_image(prompt="x", openai_client=client, ref_paths=None)
    assert isinstance(png, bytes)
    assert client.images.generate.call_count == 1
    assert budget.snapshot()["used"] == 1


def test_cap_exhausts_mid_retry_loop_without_swallowing(monkeypatch):
    """cap=2 + transient errors → 2 attempts then ImageCallBudgetExceeded raised
    directly (NOT wrapped in RuntimeError)."""
    import app.modules.pipeline.location_floor_plan as lfp
    # eliminate retry sleeps so the test is fast.
    monkeypatch.setattr(lfp.time, "sleep", lambda _s: None)

    budget = ImageCallBudget(cap=2)
    install_budget(budget)
    client = MagicMock()
    client.images.generate.side_effect = RuntimeError("upstream 503")

    with pytest.raises(ImageCallBudgetExceeded):
        generate_floor_plan_image(prompt="x", openai_client=client, ref_paths=None)
    assert client.images.generate.call_count == 2
    snap = budget.snapshot()
    assert snap["used"] == 2
    assert snap["denied"] == 1