"""Cross-cron resource locks for the AI-spending period rollup.

``ai_spending_period_rollup`` is a materialised cache rebuilt window by
window with DELETE-then-INSERT. Two writers overlapping on one window do
not collide loudly: the table carries no unique constraint (only
``ix_ai_rollup_window_fp``), so both inserts survive and every read — which
SUMs the window's rows — silently double-counts.

Until 2026-07-28 there was exactly one writer (the
``ai_spending.period_rollup_refresh`` cron) and the scheduled-jobs runner's
single-flight lock was enough. Then the drift-check cron gained an
auto-heal path that re-materialises drifted windows (#2083) — a correct fix
for a real incident, but the runner's lock is keyed ``sched_job:{job_id}``,
so it serialises a job against ITSELF and not against a different job. The
two crons sit 30 minutes apart (06:30 / 07:00) and either can be delayed by
the concurrency gate, so "they won't overlap" is a schedule coincidence,
not an invariant.

Every writer of a rollup window takes this lock. It is deliberately ONE
key rather than one per window: the windows a writer touches are decided
inside the writer (the refresh walks ``months_back``, the heal walks the
drifted set), so a per-window lock would have to be acquired in a loop and
could still interleave two writers on different windows of the same
fingerprint rebuild. Coarse and obviously correct beats fine and subtly
wrong for a job that runs once a day.
"""
from __future__ import annotations

import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

from app.core.distributed_lock import LockBusy, distributed_lock

ROLLUP_WRITE_LOCK = "ai_spending:period_rollup_write"

# The refresh cron walks up to ``months_back`` windows over a multi-million
# row facts table; the heal path rebuilds only the drifted windows. 20
# minutes covers the former with headroom, and `auto_renew` extends it if a
# backfill makes the scan unusually slow.
ROLLUP_WRITE_LOCK_TTL = 1200


class RollupLockBusy(Exception):
    """``ROLLUP_WRITE_LOCK`` was already held at *acquire* time.

    This is the benign "another writer is rebuilding right now" overlap —
    the caller may skip (refresh cron), defer (drift-check heal), or 409
    (manual refresh). It is deliberately a DIFFERENT type from the core
    :class:`app.core.distributed_lock.LockBusy`, because ``distributed_lock``
    also raises ``LockBusy`` when the auto-renew watchdog detects the lock
    was LOST mid-hold (Redis eviction / manual delete). Mid-hold loss means
    the protected DELETE+INSERT may have interleaved with a new holder and
    the window can be double-counted — that must abort loudly (handler run
    recorded ``failed``, failure_notifier pages), never be reported as a
    harmless "skipped". A plain ``except LockBusy`` around the whole
    ``async with`` cannot tell the two apart; :func:`rollup_write_lock`
    can, and callers must catch only this class.
    """

    def __init__(self) -> None:
        super().__init__(f"rollup write lock busy: {ROLLUP_WRITE_LOCK}")


@asynccontextmanager
async def rollup_write_lock() -> AsyncIterator[None]:
    """Hold ``ROLLUP_WRITE_LOCK`` (auto-renewed) for the ``async with`` body.

    * Acquire failure → :class:`RollupLockBusy` (catch it; see its docstring).
    * Lock LOST mid-hold → the core :class:`LockBusy` propagates out of the
      block unchanged so the run fails loudly. Do NOT catch it here or in
      callers.

    Every writer of a rollup window (refresh cron, drift-heal, manual
    refresh via HTTP router or MCP tool) goes through this helper so the
    acquire-vs-lost distinction is made in exactly one place.
    """
    try:
        ctx = distributed_lock(
            ROLLUP_WRITE_LOCK, ttl=ROLLUP_WRITE_LOCK_TTL, auto_renew=True,
        )
        await ctx.__aenter__()
    except LockBusy as exc:
        raise RollupLockBusy() from exc
    try:
        yield
    except BaseException:
        # ``__aexit__`` re-raises the mid-hold-loss LockBusy conversion
        # (the watchdog cancels the holder; the generator turns the
        # CancelledError into LockBusy). Returns False when the original
        # exception should propagate.
        if not await ctx.__aexit__(*sys.exc_info()):
            raise
    else:
        await ctx.__aexit__(None, None, None)