"""The learned roster: which of FabledCurator's parts have checked in, and when.

Milestone 365. `celery inspect` answers "who is here"; this answers "who is
missing", which nothing in the application could do before — see
`models/service_seen.py` for why the identity is a queue set and not a
worker hostname.

## Who does the observing, and why it is the web process

Three candidates, and the choice matters more than the code:

* **A celery beat sweep.** Rejected. If the scheduler dies, the sweep stops,
  every row goes stale, and the page reports that everything is down when one
  thing is. An alarm that cannot distinguish "one part died" from "the
  observer died" is worse than no alarm.
* **A background task in web.** Rejected on a detail of how this deploys:
  hypercorn runs `--workers 4`, so a `before_serving` loop would be FOUR
  concurrent inspect loops hammering the broker, forever, per container.
* **Refresh on demand, rate-limited by the data itself.** Taken. Whichever web
  process happens to serve a health request refreshes the roster if it is
  older than REFRESH_TTL, and otherwise reads what is already there.

The third has the property the other two lack: **the observer is the thing
serving the page.** If web is down you get a browser error rather than a
confidently green page, which is the honest failure. It also self-limits
without coordination — the TTL lives in the row everybody can see.
"""

from __future__ import annotations

import asyncio
import logging

from sqlalchemy import func
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import ServiceSeen
from .worker_lanes import LANES, lane_for_node

log = logging.getLogger(__name__)

# There is no refresh TTL any more. It existed because the HEALTH REQUEST
# refreshed the roster, rate-limited to 20s so that a page open in two tabs
# did not inspect twice as often. `size_worker_lanes` owns the refresh now, on
# `SWEEP_PERIOD_SECONDS`, so the cadence is a schedule rather than a side
# effect of someone looking.

# celery inspect is a broker round trip and this sits on a request path, so it
# gets a deadline (rule 156). A broker that has stopped answering must make the
# roster stale — which is a true statement about the system — not hang the one
# page that exists to explain it.
INSPECT_TIMEOUT_SECONDS = 2.0

# How many broadcast round trips `_inspect_celery_sync` makes. Named, because
# the wrapper's budget is derived from it and the two must not drift.
#
# `active_queues()` and `active()` are separate broadcasts, and a broadcast
# with no `destination` cannot know how many replies to expect — so each one
# waits out its full timeout rather than returning on the last reply. The sync
# call therefore costs ~2 x INSPECT_TIMEOUT_SECONDS in the ordinary case, not
# once.
INSPECT_ROUND_TRIPS = 2

# Slack for the thread handoff. `asyncio.to_thread` hands work to the default
# executor, and on a loaded web process — the operator's showcase page pulling
# ninety thumbnails a second — the thread may not even be scheduled inside the
# budget, let alone finish.
#
# This exists because the wrapper used to allow `INSPECT_TIMEOUT_SECONDS * 2`,
# which LOOKS like a safety factor and is exactly the worst case with nothing
# left over. Observed on the operator's first consolidated deploy, 2026-09-23:
# a TimeoutError traceback per refresh while the two inspect calls were
# working perfectly. A budget equal to the work is a budget that fails under
# any load at all.
INSPECT_SLACK_SECONDS = 3.0

# Queue set -> the name an operator recognises. Sorted-tuple keys, because the
# order celery reports them in is not guaranteed.
#
# DERIVED from `worker_lanes.LANES` (milestone 422 step 1) rather than written
# out here. It was a hand-kept second copy of the same fact, and it had already
# drifted: `maintenance_long` is a live lane with four task routes pointing at
# it and a dedicated worker in the operator's stack, and this map did not know
# it — so the System tab labelled it `Worker (maintenance_long)`. One list of
# lanes now names them everywhere.
#
# A deployment that slices CELERY_QUEUES differently still falls through to the
# raw queue list rather than being given a name this code invented for it: a
# wrong-but-confident label on a status page is worse than an ugly true one.
ROLE_NAMES: dict[tuple[str, ...], str] = {
    lane.queue_key: lane.display_name for lane in LANES
}


def role_display_name(queues: tuple[str, ...]) -> str:
    known = ROLE_NAMES.get(queues)
    if known:
        return known
    return "Worker (" + ", ".join(queues) + ")"


def _inspect_celery_sync() -> dict[tuple[str, ...], dict]:
    """celery inspect, grouped by queue set rather than by worker.

    Returns {queue_set: {"hostnames": [...], "active": int}}. Two replicas of
    one role collapse into one entry on purpose — the question is whether the
    role is being served, not how many containers exist.
    """
    from ..celery_app import celery as celery_app

    insp = celery_app.control.inspect(timeout=INSPECT_TIMEOUT_SECONDS)
    # TWO broadcasts, each waiting out its own timeout — see
    # INSPECT_ROUND_TRIPS, which the caller's budget is derived from. Adding a
    # third call here without updating that constant puts the wrapper back
    # under the work it is waiting for.
    active_queues = insp.active_queues() or {}
    active_tasks = insp.active() or {}

    grouped: dict[tuple[str, ...], dict] = {}
    for hostname, queues in active_queues.items():
        # Keyed on the LANE's queue set when the node name identifies one, so
        # a lane keeps the same roster row whether or not it is consuming.
        #
        # Grouping on the ACTIVE queues alone meant a lane at cap 0 — which
        # cancels its consumers — reported an empty set, landed under the key
        # `celery:`, and rendered as a phantom row named `Worker ()` while its
        # real row went stale beside it. Both symptoms on the operator's
        # screen, 2026-09-23, from this one line.
        #
        # Deriving the key from `lane.queue_key` rather than inventing a new
        # one keeps every existing row: it is the same string the lane already
        # had while it was running.
        lane = lane_for_node(hostname)
        key = lane.queue_key if lane else tuple(sorted({q["name"] for q in queues}))
        entry = grouped.setdefault(key, {"hostnames": [], "active": 0})
        entry["hostnames"].append(hostname)
        entry["active"] += len(active_tasks.get(hostname, []))
    for entry in grouped.values():
        entry["hostnames"].sort()
    return grouped


def touch_service_stmt(*, key: str, kind: str, display_name: str, details: dict):
    """The upsert that records a check-in, as a statement.

    Built here rather than inline so the async caller (an agent lease, over
    the API) and the sync one (the sizing sweep, in a celery task) run the
    SAME write. Two spellings of one upsert is the kind of duplication that
    stays correct right up until one of them gains a column.

    Upsert rather than read-modify-write: several web processes and several
    agents can be doing this at once, and the last writer is simply the most
    recent sighting. `first_seen_at` is deliberately NOT updated — it is the
    one field that answers "has this ever run", which the learned-roster design
    depends on.
    """
    stmt = pg_insert(ServiceSeen).values(
        key=key, kind=kind, display_name=display_name, details=details,
    )
    return stmt.on_conflict_do_update(
        index_elements=[ServiceSeen.key],
        set_={
            "kind": stmt.excluded.kind,
            "display_name": stmt.excluded.display_name,
            "details": stmt.excluded.details,
            "last_seen_at": func.now(),
        },
    )


async def touch_service(
    session: AsyncSession, *, key: str, kind: str, display_name: str, details: dict
) -> None:
    """Record that a part checked in just now."""
    await session.execute(touch_service_stmt(
        key=key, kind=kind, display_name=display_name, details=details,
    ))


def _roster_rows(grouped: dict[tuple[str, ...], dict]) -> list[dict]:
    """The `touch_service` arguments for everything that answered.

    Split from the write so the async and sync refreshes below share the
    mapping as well as the statement — what a roster row IS should not depend
    on which kind of session is writing it.
    """
    return [
        {
            "key": "celery:" + ",".join(queues),
            "kind": "celery",
            "display_name": role_display_name(queues),
            "details": {
                "queues": list(queues),
                "hostnames": entry["hostnames"],
                "replicas": len(entry["hostnames"]),
                "active": entry["active"],
            },
        }
        for queues, entry in grouped.items()
    ]


def refresh_celery_roster_sync(session) -> None:
    """The roster refresh, from the sizing sweep's sync session.

    ## Why the sweep owns this now

    It used to run on the request path, rate-limited to once every 20s by the
    newest celery row. So the roster only advanced while somebody had a
    browser open — the liveness of the workers was a function of whether
    anyone was looking at them, which is the observer-effect version of the
    bug this roster exists to prevent.

    Operator, 2026-09-23: *"there is a repull every time this page loads — is
    there a reason this info isn't being tracked in the background and stored
    in some way?"*

    Now a timer writes it and the page only reads. The cadence is
    `SWEEP_PERIOD_SECONDS`, and `api/system_health` asserts it leaves headroom
    under the staleness thresholds — because a sweep period and a stale
    threshold chosen in different files and never compared is exactly how the
    idle GPU agent came to read as stopped (lesson #4355).

    Never raises. A failure means the roster does not advance, and the rows
    going stale is then a TRUE report about a broker nobody can reach.
    """
    try:
        grouped = _inspect_celery_sync()
    except Exception:
        log.warning(
            "service roster: celery inspect failed; roster not refreshed",
            exc_info=True,
        )
        return
    for row in _roster_rows(grouped):
        session.execute(touch_service_stmt(**row))
    session.commit()


async def refresh_celery_roster(session: AsyncSession) -> None:
    """Inspect the broker and record what answered. Never raises.

    A failure here means the roster does not advance, and the rows going stale
    is then a TRUE report about a broker nobody can reach. Letting the
    exception out would instead break the health endpoint, which is the one
    thing that must keep answering when the stack is unwell.
    """
    try:
        grouped = await asyncio.wait_for(
            asyncio.to_thread(_inspect_celery_sync),
            timeout=(
                INSPECT_TIMEOUT_SECONDS * INSPECT_ROUND_TRIPS
                + INSPECT_SLACK_SECONDS
            ),
        )
    except Exception:
        log.warning("service roster: celery inspect failed; roster not refreshed", exc_info=True)
        return

    for row in _roster_rows(grouped):
        await touch_service(session, **row)