"""The worker lanes: what they are, and how many slots each may be given.

Milestone 422 step 1. This module is the ONE place that knows the lane set;
`models/worker_lane.py` holds only what the operator can change about them.

## Why the queues are here and not in the table

A lane's queue set is not a preference — it is decided by `celery_app.py`'s
`task_routes`, which is what puts a backup on `maintenance_long` and a
thumbnail on `thumbnail`. An operator cannot move a task to another lane, so
storing the queues as settings would create a row that can disagree with the
routing table, and nothing would notice until a queue had no consumer.

So: queues and display names are code, slots and caps are data. The table
stores three numbers and a flag, and nothing that could contradict celery.

This also collapses a duplicate rather than adding one.
`service_roster.ROLE_NAMES` was a second copy of "queue set -> the name an
operator recognises", and it had already drifted: `maintenance_long` is a
live lane with four task routes pointing at it, and the roster did not know
its name, so the System tab rendered it as `Worker (maintenance_long)`. That
map is now derived from `LANES` below, so a lane added here is named
everywhere at once.

## Why the ceiling is derived rather than configured

Consolidating the stack into one container (step 5) widens the OOM blast
radius: today an ml-worker that exhausts memory is killed by Docker on its
own, and web keeps serving. In one container the kernel picks a victim from
the whole cgroup, and it may pick hypercorn — so a tagging task can take the
UI down with it, on exactly the modest hardware least able to spare the
memory.

Operator, 2026-09-22: *"ram isn't an issue for me but some users might run
this on weaker hardware and I don't want it to kill their servers."*

So the maximum is computed from what the container actually has, and the
operator's own `slots_cap` must fit under it. Three numbers, not two, and the
ordering is the point:

    slots  <=  slots_cap  <=  derived_ceiling
    (live)     (operator)     (this module)

The operator can always lower their cap. They cannot raise it past what the
box can hold. The derived ceiling is never stored — a row that outlived a
change in container limits must not carry a stale one.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from pathlib import Path

log = logging.getLogger(__name__)

# --- the lanes ---------------------------------------------------------------


GIB = 1024 ** 3


@dataclass(frozen=True)
class ModelRequirement:
    """A model a lane must download before it can do anything.

    Surfaced to the UI so the operator is told WHICH model, how big, and what
    it costs to hold — before they turn the lane on, not after a multi-GB
    download has already started. The lane is optional and its cost is not
    obvious from its name, which is the whole reason this is structured data
    rather than a sentence in a component.

    `measured=False` means the numbers are ESTIMATES and the UI must say so.
    They come from the checkpoint's parameter count and dtype, not from a
    build — and a number presented as fact decides whether someone's server
    survives, so it is labelled rather than rounded confidently.
    """

    # The Hugging Face repo id, which is the honest answer to "which model".
    repo: str
    # Roughly what the download costs, for the operator's bandwidth and disk.
    approx_download_bytes: int
    # Roughly what ONE slot holds while running. Prefork forks a child per
    # slot and each loads its own copy, so this multiplies.
    approx_resident_bytes: int
    measured: bool = False


# SigLIP so400m — the only model FabledCurator itself downloads.
#
# What it is for, which is NOT obvious from the lane's name: it produces the
# image embeddings that back similarity search, duplicate grouping and the
# tag heads. WD14 tagging is the GPU AGENT's job, not this lane's — the
# comment in celery_app.py naming both is stale since B3 (#1238), when the
# agent took over and this lane was left as the CPU embed fallback for stacks
# running no agent at all (see MLSettings.cpu_embed_enabled).
#
# Both numbers are ESTIMATES, derived from the checkpoint rather than from a
# build: ~877M parameters at fp32 is ~3.5GB of weights, and holding them plus
# activations and the torch runtime is what the resident figure covers. They
# err high. Replace them with measurements — download the repo and read its
# size; run one embed and read the worker child's VmHWM — and set
# `measured=True` when you do.
SIGLIP_MODEL = ModelRequirement(
    repo="google/siglip-so400m-patch14-384",
    approx_download_bytes=3_500_000_000,
    approx_resident_bytes=4 * GIB,
    measured=False,
)


@dataclass(frozen=True)
class Lane:
    """A worker lane. `name` is the stable key the settings row is keyed on.

    Keyed on a lane NAME rather than a container hostname for the reason
    `models/service_seen.py` gives at length: celery's worker names here are
    `celery@<container id>` and are minted fresh on every deploy, so anything
    keyed on them records a death and a birth every time the stack updates.
    """

    name: str
    display_name: str
    queues: tuple[str, ...]
    # Which `entrypoint.sh` role starts this lane. NOT always the lane name:
    # `maintenance_long` is the plain `worker` role pointed at a different
    # queue, exactly as docker-compose starts it today (`command: ["worker"]`
    # with CELERY_QUEUES=maintenance_long). Recorded here so the generated
    # supervisord config and the compose file cannot disagree about it.
    entrypoint_role: str
    default_slots: int
    # The cap a lane STARTS with, which is not the ceiling. Set low enough
    # that raising slots within it is an ordinary adjustment, and raising the
    # cap itself is a deliberate act — a cap that begins at the ceiling is a
    # rubber stamp and protects nobody.
    default_slots_cap: int
    default_enabled: bool
    # True when a slot costs a copy of the ML model rather than just a process.
    # The only lane whose ceiling is decided by memory instead of by cores.
    memory_bound: bool = False
    # Models this lane downloads the first time it is enabled. Empty for every
    # lane that needs none, which is how the UI knows whether to warn at all.
    models: tuple[ModelRequirement, ...] = ()
    # An optional lane is one the product works without. Shown as such, so
    # nobody turns on a multi-GB download believing it is required.
    optional: bool = False

    @property
    def queue_key(self) -> tuple[str, ...]:
        """The sorted queue set, which is how `service_seen` identifies a
        running worker. The join between what is configured here and what
        `celery inspect` reports."""
        return tuple(sorted(self.queues))


# Defaults are ONE OF EACH, with ML off — operator, 2026-09-22: *"that
# starting value should be one of each."* Deliberately far below the
# operator's own production numbers (worker 8, ml 2), which are tuned for
# their hardware and are not a sane first boot for a stranger.
#
# ML ships disabled because enabling it is what triggers the SigLIP download
# (milestone 422 step 6) — rule 164 allows a feature that needs a fetch only
# when it is "optional and clearly off", and off-by-default is also what keeps
# a small box from loading a multi-GB model it was never asked to load.
LANES: tuple[Lane, ...] = (
    Lane(
        name="worker",
        display_name="Worker",
        queues=("default", "import", "thumbnail", "download"),
        entrypoint_role="worker",
        default_slots=1,
        default_slots_cap=4,
        default_enabled=True,
    ),
    Lane(
        name="scheduler",
        display_name="Scheduler",
        queues=("maintenance", "scan"),
        entrypoint_role="scheduler",
        default_slots=1,
        default_slots_cap=2,
        default_enabled=True,
    ),
    Lane(
        name="maintenance_long",
        display_name="Long maintenance",
        queues=("maintenance_long",),
        entrypoint_role="worker",
        default_slots=1,
        default_slots_cap=2,
        default_enabled=True,
    ),
    Lane(
        name="ml",
        display_name="ML tagging",
        queues=("ml",),
        entrypoint_role="ml-worker",
        default_slots=0,
        default_slots_cap=1,
        default_enabled=False,
        memory_bound=True,
        models=(SIGLIP_MODEL,),
        optional=True,
    ),
)

LANES_BY_NAME: dict[str, Lane] = {lane.name: lane for lane in LANES}
LANES_BY_QUEUE_KEY: dict[tuple[str, ...], Lane] = {
    lane.queue_key: lane for lane in LANES
}


# --- what the container actually has -----------------------------------------

# cgroup v2 first, then v1. A container started without an explicit memory
# limit reports "max" on v2 and a sentinel near 2**63 on v1; both mean "no
# limit", and the answer then is the host's RAM.
_CGROUP_V2_MEMORY = Path("/sys/fs/cgroup/memory.max")
_CGROUP_V1_MEMORY = Path("/sys/fs/cgroup/memory/memory.limit_in_bytes")
_CGROUP_V2_CPU = Path("/sys/fs/cgroup/cpu.max")
_CGROUP_V1_CPU_QUOTA = Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
_CGROUP_V1_CPU_PERIOD = Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us")

# A v1 "unlimited" is PAGE_SIZE-aligned LONG_MAX, not a round number, so it is
# recognised by magnitude rather than by equality. Anything claiming more than
# a petabyte is a sentinel, not a machine.
_UNLIMITED_ABOVE = 1 << 50

# DERIVED from the model requirement above, never restated. The ceiling and
# the number shown to the operator before they enable the lane have to be the
# same figure, or the UI promises something the cap will then refuse.
ML_BYTES_PER_SLOT = SIGLIP_MODEL.approx_resident_bytes

# Held back for hypercorn and the non-ML lanes before any ML slot is offered.
# In the consolidated container these share one cgroup with ML, and they are
# the processes an OOM kill must not take (see the module docstring).
RESERVED_BYTES = 2 * GIB

# The floor a cores-derived ceiling never goes below. A single-core box still
# needs to be able to run its lanes; the ceiling exists to stop absurd values,
# not to make a small machine unusable.
MIN_CEILING = 1

# What an unreadable limit yields. Low rather than unlimited, on purpose: not
# knowing how much memory there is must never read as "plenty". An unswept
# absence is not a verdict.
UNKNOWN_CEILING = 1


def _read_int(path: Path) -> int | None:
    try:
        raw = path.read_text().strip()
    except OSError:
        return None
    if raw == "max":
        return None
    try:
        return int(raw)
    except ValueError:
        return None


def container_memory_bytes() -> int | None:
    """The memory this container may use, or None when it cannot be read.

    None means UNKNOWN, never UNLIMITED. Every caller must treat it as the
    conservative case — the whole point of the ceiling is to protect a machine
    whose size we are unsure of.
    """
    for path in (_CGROUP_V2_MEMORY, _CGROUP_V1_MEMORY):
        value = _read_int(path)
        if value is not None and value < _UNLIMITED_ABOVE:
            return value
        if value is not None:
            # A sentinel: the cgroup exists but sets no limit, so the real
            # bound is the host's.
            break
    try:
        return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")
    except (ValueError, OSError, AttributeError):
        return None


def container_cpu_count() -> int | None:
    """Effective cores, honouring a cgroup CPU quota.

    `os.cpu_count()` reports the HOST's cores from inside a container, so a
    quota of 2.0 on a 32-core host would otherwise offer 32 slots. The
    operator's own stack sets `cpus: '4.0'` on ml-worker, so this is a real
    configuration here and not a hypothetical.
    """
    quota: float | None = None
    try:
        raw = _CGROUP_V2_CPU.read_text().strip().split()
        if raw and raw[0] != "max":
            quota = int(raw[0]) / int(raw[1])
    except (OSError, ValueError, IndexError, ZeroDivisionError):
        pass
    if quota is None:
        q = _read_int(_CGROUP_V1_CPU_QUOTA)
        p = _read_int(_CGROUP_V1_CPU_PERIOD)
        if q is not None and p and q > 0:
            quota = q / p
    if quota is not None and quota > 0:
        return max(1, int(quota))
    return os.cpu_count()


def derived_ceiling(lane: Lane) -> int:
    """The most slots `lane` may be given on this container.

    Never stored. Recomputed on every read so a container whose limits changed
    is bounded by what it has NOW rather than by what it had when its row was
    written.
    """
    if lane.memory_bound:
        total = container_memory_bytes()
        if total is None:
            log.warning(
                "worker_lanes: cannot read a memory limit; capping %s at %d",
                lane.name, UNKNOWN_CEILING,
            )
            return UNKNOWN_CEILING
        usable = total - RESERVED_BYTES
        if usable < ML_BYTES_PER_SLOT:
            # Honestly zero. A box that cannot hold one model alongside the web
            # process must be told it cannot run tagging, not sold a slot that
            # will OOM the container the first time it is used.
            return 0
        return int(usable // ML_BYTES_PER_SLOT)

    cores = container_cpu_count()
    if cores is None:
        return UNKNOWN_CEILING
    return max(MIN_CEILING, cores)


def ceilings() -> dict[str, int]:
    """Every lane's ceiling, for the settings API and the UI."""
    return {lane.name: derived_ceiling(lane) for lane in LANES}