"""W20E5 — runtime image call budget primitive.

A thread-safe counter installed for the duration of a background run that
acts as the authoritative cap at every gpt-image-2 / Gemini image / fal
provider call site. ``reserve()`` raises ``ImageCallBudgetExceeded`` BEFORE
the network request is allowed, so cap=0 / exhausted runs never reach the
provider.

The primitive is stdlib-only; no ``app.*`` imports and no provider SDK
imports (verified by an AST test). Call sites import
``reserve_current_call`` which is a no-op when no budget is installed on
the current thread, keeping read-only paths and existing tests unaffected.
"""
from __future__ import annotations

import threading
from typing import Any, Dict, Optional


class ImageCallBudgetExceeded(RuntimeError):
    """Raised by ``ImageCallBudget.reserve`` when the cap is reached.

    The exception is raised *before* the provider/network call so the
    caller can abort without touching the SDK.
    """

    def __init__(self, *, cap: int, used: int, source: str) -> None:
        self.cap = int(cap)
        self.used = int(used)
        self.source = str(source)
        super().__init__(
            f"image_call_budget exceeded — cap={self.cap}, used={self.used}, "
            f"denied source={self.source!r}"
        )


class ImageCallBudget:
    """Counter with a hard cap on image-provider calls.

    Counts are managed under a re-entrant lock; ``reserve`` and ``snapshot``
    are safe to call from multiple threads against the same instance.
    """

    __slots__ = ("_cap", "_used", "_denied", "_lock")

    def __init__(self, cap: int) -> None:
        cap_int = int(cap)
        if cap_int < 0:
            raise ValueError(f"cap must be >= 0 (got {cap_int})")
        self._cap = cap_int
        self._used = 0
        self._denied = 0
        self._lock = threading.RLock()

    def reserve(self, *, source: str) -> None:
        """Reserve one unit of the budget for a provider call.

        Must be called *before* the actual provider/network call. Raises
        :class:`ImageCallBudgetExceeded` and increments the denied counter
        when the cap has already been reached; the used counter is left
        unchanged in that case.
        """
        with self._lock:
            if self._used >= self._cap:
                self._denied += 1
                raise ImageCallBudgetExceeded(
                    cap=self._cap, used=self._used, source=source,
                )
            self._used += 1

    def snapshot(self) -> Dict[str, int]:
        """Return a JSON-serialisable copy of the current counts."""
        with self._lock:
            return {
                "cap": self._cap,
                "used": self._used,
                "denied": self._denied,
                "remaining": self._cap - self._used,
            }


# ─────────────────────────────────────────────────────────────────────────────
# Per-thread install. We use threading.local rather than ContextVar so the
# install is explicit and isolated to the worker thread that performs image
# calls; child threads start clean and main threads cannot leak budgets
# across runs.
# ─────────────────────────────────────────────────────────────────────────────


_local = threading.local()


def install_budget(budget: ImageCallBudget) -> None:
    """Install ``budget`` as the current-thread budget.

    The previously installed budget (if any) is replaced. Callers should
    pair this with :func:`uninstall_budget` in a try/finally.
    """
    _local.budget = budget


def uninstall_budget() -> None:
    """Clear the current-thread budget (no error if nothing installed)."""
    _local.budget = None


def get_current_budget() -> Optional[ImageCallBudget]:
    """Return the budget installed on this thread, or ``None``."""
    return getattr(_local, "budget", None)


def reserve_current_call(*, source: str) -> None:
    """Reserve one image-call unit against the current-thread budget.

    No-op when no budget is installed — keeps non-budgeted code paths
    (legacy unit tests, dry-runs, plan-only runs) unaffected.

    Raises :class:`ImageCallBudgetExceeded` when the installed budget is
    exhausted; callers should let the exception propagate so that
    ``run_steps_batch`` can record it and abort the batch.
    """
    budget = get_current_budget()
    if budget is None:
        return
    budget.reserve(source=source)


# ─────────────────────────────────────────────────────────────────────────────
# Thread-pool propagation helpers (W20E5 Codex B1 fix).
#
# ``install_budget`` is per-thread, so a ``ThreadPoolExecutor`` worker that
# inherits no per-thread state would see no budget — making the cap a no-op
# at every image call site reached through a child thread. The helpers
# below explicitly propagate the *parent's* current budget into the worker
# for the duration of a callable, then restore whatever was previously
# installed (typically ``None``) so the worker thread does not leak the
# budget into later, unrelated submissions on the same pool.
# ─────────────────────────────────────────────────────────────────────────────


def run_with_budget(
    budget: Optional[ImageCallBudget],
    fn: Any,
    *args: Any,
    **kwargs: Any,
) -> Any:
    """Run ``fn(*args, **kwargs)`` with ``budget`` installed on this thread.

    ``budget=None`` is a no-op for installation (preserves the current
    install). The previous thread-local budget — if any — is restored in a
    ``finally`` block so exceptions cannot leak install state.
    """
    if budget is None:
        return fn(*args, **kwargs)
    previous = get_current_budget()
    install_budget(budget)
    try:
        return fn(*args, **kwargs)
    finally:
        if previous is None:
            uninstall_budget()
        else:
            install_budget(previous)


def bind_current_budget(fn: Any) -> Any:
    """Return a callable that re-installs *this thread's* current budget
    when invoked on another thread.

    Intended use::

        from concurrent.futures import ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=4) as pool:
            pool.submit(bind_current_budget(_worker), arg1, arg2)

    The budget is captured at ``bind_current_budget`` call time, not at
    callable invocation time — so the parent may safely clear its own
    install before the worker actually runs.
    """
    captured = get_current_budget()

    def _wrapped(*args: Any, **kwargs: Any) -> Any:
        return run_with_budget(captured, fn, *args, **kwargs)

    return _wrapped


__all__ = [
    "ImageCallBudget",
    "ImageCallBudgetExceeded",
    "install_budget",
    "uninstall_budget",
    "get_current_budget",
    "reserve_current_call",
    "run_with_budget",
    "bind_current_budget",
]


# Static guard helper for callers that want to surface a friendly state in
# logs/reports without raising. Returns ``None`` when no budget installed.
def current_snapshot() -> Optional[Dict[str, Any]]:
    budget = get_current_budget()
    if budget is None:
        return None
    return budget.snapshot()