"""W20E5 — image_call_budget primitive (RED).

Pure runtime image call budget used as the authoritative cap at every provider
call site (OpenAI gpt-image-2, Gemini image, fal). No-op when no budget is
installed in the current execution context.

Failure mode: ``reserve()`` raises ``ImageCallBudgetExceeded`` BEFORE the
provider call is allowed; ``denied`` counter increments, ``used`` does not.

Module invariants (asserted via AST):
- stdlib only — no ``app.*`` import, no provider SDK import
  (``openai`` / ``litellm`` / ``fal`` / ``google`` / ``anthropic`` / ``PIL``
  / ``requests`` / ``httpx``).
"""
from __future__ import annotations

import ast
import threading
from pathlib import Path

import pytest

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


# ─────────────────────────────────────────────────────────────────────────────
# Primitive behaviour
# ─────────────────────────────────────────────────────────────────────────────


def test_initial_snapshot_has_cap_used_zero_denied_zero():
    b = ImageCallBudget(cap=3)
    snap = b.snapshot()
    assert snap["cap"] == 3
    assert snap["used"] == 0
    assert snap["denied"] == 0
    assert snap["remaining"] == 3


def test_reserve_increments_used():
    b = ImageCallBudget(cap=2)
    b.reserve(source="t.a")
    assert b.snapshot() == {"cap": 2, "used": 1, "denied": 0, "remaining": 1}
    b.reserve(source="t.a")
    assert b.snapshot() == {"cap": 2, "used": 2, "denied": 0, "remaining": 0}


def test_reserve_at_cap_raises_and_increments_denied_not_used():
    b = ImageCallBudget(cap=1)
    b.reserve(source="t.a")
    with pytest.raises(ImageCallBudgetExceeded) as ei:
        b.reserve(source="t.b")
    assert b.snapshot() == {"cap": 1, "used": 1, "denied": 1, "remaining": 0}
    msg = str(ei.value)
    assert "cap" in msg.lower() and "t.b" in msg


def test_reserve_cap_zero_raises_immediately():
    b = ImageCallBudget(cap=0)
    with pytest.raises(ImageCallBudgetExceeded):
        b.reserve(source="t.a")
    assert b.snapshot() == {"cap": 0, "used": 0, "denied": 1, "remaining": 0}


def test_negative_cap_rejected_at_construction():
    with pytest.raises(ValueError):
        ImageCallBudget(cap=-1)


def test_reserve_requires_source_kwarg():
    b = ImageCallBudget(cap=1)
    with pytest.raises(TypeError):
        b.reserve()  # type: ignore[call-arg]


def test_snapshot_is_independent_dict():
    b = ImageCallBudget(cap=2)
    b.reserve(source="t.a")
    snap1 = b.snapshot()
    snap1["used"] = 999  # mutate copy
    assert b.snapshot()["used"] == 1  # internal state unaffected


# ─────────────────────────────────────────────────────────────────────────────
# Context install / uninstall
# ─────────────────────────────────────────────────────────────────────────────


def test_get_current_budget_default_is_none():
    uninstall_budget()  # defensive
    assert get_current_budget() is None


def test_install_then_get_returns_budget():
    b = ImageCallBudget(cap=1)
    install_budget(b)
    try:
        assert get_current_budget() is b
    finally:
        uninstall_budget()


def test_uninstall_clears_current():
    b = ImageCallBudget(cap=1)
    install_budget(b)
    uninstall_budget()
    assert get_current_budget() is None


def test_reserve_current_call_is_noop_when_no_budget():
    uninstall_budget()
    # Must not raise, must not register anything.
    reserve_current_call(source="t.a")


def test_reserve_current_call_uses_installed_budget():
    b = ImageCallBudget(cap=2)
    install_budget(b)
    try:
        reserve_current_call(source="t.a")
        reserve_current_call(source="t.b")
        assert b.snapshot() == {"cap": 2, "used": 2, "denied": 0, "remaining": 0}
    finally:
        uninstall_budget()


def test_reserve_current_call_raises_when_exhausted():
    b = ImageCallBudget(cap=1)
    install_budget(b)
    try:
        reserve_current_call(source="t.a")
        with pytest.raises(ImageCallBudgetExceeded):
            reserve_current_call(source="t.b")
        assert b.snapshot() == {"cap": 1, "used": 1, "denied": 1, "remaining": 0}
    finally:
        uninstall_budget()


def test_raw_thread_child_starts_clean_without_propagation_helper():
    """Without an explicit propagation helper, a raw ``threading.Thread``
    child must start with no installed budget. This documents the *default*
    behaviour — production callers that need the parent's budget inside a
    worker thread must use :func:`bind_current_budget` /
    :func:`run_with_budget`.
    """
    b_main = ImageCallBudget(cap=5)
    install_budget(b_main)

    seen_in_child: dict = {}

    def _child():
        seen_in_child["budget"] = get_current_budget()

    t = threading.Thread(target=_child)
    t.start()
    t.join()
    try:
        assert seen_in_child["budget"] is None
        assert get_current_budget() is b_main
    finally:
        uninstall_budget()


# ─────────────────────────────────────────────────────────────────────────────
# Propagation helpers — run_with_budget / bind_current_budget
# ─────────────────────────────────────────────────────────────────────────────


def test_run_with_budget_installs_for_duration_then_restores_previous():
    """``run_with_budget`` installs the given budget for the call and
    restores whatever was previously installed (None or another budget).
    """
    uninstall_budget()
    assert get_current_budget() is None
    b = ImageCallBudget(cap=3)

    seen: list = []

    def _fn():
        seen.append(get_current_budget())

    run_with_budget(b, _fn)
    assert seen == [b]
    # Previous state (None) restored.
    assert get_current_budget() is None


def test_run_with_budget_restores_prior_budget_not_just_none():
    outer = ImageCallBudget(cap=2)
    inner = ImageCallBudget(cap=99)
    install_budget(outer)
    try:
        seen: list = []
        def _fn():
            seen.append(get_current_budget())
        run_with_budget(inner, _fn)
        assert seen == [inner]
        # restored to outer, not None.
        assert get_current_budget() is outer
    finally:
        uninstall_budget()


def test_run_with_budget_restores_even_on_exception():
    b = ImageCallBudget(cap=1)
    uninstall_budget()
    with pytest.raises(RuntimeError):
        run_with_budget(b, lambda: (_ for _ in ()).throw(RuntimeError("boom")))
    assert get_current_budget() is None


def test_run_with_budget_forwards_args_kwargs_and_return_value():
    b = ImageCallBudget(cap=1)
    def _add(a, b, *, c):
        return a + b + c
    result = run_with_budget(b, _add, 1, 2, c=3)
    assert result == 6


def test_run_with_budget_passes_none_as_no_install():
    install_budget(ImageCallBudget(cap=2))
    try:
        seen: list = []
        def _fn():
            seen.append(get_current_budget())
        run_with_budget(None, _fn)
        # When budget=None, current install must be preserved (caller may
        # legitimately want to run with whatever the caller had).
        assert seen[0] is get_current_budget()
    finally:
        uninstall_budget()


def test_bind_current_budget_captures_parent_thread_budget_at_call_time():
    """``bind_current_budget`` snapshots the *parent's* budget when the
    helper is called and propagates that captured value into the worker
    thread when the returned callable executes there.
    """
    parent_budget = ImageCallBudget(cap=4)
    install_budget(parent_budget)
    try:
        def _worker():
            return get_current_budget()
        bound = bind_current_budget(_worker)
    finally:
        # Even if the parent later clears its own budget, the bound
        # callable must still see the captured one.
        uninstall_budget()

    # Run the bound callable in a fresh thread.
    seen_in_thread: dict = {}
    def _thread_body():
        seen_in_thread["budget"] = bound()
        seen_in_thread["after"] = get_current_budget()
    t = threading.Thread(target=_thread_body)
    t.start()
    t.join()
    assert seen_in_thread["budget"] is parent_budget
    # After the bound call exits the worker thread, that thread is left
    # clean (it never had a budget before).
    assert seen_in_thread["after"] is None


def test_bind_current_budget_no_op_when_no_parent_budget():
    uninstall_budget()
    def _worker():
        return get_current_budget()
    bound = bind_current_budget(_worker)
    seen: dict = {}
    def _thread_body():
        seen["budget"] = bound()
    t = threading.Thread(target=_thread_body)
    t.start()
    t.join()
    assert seen["budget"] is None


def test_bind_current_budget_propagates_into_thread_pool_executor():
    """End-to-end: ThreadPoolExecutor child workers must see the captured
    parent budget when callables are wrapped with ``bind_current_budget``.
    """
    from concurrent.futures import ThreadPoolExecutor

    parent_budget = ImageCallBudget(cap=10)
    install_budget(parent_budget)

    def _worker(idx: int) -> tuple[int, object, str]:
        # Returns (idx, budget-seen-in-child, thread-name).
        return (idx, get_current_budget(), threading.current_thread().name)

    try:
        with ThreadPoolExecutor(max_workers=3) as pool:
            futures = [pool.submit(bind_current_budget(_worker), i) for i in range(6)]
            results = [f.result() for f in futures]
    finally:
        uninstall_budget()

    for idx, seen_budget, name in results:
        assert seen_budget is parent_budget, (
            f"child worker (thread={name}, idx={idx}) did not see the "
            f"propagated parent budget"
        )


def test_bind_current_budget_does_not_leak_into_subsequent_unbound_submissions():
    """A worker thread that ran a bound call must not retain the budget
    for later raw (unbound) submissions on the same pool — ``run_with_budget``
    restores the previous (None) value at the end of each bound call.
    """
    from concurrent.futures import ThreadPoolExecutor

    parent_budget = ImageCallBudget(cap=5)
    install_budget(parent_budget)
    bound = bind_current_budget(lambda: get_current_budget())
    try:
        with ThreadPoolExecutor(max_workers=1) as pool:
            # Submit bound first so the worker definitely installs/uninstalls.
            assert pool.submit(bound).result() is parent_budget
            # Now submit an unbound callable on the same worker.
            unbound_seen = pool.submit(lambda: get_current_budget()).result()
    finally:
        uninstall_budget()
    assert unbound_seen is None, (
        "raw unbound call after a bound call must see no propagated budget"
    )


def test_reserve_inside_bound_thread_counts_against_parent_budget():
    """Reservations made inside a propagated worker increment the *parent's*
    budget counter — the cap is authoritative even across thread hops.
    """
    from concurrent.futures import ThreadPoolExecutor

    parent_budget = ImageCallBudget(cap=3)
    install_budget(parent_budget)

    def _reserve_in_child(idx: int) -> str:
        try:
            reserve_current_call(source=f"child.{idx}")
            return "ok"
        except ImageCallBudgetExceeded:
            return "denied"

    try:
        with ThreadPoolExecutor(max_workers=4) as pool:
            futures = [
                pool.submit(bind_current_budget(_reserve_in_child), i)
                for i in range(5)
            ]
            outcomes = sorted(f.result() for f in futures)
    finally:
        uninstall_budget()

    snap = parent_budget.snapshot()
    assert snap["cap"] == 3
    assert snap["used"] == 3
    assert snap["denied"] == 2
    assert outcomes.count("ok") == 3
    assert outcomes.count("denied") == 2


# ─────────────────────────────────────────────────────────────────────────────
# Thread-safety of the counter object itself (shared budget, many threads)
# ─────────────────────────────────────────────────────────────────────────────


def test_concurrent_reserve_respects_cap_and_does_not_double_count():
    """Single budget shared across threads — no race past cap."""
    cap = 50
    workers = 16
    attempts_per_worker = 20  # 16 * 20 = 320 attempts, cap = 50

    b = ImageCallBudget(cap=cap)
    barrier = threading.Barrier(workers)

    successes = [0]
    denials = [0]
    lock = threading.Lock()

    def _worker():
        barrier.wait()
        local_ok = 0
        local_denied = 0
        for _ in range(attempts_per_worker):
            try:
                b.reserve(source="t.x")
                local_ok += 1
            except ImageCallBudgetExceeded:
                local_denied += 1
        with lock:
            successes[0] += local_ok
            denials[0] += local_denied

    threads = [threading.Thread(target=_worker) for _ in range(workers)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    snap = b.snapshot()
    assert snap["used"] == cap
    assert successes[0] == cap
    assert denials[0] == workers * attempts_per_worker - cap
    assert snap["denied"] == denials[0]


# ─────────────────────────────────────────────────────────────────────────────
# Static module invariants (no forbidden imports)
# ─────────────────────────────────────────────────────────────────────────────


_FORBIDDEN_TOP_LEVEL_IMPORT_ROOTS = frozenset({
    "openai",
    "litellm",
    "fal",
    "fal_client",
    "google",
    "anthropic",
    "PIL",
    "requests",
    "httpx",
})


def _budget_module_source() -> str:
    p = Path(__file__).resolve().parents[2] / "app" / "core" / "image_call_budget.py"
    return p.read_text(encoding="utf-8")


def test_module_has_no_app_or_provider_imports():
    tree = ast.parse(_budget_module_source())
    offenders: list[str] = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                root = alias.name.split(".", 1)[0]
                if root == "app":
                    offenders.append(f"import {alias.name}")
                if root in _FORBIDDEN_TOP_LEVEL_IMPORT_ROOTS:
                    offenders.append(f"import {alias.name}")
        elif isinstance(node, ast.ImportFrom):
            mod = (node.module or "")
            root = mod.split(".", 1)[0]
            if root == "app":
                offenders.append(f"from {mod} import ...")
            if root in _FORBIDDEN_TOP_LEVEL_IMPORT_ROOTS:
                offenders.append(f"from {mod} import ...")
    assert offenders == [], (
        "image_call_budget.py must be stdlib-only — found: " + ", ".join(offenders)
    )