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

``GeminiImageClient.generate_image`` opens an HTTPS request via
``urllib.request.urlopen`` in a retry loop; each urlopen attempt is a
distinct network call and must reserve one budget unit beforehand. Cap
exhaustion raises ``ImageCallBudgetExceeded`` and the helper must not
swallow it via the broad ``except (URLError, socket.timeout)`` arm.
"""
from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from app.core.image_call_budget import (
    ImageCallBudget,
    ImageCallBudgetExceeded,
    bind_current_budget,
    install_budget,
    uninstall_budget,
)


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


def _make_client():
    """Build a GeminiImageClient with a fixed key (skip rotating pool)."""
    from app.modules.llm.gemini_image_client import GeminiImageClient
    return GeminiImageClient(api_key="dummy-key", model="gemini-3.1-flash-image-preview")


def _patch_log_pieces(monkeypatch):
    """Stub DB/Opik loggers so no network/DB writes occur."""
    import app.modules.llm.gemini_image_client as gi
    monkeypatch.setattr(gi, "log_llm_call", lambda **kw: None)

    fake_tracer = MagicMock()
    fake_tracer.log = MagicMock()
    monkeypatch.setattr(
        "app.modules.llm.image_tracer.get_image_tracer",
        lambda: fake_tracer,
    )


def test_no_budget_keeps_legacy_call_path(monkeypatch):
    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    calls = {"n": 0}

    class _FakeResponse:
        def __enter__(self):
            return self
        def __exit__(self, *a):
            return False
        def read(self):
            import base64, json
            png_b64 = base64.b64encode(b"x" * 32).decode()
            return json.dumps({
                "candidates": [{"content": {"parts": [
                    {"inlineData": {"data": png_b64}},
                ]}}],
            }).encode()

    def _fake_urlopen(req, timeout=None):
        calls["n"] += 1
        return _FakeResponse()

    monkeypatch.setattr(gi.urllib.request, "urlopen", _fake_urlopen)

    client = _make_client()
    png, _ms = client.generate_image(prompt="x")
    assert isinstance(png, bytes) and len(png) > 0
    assert calls["n"] == 1


def test_cap_zero_blocks_urlopen(monkeypatch):
    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    fake_urlopen = MagicMock()
    monkeypatch.setattr(gi.urllib.request, "urlopen", fake_urlopen)

    install_budget(ImageCallBudget(cap=0))
    client = _make_client()
    with pytest.raises(ImageCallBudgetExceeded):
        client.generate_image(prompt="x")
    assert fake_urlopen.call_count == 0


def test_cap_one_allows_exactly_one_urlopen(monkeypatch):
    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    calls = {"n": 0}

    class _FakeResponse:
        def __enter__(self): return self
        def __exit__(self, *a): return False
        def read(self):
            import base64, json
            png_b64 = base64.b64encode(b"x" * 32).decode()
            return json.dumps({
                "candidates": [{"content": {"parts": [
                    {"inlineData": {"data": png_b64}},
                ]}}],
            }).encode()

    def _fake_urlopen(req, timeout=None):
        calls["n"] += 1
        return _FakeResponse()

    monkeypatch.setattr(gi.urllib.request, "urlopen", _fake_urlopen)

    budget = ImageCallBudget(cap=1)
    install_budget(budget)
    client = _make_client()
    png, _ms = client.generate_image(prompt="x")
    assert isinstance(png, bytes)
    assert calls["n"] == 1
    assert budget.snapshot()["used"] == 1


# ─────────────────────────────────────────────────────────────────────────────
# W20E5 Codex B2 — nested ThreadPool propagation for GeminiImageClient.
#
# The Gemini image client is invoked from several nested ThreadPool sites
# (scene_generation_coordinator variation pool, reference_phase1/2/3 pools).
# This test proves that with ``bind_current_budget`` wrapping the submit,
# a parent-installed budget is enforced at the urlopen call site inside
# the child thread. cap=0 in parent → 0 urlopen calls in any child.
# ─────────────────────────────────────────────────────────────────────────────


def test_gemini_image_client_in_propagated_thread_pool_respects_cap_zero(monkeypatch):
    """Nested-pool regression — GeminiImageClient.generate_image called
    inside a ThreadPoolExecutor child must observe the parent budget when
    the submit is wrapped with ``bind_current_budget``.
    """
    from concurrent.futures import ThreadPoolExecutor, as_completed

    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    fake_urlopen = MagicMock()
    monkeypatch.setattr(gi.urllib.request, "urlopen", fake_urlopen)

    install_budget(ImageCallBudget(cap=0))

    def _worker(idx: int):
        client = _make_client()
        client.generate_image(prompt=f"prompt {idx}")
        return idx

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

    # Every child blocked before its urlopen call.
    assert fake_urlopen.call_count == 0


def test_gemini_image_client_in_propagated_thread_pool_caps_to_n(monkeypatch):
    """cap=2 across 5 nested-pool submissions → exactly 2 urlopen calls."""
    from concurrent.futures import ThreadPoolExecutor, as_completed
    import threading

    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    calls = {"n": 0}
    lock = threading.Lock()

    class _FakeResponse:
        def __enter__(self): return self
        def __exit__(self, *a): return False
        def read(self):
            import base64, json
            png_b64 = base64.b64encode(b"x" * 32).decode()
            return json.dumps({
                "candidates": [{"content": {"parts": [
                    {"inlineData": {"data": png_b64}},
                ]}}],
            }).encode()

    def _fake_urlopen(req, timeout=None):
        with lock:
            calls["n"] += 1
        return _FakeResponse()

    monkeypatch.setattr(gi.urllib.request, "urlopen", _fake_urlopen)

    budget = ImageCallBudget(cap=2)
    install_budget(budget)

    barrier = threading.Barrier(4)
    seen_threads: set[str] = set()

    def _worker(idx: int):
        # Wait at a barrier so all submitted workers run concurrently on
        # distinct pool threads — proves cross-thread propagation rather
        # than pool reuse of a single worker after fast returns.
        try:
            barrier.wait(timeout=2.0)
        except threading.BrokenBarrierError:
            # 5th submission may not be admitted concurrently — barrier
            # still releases the rest. Continue without failing the test
            # purely on scheduling timing.
            pass
        with lock:
            seen_threads.add(threading.current_thread().name)
        client = _make_client()
        return client.generate_image(prompt=f"prompt {idx}")

    successes = 0
    denials = 0
    try:
        with ThreadPoolExecutor(max_workers=4) as pool:
            futures = [pool.submit(bind_current_budget(_worker), i) for i in range(5)]
            for fut in as_completed(futures):
                try:
                    fut.result()
                    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 calls["n"] == 2
    # Sanity — multiple child threads actually used (barrier forces concurrency).
    assert len(seen_threads) >= 2


def test_cap_exhausts_on_retry_loop_without_swallowing(monkeypatch):
    """cap=2 + transient URLError → 2 urlopen attempts then ImageCallBudgetExceeded.

    The broad ``except (URLError, socket.timeout)`` must not swallow the
    budget error.
    """
    _patch_log_pieces(monkeypatch)
    import app.modules.llm.gemini_image_client as gi

    calls = {"n": 0}

    def _flaky_urlopen(req, timeout=None):
        calls["n"] += 1
        # ★2026-08-26: 예외를 `URLError("transient")` 에서 「연결 거부」로
        #  바꿨다. 이 시험이 보는 것은 **예산 상한이 재시도 루프에서
        #  소진되는가**인데, 새 전송 판정에서 원인 문자열만 있는 URLError 는
        #  「보냈는지 모른다」라 재시도 자체를 안 한다(요금 보호). 재시도를
        #  정당하게 유발하는 예외로 바꿔야 원래 계약을 계속 잰다.
        raise gi.urllib.error.URLError(ConnectionRefusedError())

    monkeypatch.setattr(gi.urllib.request, "urlopen", _flaky_urlopen)
    monkeypatch.setattr(gi.time, "sleep", lambda _s: None)
    # ensure max_retries is permissive — settings default may be small,
    # we just need the loop to *want* more attempts than cap allows.
    monkeypatch.setattr(gi._settings, "llm_max_retries", 5, raising=False)

    budget = ImageCallBudget(cap=2)
    install_budget(budget)
    client = _make_client()
    with pytest.raises(ImageCallBudgetExceeded):
        client.generate_image(prompt="x")
    assert calls["n"] == 2
    assert budget.snapshot()["used"] == 2
    assert budget.snapshot()["denied"] == 1