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

Authoritative provider-side cap. Each ``openai_client.images.edit`` /
``openai_client.images.generate`` attempt must reserve one budget unit
*before* the call; cap exhaustion raises ``ImageCallBudgetExceeded`` and
the helper must not swallow it.

These tests use ``MagicMock`` openai clients — no network, no real
gpt-image-2 call. They run independently of the existing
``test_background_render.py`` suite (which deliberately uses no installed
budget and must keep passing).
"""
from __future__ import annotations

import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from unittest.mock import MagicMock

import pytest

from app.core.image_call_budget import (
    ImageCallBudget,
    ImageCallBudgetExceeded,
    bind_current_budget,
    install_budget,
    uninstall_budget,
)
from app.modules.pipeline.background_render import render_one_background


@pytest.fixture(autouse=True)
def _isolate_budget():
    """Ensure each test starts with no installed budget and cleans up."""
    uninstall_budget()
    yield
    uninstall_budget()


def _ok_response(b64: str = "aGVsbG8="):
    resp = MagicMock()
    resp.data = [MagicMock(b64_json=b64)]
    return resp


def test_no_budget_installed_keeps_legacy_behaviour(tmp_path):
    """Existing call paths must be unaffected when no budget is installed."""
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    out = tmp_path / "out.png"

    res = render_one_background(
        openai_client=client,
        image_model="gpt-image-2.5-sunburst",
        prompt="x",
        out_path=out,
        fp_path=None,
        prior_bg_paths=[],
        max_attempts=1,
    )
    assert res["status"] == "ok"
    assert client.images.generate.call_count == 1
    assert client.images.edit.call_count == 0


def test_cap_zero_blocks_provider_call_for_text_only(tmp_path):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    out = tmp_path / "out.png"

    with pytest.raises(ImageCallBudgetExceeded):
        render_one_background(
            openai_client=client,
            image_model="gpt-image-2.5-sunburst",
            prompt="x",
            out_path=out,
            fp_path=None,
            prior_bg_paths=[],
            max_attempts=1,
        )
    assert client.images.generate.call_count == 0
    assert client.images.edit.call_count == 0


def test_cap_zero_blocks_provider_call_for_single_edit(tmp_path):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.edit.return_value = _ok_response()
    fp = tmp_path / "fp.png"
    fp.write_bytes(b"FP")
    out = tmp_path / "out.png"

    with pytest.raises(ImageCallBudgetExceeded):
        render_one_background(
            openai_client=client,
            image_model="gpt-image-2.5-sunburst",
            prompt="x",
            out_path=out,
            fp_path=fp,
            prior_bg_paths=[],
            max_attempts=1,
        )
    assert client.images.edit.call_count == 0
    assert client.images.generate.call_count == 0


def test_cap_zero_blocks_provider_call_for_multi_edit(tmp_path):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.edit.return_value = _ok_response()
    fp = tmp_path / "fp.png"
    fp.write_bytes(b"FP")
    prior = tmp_path / "prior.png"
    prior.write_bytes(b"PR")
    out = tmp_path / "out.png"

    with pytest.raises(ImageCallBudgetExceeded):
        render_one_background(
            openai_client=client,
            image_model="gpt-image-2.5-sunburst",
            prompt="x",
            out_path=out,
            fp_path=fp,
            prior_bg_paths=[prior],
            max_attempts=1,
        )
    assert client.images.edit.call_count == 0
    assert client.images.generate.call_count == 0


def test_cap_one_allows_exactly_one_successful_generate(tmp_path):
    budget = ImageCallBudget(cap=1)
    install_budget(budget)
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    out = tmp_path / "out.png"

    res = render_one_background(
        openai_client=client,
        image_model="gpt-image-2.5-sunburst",
        prompt="x",
        out_path=out,
        fp_path=None,
        prior_bg_paths=[],
        max_attempts=4,
    )
    assert res["status"] == "ok"
    assert client.images.generate.call_count == 1
    snap = budget.snapshot()
    assert snap["used"] == 1
    assert snap["denied"] == 0


def test_cap_two_exhausts_on_third_non_moderation_retry(tmp_path):
    """cap=2 + max_attempts=4 + transient errors → 2 provider calls then raise."""
    budget = ImageCallBudget(cap=2)
    install_budget(budget)
    client = MagicMock()
    client.images.generate.side_effect = RuntimeError("upstream transient 502")
    out = tmp_path / "out.png"

    with pytest.raises(ImageCallBudgetExceeded):
        render_one_background(
            openai_client=client,
            image_model="gpt-image-2.5-sunburst",
            prompt="x",
            out_path=out,
            fp_path=None,
            prior_bg_paths=[],
            max_attempts=4,
        )
    assert client.images.generate.call_count == 2
    snap = budget.snapshot()
    assert snap["used"] == 2
    assert snap["denied"] == 1


def test_reserve_source_label_distinguishes_modes(tmp_path, monkeypatch):
    """Recorded source on denial must identify the call mode."""
    seen_sources: list[str] = []

    from app.core.image_call_budget import ImageCallBudget as _B

    real_reserve = _B.reserve

    def _spy(self, *, source: str) -> None:
        seen_sources.append(source)
        real_reserve(self, source=source)

    monkeypatch.setattr(_B, "reserve", _spy)

    budget = _B(cap=3)
    install_budget(budget)

    # mode 1: text-only
    out1 = tmp_path / "a.png"
    c1 = MagicMock()
    c1.images.generate.return_value = _ok_response()
    render_one_background(
        openai_client=c1, image_model="gpt-image-2.5-sunburst", prompt="x",
        out_path=out1, fp_path=None, prior_bg_paths=[], max_attempts=1,
    )

    # mode 2: single edit
    out2 = tmp_path / "b.png"
    fp = tmp_path / "fp.png"
    fp.write_bytes(b"FP")
    c2 = MagicMock()
    c2.images.edit.return_value = _ok_response()
    render_one_background(
        openai_client=c2, image_model="gpt-image-2.5-sunburst", prompt="x",
        out_path=out2, fp_path=fp, prior_bg_paths=[], max_attempts=1,
    )

    # mode 3: multi edit
    out3 = tmp_path / "c.png"
    prior = tmp_path / "prior.png"
    prior.write_bytes(b"PR")
    c3 = MagicMock()
    c3.images.edit.return_value = _ok_response()
    render_one_background(
        openai_client=c3, image_model="gpt-image-2.5-sunburst", prompt="x",
        out_path=out3, fp_path=fp, prior_bg_paths=[prior], max_attempts=1,
    )

    # All three sources must be distinct and prefixed with module name.
    assert len(seen_sources) == 3
    assert all(s.startswith("background_render.") for s in seen_sources)
    assert len(set(seen_sources)) == 3


# ─────────────────────────────────────────────────────────────────────────────
# W20E5 Codex B1 fix — production-path regression for ThreadPool propagation.
#
# The production background_render step submits ``_process`` (which calls
# ``render_one_background`` → openai.images.*) to a ThreadPoolExecutor. The
# parent-thread budget is invisible to those workers unless the submit site
# wraps the callable in ``bind_current_budget``. The tests below mimic the
# production submit pattern with the actual ``render_one_background`` helper
# and a fake openai client, and prove that — with the binding — cap=0
# blocks the provider call from the child thread, and cap=N is authoritative
# across many workers.
# ─────────────────────────────────────────────────────────────────────────────


def test_thread_pool_child_with_propagated_cap_zero_blocks_provider(tmp_path):
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.generate.return_value = _ok_response()

    def _process(idx: int):
        out = tmp_path / f"bg_{idx}.png"
        return render_one_background(
            openai_client=client, image_model="gpt-image-2.5-sunburst", prompt="x",
            out_path=out, fp_path=None, prior_bg_paths=[], max_attempts=1,
        )

    try:
        with ThreadPoolExecutor(max_workers=3) as pool:
            futures = [pool.submit(bind_current_budget(_process), i) for i in range(4)]
            for fut in as_completed(futures):
                with pytest.raises(ImageCallBudgetExceeded):
                    fut.result()
    finally:
        uninstall_budget()

    # Every child thread saw the propagated cap=0 budget — no provider call.
    assert client.images.generate.call_count == 0
    assert client.images.edit.call_count == 0


def test_thread_pool_child_with_propagated_cap_n_is_authoritative(tmp_path):
    """cap=2 with 5 submitted jobs → exactly 2 provider calls allowed."""
    budget = ImageCallBudget(cap=2)
    install_budget(budget)
    client = MagicMock()
    client.images.generate.return_value = _ok_response()
    seen_threads: set[str] = set()
    lock = threading.Lock()

    def _process(idx: int):
        with lock:
            seen_threads.add(threading.current_thread().name)
        out = tmp_path / f"bg_{idx}.png"
        return render_one_background(
            openai_client=client, image_model="gpt-image-2.5-sunburst", prompt="x",
            out_path=out, fp_path=None, prior_bg_paths=[], max_attempts=1,
        )

    successes = 0
    denials = 0
    try:
        with ThreadPoolExecutor(max_workers=4) as pool:
            futures = [pool.submit(bind_current_budget(_process), i) for i in range(5)]
            for fut in as_completed(futures):
                try:
                    res = fut.result()
                    if res.get("status") == "ok":
                        successes += 1
                except ImageCallBudgetExceeded:
                    denials += 1
    finally:
        uninstall_budget()

    snap = budget.snapshot()
    assert snap["used"] == 2
    assert snap["denied"] == 3
    assert successes == 2
    assert denials == 3
    assert client.images.generate.call_count == 2
    # Sanity: real worker threads were used (otherwise the test does not
    # exercise the propagation hop).
    assert len(seen_threads) >= 2


def test_thread_pool_child_without_binding_misses_parent_cap(tmp_path):
    """Negative control — without ``bind_current_budget``, a child thread
    sees no budget and the cap is silently bypassed. This is the bug
    Codex B1 found; the production submit sites must use the binding.
    """
    install_budget(ImageCallBudget(cap=0))
    client = MagicMock()
    client.images.generate.return_value = _ok_response()

    def _process(idx: int):
        out = tmp_path / f"bg_{idx}.png"
        return render_one_background(
            openai_client=client, image_model="gpt-image-2.5-sunburst", prompt="x",
            out_path=out, fp_path=None, prior_bg_paths=[], max_attempts=1,
        )

    try:
        with ThreadPoolExecutor(max_workers=2) as pool:
            # NOTE: deliberately NOT wrapped in bind_current_budget.
            futures = [pool.submit(_process, i) for i in range(2)]
            results = [f.result() for f in futures]
    finally:
        uninstall_budget()

    # Without the binding, child threads bypass the cap → provider was hit.
    assert client.images.generate.call_count == 2
    for r in results:
        assert r.get("status") == "ok"