"""The lease → fetch → detect+embed → submit loop, run by a pool of worker
slots whose count is tunable live from the UI.

Each slot is an independent loop (its own leases; the server's SKIP-LOCKED lease
keeps them from colliding). More slots = more GPU load + throughput; the model is
loaded once and shared, so slots add concurrent inference, not N× model VRAM.
That's the dial the operator turns to trade desktop responsiveness for speed.

Stop (or shrinking the pool) RELEASES a slot's still-leased jobs immediately so
orphaned work is re-picked at once rather than waiting out the lease.
"""
import logging
import threading
import time

import numpy as np
import requests

from . import media, models
from .client import FcClient
from .config import Config
from .crops import crop_region

# Cap on the lease-retry backoff: when curator is unreachable (e.g. you redeploy
# it while away), each slot retries leasing with exponential backoff up to this
# many seconds, then resumes within this window once the server is back — no
# restart needed.
MAX_BACKOFF_SECONDS = 60.0


def _is_transient(exc: requests.RequestException) -> bool:
    """A server/transport problem (wait it out) vs a job-specific fault (fail it).
    No response → connection refused/timeout → curator is down → transient. With
    a response: 5xx, auth (401/403, e.g. a token blip on redeploy), 408/409/429
    (timeout / our lease reclaimed / rate-limited) are all 'not this job's fault'.
    A specific 4xx like 404 (image gone) / 400 IS the job's fault → fail it."""
    resp = getattr(exc, "response", None)
    if resp is None:
        return True
    return resp.status_code >= 500 or resp.status_code in (401, 403, 408, 409, 429)

# Generous cap: the pipeline is usually I/O-bound (downloading + decoding images
# over HTTP), so the GPU stays underused until many workers overlap that I/O.
# Push it up while watching the GPU util + VRAM in the UI.
MAX_CONCURRENCY = 32

# Fallbacks only — the server ANNOUNCES the embedding model (name + version) in
# the lease so the agent stays model-agnostic and in lock-step with the space
# the heads were trained in. These cover an older server that doesn't send them.
DEFAULT_EMBED_MODEL = "google/siglip-so400m-patch14-384"
DEFAULT_EMBED_VERSION = "siglip-so400m-patch14-384"

# Autoscaler (when Auto is on): a SMOOTHED, throughput-aware climb that SETTLES.
# Raw GPU util swings wildly (a batched embed pegs it ~99%, then image decode/IO
# drops it ~0%), so a single reading is meaningless — util is sampled often and
# EWMA-smoothed. Each decision (spaced ~24s) grows the pool by one only while
# doing so keeps lifting *throughput* (jobs/s, also smoothed); when a grow stops
# helping the pool is IO/CPU/curator-bound, so it backs off one and SETTLES,
# holding there before an occasional re-probe. This finds the worker count that
# maximises real work — instead of flopping every cycle, or growing forever
# because util never reaches a fixed threshold on an IO-bound load.
CONTROL_INTERVAL = 2.0        # util sampling cadence (seconds)
SAMPLES_PER_DECISION = 12     # decide ~every 24s (12 × 2s) on averaged signals
UTIL_HI = 92                  # smoothed util above this = saturated → shrink
UTIL_START = 85               # only begin a climb when smoothed util is below this
VRAM_HI = 0.88                # shrink above this fraction of VRAM (memory pressure)
VRAM_GROW_MAX = 0.80          # don't grow past this VRAM
UTIL_ALPHA = 0.25             # util EWMA weight on the newest sample (smoother)
TPUT_ALPHA = 0.5              # throughput EWMA weight
TPUT_MARGIN = 0.08            # a grow must lift smoothed jobs/s by this to "help"
REPROBE_TICKS = 8             # decisions to hold after settling before re-probing

# How often to log the per-stage timing breakdown (lease/download/decode/gpu/
# submit) so the operator can see where a job's wall-clock actually goes — the
# data that decides whether a download/compute split is worth building.
STATS_INTERVAL = 30.0

# The queue snapshot exists only to populate the UI's counts, so it's polled
# lazily — only while a browser is actually watching (a /status hit in the last
# UI_IDLE_GRACE seconds), and not on a tight loop. The work loop's own lease/
# submit calls are the real "is curator up?" signal; nothing polls just to poll.
QUEUE_POLL_INTERVAL = 5.0
UI_IDLE_GRACE = 20.0

log = logging.getLogger("fc_agent.worker")


class _Slot:
    """One worker loop. `inflight` = jobs leased but not yet processed, so a
    graceful stop can hand them back."""
    __slots__ = ("stop", "inflight")

    def __init__(self):
        self.stop = threading.Event()
        self.inflight: list[int] = []


class Worker:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
        self._lock = threading.Lock()
        self._running = False
        self._target = max(1, min(MAX_CONCURRENCY, cfg.concurrency))
        self._auto = bool(cfg.auto_scale)      # autoscale worker count
        self._ctrl_stop = threading.Event()
        self._ctrl_thread: threading.Thread | None = None
        self._slots: list[_Slot] = []
        self.processed = 0
        self.errors = 0
        self.transient = 0        # jobs handed back due to a server outage (NOT
        #                           failed) — the "waiting out curator" counter
        self._active = 0          # slots currently mid-image
        self._util_smooth: float | None = None   # EWMA GPU util (set by control loop)
        # Curator queue snapshot, refreshed by a background poller so the UI
        # /status read is instant — never an inline curator HTTP call (which
        # stalls the whole status view when curator is busy).
        self._queue: dict | None = None
        self._ui_seen = 0.0        # monotonic time of the last UI /status hit
        self._queue_thread = threading.Thread(target=self._queue_poll_loop, daemon=True)
        self._queue_thread.start()
        # Per-stage timing: stage -> [sum_seconds, count], summarised to the log
        # every STATS_INTERVAL so we can see where wall-clock goes per job.
        self._timing: dict[str, list[float]] = {}
        self._timing_lock = threading.Lock()
        self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True)
        self._stats_thread.start()
        # The crop embedder (SigLIP-family) is built lazily on the first job that
        # needs it, from the model the server announces — one shared instance.
        self._embedder = None
        self._embedder_lock = threading.Lock()
        # Region proposers (extra YOLO detectors) — lazily built once, shared.
        self._proposers = None
        self._proposers_lock = threading.Lock()

    def _queue_poll_loop(self):
        """Refresh the curator queue snapshot so /status is a pure in-memory read
        — but ONLY while the UI is being watched (a recent /status hit). No
        browser open → no polling; the work loop is curator's only visitor.
        Errors just leave the last snapshot (or None) — never blocks the UI."""
        while True:
            if time.monotonic() - self._ui_seen <= UI_IDLE_GRACE:
                try:
                    self._queue = self.client.queue_status()
                except Exception:
                    self._queue = None
            time.sleep(QUEUE_POLL_INTERVAL)

    def note_ui(self) -> None:
        """The UI polled /status — keep the queue snapshot warm for a while."""
        self._ui_seen = time.monotonic()

    def latest_queue(self) -> dict | None:
        return self._queue

    def util_smooth(self) -> float | None:
        return self._util_smooth

    def _record(self, stage: str, seconds: float) -> None:
        with self._timing_lock:
            s = self._timing.get(stage)
            if s is None:
                self._timing[stage] = [seconds, 1]
            else:
                s[0] += seconds
                s[1] += 1

    def _stats_loop(self) -> None:
        """Log a per-stage timing breakdown every STATS_INTERVAL (only when there
        was work), so the operator can see the download/decode/gpu/submit split."""
        while True:
            time.sleep(STATS_INTERVAL)
            with self._timing_lock:
                snap = {k: (v[0], v[1]) for k, v in self._timing.items() if v[1]}
                self._timing = {}
            if not snap:
                continue
            order = ["lease", "download", "decode", "gpu", "submit"]
            parts = [
                f"{st} {1000 * snap[st][0] / snap[st][1]:.0f}ms"
                for st in order if st in snap
            ]
            jobs = (snap.get("gpu") or snap.get("download") or (0, 0))[1]
            # Per-job wall time across the compute path (lease is per-batch, so
            # it's shown separately above, not folded into this figure).
            per_job = sum(
                snap[st][0] for st in ("download", "decode", "gpu", "submit")
                if st in snap
            )
            pj_ms = 1000 * per_job / jobs if jobs else 0
            log.info("timing/%ds — %s | wall/job %.0fms (%d jobs)",
                     int(STATS_INTERVAL), " · ".join(parts), pj_ms, jobs)

    # --- control -----------------------------------------------------------
    def start(self):
        with self._lock:
            self._running = True
            self._reconcile_locked()
        # (Re)start the autoscaler control loop.
        if self._ctrl_thread is None or not self._ctrl_thread.is_alive():
            self._ctrl_stop.clear()
            self._ctrl_thread = threading.Thread(target=self._control_loop, daemon=True)
            self._ctrl_thread.start()

    def stop(self):
        # Flip the flag FIRST (atomic bool), before any lock, so /status and the
        # worker loops observe "stopped" immediately even if _lock is momentarily
        # held — the state can never lag behind the click.
        self._running = False
        self._ctrl_stop.set()
        with self._lock:
            slots, self._slots = self._slots, []
        for s in slots:
            s.stop.set()          # each slot releases its inflight on exit

    def set_auto(self, on: bool):
        with self._lock:
            self._auto = bool(on)

    def set_concurrency(self, n: int):
        # A manual set is an override → leave Auto.
        with self._lock:
            self._auto = False
            self._target = max(1, min(MAX_CONCURRENCY, int(n)))
            if self._running:
                self._reconcile_locked()

    def _apply_step(self, delta: int) -> bool:
        """Nudge the target by delta (bounded). Returns True if it changed."""
        with self._lock:
            new = max(1, min(MAX_CONCURRENCY, self._target + delta))
            if new == self._target:
                return False
            self._target = new
            if self._running:
                self._reconcile_locked()
            return True

    def _reconcile_locked(self):
        while len(self._slots) < self._target:
            slot = _Slot()
            self._slots.append(slot)
            threading.Thread(target=self._loop, args=(slot,), daemon=True).start()
        while len(self._slots) > self._target:
            self._slots.pop().stop.set()

    def status(self) -> dict:
        # Lock-free on purpose: these are plain int / bool reads (atomic under the
        # GIL) and this backs the UI poll — it must NEVER be able to block behind
        # a worker holding _lock, or the whole status view freezes.
        return {
            "state": "running" if self._running else "stopped",
            "concurrency": self._target,
            "max_concurrency": MAX_CONCURRENCY,
            "auto": self._auto,
            "workers": len(self._slots),
            "active": self._active,
            "processed": self.processed,
            "errors": self.errors,
            "transient": self.transient,
        }

    def _bump(self, *, processed=0, errors=0, active=0, transient=0):
        with self._lock:
            self.processed += processed
            self.errors += errors
            self.transient += transient
            self._active += active

    # --- per-slot loop -----------------------------------------------------
    def _loop(self, slot: _Slot):
        backoff = self.cfg.poll_idle_seconds
        while not slot.stop.is_set() and self._running:
            try:
                _t = time.monotonic()
                jobs = self.client.lease(self.cfg.batch_size)
                self._record("lease", time.monotonic() - _t)
                backoff = self.cfg.poll_idle_seconds      # server answered → reset
            except Exception:
                # curator unreachable (redeploy, network drop): wait it out with
                # exponential backoff, capped — resume on our own when it returns.
                self._interruptible_sleep(slot, backoff)
                backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
                continue
            if not jobs:
                self._interruptible_sleep(slot, self.cfg.poll_idle_seconds)
                continue
            slot.inflight = [j["job_id"] for j in jobs]
            for job in jobs:
                if slot.stop.is_set() or not self._running:
                    break
                ok = self._process(job, slot)
                slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
                if not ok:
                    # Server went away mid-batch: hand the rest back (best effort)
                    # and back off instead of hammering a recovering server or
                    # burning the jobs' attempt budgets on fail().
                    if slot.inflight:
                        self.client.release(slot.inflight)
                        slot.inflight = []
                    self._interruptible_sleep(slot, backoff)
                    backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
                    break
                if slot.inflight:
                    self.client.heartbeat(slot.inflight)
        # Graceful hand-back of anything leased but not processed.
        if slot.inflight:
            self.client.release(slot.inflight)
            slot.inflight = []

    def _interruptible_sleep(self, slot: _Slot, seconds: float):
        """Sleep, but wake immediately if the slot is told to stop — so a Stop or
        a pool-shrink doesn't hang for a full backoff window."""
        slot.stop.wait(timeout=seconds)

    # --- autoscaler --------------------------------------------------------
    def _control_loop(self):
        """Smoothed, throughput-aware climb that settles (Auto mode). Samples GPU
        util often and EWMA-smooths it (raw util swings 0↔99 between a batched
        embed and the IO/decode around it, so one reading is noise). Every
        SAMPLES_PER_DECISION ticks it makes ONE move: grow by one while each grow
        keeps lifting smoothed throughput; when a grow stops helping (IO/CPU/
        curator-bound) back off one and SETTLE, holding before an occasional
        re-probe. Memory pressure and saturation always shrink immediately."""
        from . import gpu as gpumod

        util_ewma: float | None = None
        tput_ewma: float | None = None
        prev_p, prev_t = self.processed, time.monotonic()
        tick = 0
        settled = False
        grew_last = False        # did the previous decision grow the pool?
        tput_before = 0.0        # smoothed jobs/s at the count before that grow
        hold = 0                 # decisions left to hold while settled
        while not self._ctrl_stop.wait(CONTROL_INTERVAL):
            if not (self._running and self._auto):
                util_ewma = tput_ewma = None
                prev_p, prev_t = self.processed, time.monotonic()
                tick = 0
                settled = grew_last = False
                hold = 0
                self._util_smooth = None
                continue

            g = gpumod.read_gpu() or {}
            mt = g.get("mem_total_mb") or 0
            vram = (g.get("mem_used_mb", 0) / mt) if mt else 0.0
            util = g.get("util_pct", 0) or 0
            util_ewma = util if util_ewma is None else (
                UTIL_ALPHA * util + (1 - UTIL_ALPHA) * util_ewma
            )
            self._util_smooth = util_ewma

            # Memory pressure overrides the cadence — react immediately.
            if vram >= VRAM_HI:
                if self._apply_step(-1):
                    log.info("autoscale: -1 → %d workers (vram %d%% — memory pressure)",
                             self._target, round(vram * 100))
                tick, settled, grew_last, hold = 0, True, False, REPROBE_TICKS
                continue

            tick += 1
            if tick < SAMPLES_PER_DECISION:
                continue
            tick = 0

            now = time.monotonic()
            inst = (self.processed - prev_p) / max(1e-3, now - prev_t)
            prev_p, prev_t = self.processed, now
            tput_ewma = inst if tput_ewma is None else (
                TPUT_ALPHA * inst + (1 - TPUT_ALPHA) * tput_ewma
            )

            t0 = self._target
            if util_ewma > UTIL_HI:                       # saturated → ease off
                self._apply_step(-1)
                settled, grew_last, hold = True, False, REPROBE_TICKS
            elif settled:
                hold -= 1
                if hold <= 0:                             # re-probe: try one grow
                    if util_ewma < UTIL_START and vram < VRAM_GROW_MAX:
                        tput_before = tput_ewma
                        grew_last = self._apply_step(+1)
                        settled = not grew_last
                    else:
                        hold = REPROBE_TICKS              # still no room → keep holding
            elif grew_last:
                if tput_ewma > tput_before * (1 + TPUT_MARGIN):   # the grow helped
                    tput_before = tput_ewma
                    if util_ewma < UTIL_START and vram < VRAM_GROW_MAX:
                        grew_last = self._apply_step(+1)
                        settled = not grew_last
                    else:
                        settled, grew_last, hold = True, False, REPROBE_TICKS
                else:                                     # overshot → back off + settle
                    self._apply_step(-1)
                    settled, grew_last, hold = True, False, REPROBE_TICKS
            elif util_ewma < UTIL_START and vram < VRAM_GROW_MAX:  # start a climb
                tput_before = tput_ewma
                grew_last = self._apply_step(+1)
                settled = not grew_last
            else:
                settled, hold = True, REPROBE_TICKS       # nothing to do → settle

            if self._target != t0:
                log.info("autoscale: %d→%d workers (util~%d%% · %.2f j/s · vram %d%%)",
                         t0, self._target, round(util_ewma), tput_ewma, round(vram * 100))

    def _ensure_embedder(self, model_name: str):
        if self._embedder is not None:
            return self._embedder
        with self._embedder_lock:
            if self._embedder is None:
                from .embedder import CropEmbedder
                self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
        return self._embedder

    def _ensure_proposers(self):
        if self._proposers is not None:
            return self._proposers
        with self._proposers_lock:
            if self._proposers is None:
                from .detectors import Proposers
                self._proposers = Proposers(self.cfg)
        return self._proposers

    def _process(self, job: dict, slot: _Slot) -> bool:
        """Process one job. Returns True when handled (completed, or hard-failed
        because the job itself is bad) and False on a TRANSPORT error (curator
        unreachable / 5xx / our lease was reclaimed mid-flight) — which is not
        the job's fault, so the caller backs off and the job is left to be
        re-leased rather than fail()ed into its attempt budget."""
        self._bump(active=1)
        try:
            _t = time.monotonic()
            data = self.client.fetch_image(job["image_url"])
            self._record("download", time.monotonic() - _t)

            _t = time.monotonic()
            if media.is_video(job.get("mime", "")):
                frames = media.sample_frames(
                    data, job.get("frame_interval_seconds", 4.0),
                    job.get("max_frames", 64),
                ) or [(None, media.load_image(data))]
            else:
                frames = [(None, media.load_image(data))]
            self._record("decode", time.monotonic() - _t)

            # Stop/shrink checkpoint: bail BEFORE the expensive detect+embed so a
            # Stop finishes promptly instead of waiting out heavy GPU work. Hand
            # the job back to pending for another agent.
            if not self._running or slot.stop.is_set():
                self.client.release([job["job_id"]])
                return True

            task = job.get("task") or "ccip"
            embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION
            model_name = (
                self.cfg.embed_model_override
                or job.get("embed_model_name")
                or DEFAULT_EMBED_MODEL
            )

            # 'embed' = WHOLE-IMAGE SigLIP embedding (re-embed the library under a
            # new model, #1190) → image_record.siglip_embedding. Mean-pool video
            # frames, matching the server's tag_and_embed. No regions.
            if task == "embed":
                embedder = self._ensure_embedder(model_name)   # one-time model load
                _t = time.monotonic()
                vecs = [embedder.embed(frame) for _, frame in frames]
                if len(vecs) > 1:
                    vec = np.mean(
                        np.asarray(vecs, dtype=np.float32), axis=0
                    ).tolist()
                else:
                    vec = vecs[0]
                self._record("gpu", time.monotonic() - _t)
                _t = time.monotonic()
                self.client.submit_embedding(job["job_id"], vec, embed_version)
                self._record("submit", time.monotonic() - _t)
                self._bump(processed=1)
                return True

            # task picks what to produce per crop:
            #   'siglip' (backfill existing images) → concept (SigLIP) regions
            #     ONLY, so it never churns their figure/CCIP regions or the
            #     character-reference cache.
            #   'ccip' / 'both' (a new image's first pass) → figure (CCIP) AND
            #     concept (SigLIP) in one go, off the same crop.
            want_ccip = task in ("ccip", "both")
            want_siglip = task in ("ccip", "siglip", "both")
            replace_kinds = (
                ["concept", "panel"] if task == "siglip"
                else ["figure", "face", "concept", "panel"]
            )
            embedder = self._ensure_embedder(model_name) if want_siglip else None
            proposers = self._ensure_proposers()

            regions = []
            ccip_ev = self.cfg.ccip_model or "ccip-default"
            dv = f"person-{self.cfg.detector_level}"

            _t_gpu = time.monotonic()   # detect + CCIP + batched embed = "gpu"
            for t, frame in frames:
                # FIGURE boxes: imgutils detect_person ∪ general COCO person,
                # NMS-merged → CCIP identity (+ a concept crop). Covers anime +
                # Western/realistic figures.
                base = models.detect_figures(frame, self.cfg.detector_level)
                figs = proposers.figures(frame, base)
                if not figs:
                    figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")]  # whole-frame fallback

                # Collect every crop that needs a SigLIP embedding, then embed
                # them in ONE batched forward pass (huge GPU-util + throughput
                # win vs one forward per crop). CCIP runs per figure inline.
                pending = []   # (crop, region-template-without-embedding)
                for bbox, score, _label in figs:
                    crop = crop_region(frame, bbox)
                    if crop is None:
                        continue
                    if want_ccip:
                        regions.append({
                            "kind": "figure", "bbox": list(bbox), "frame_time": t,
                            "score": score,
                            "ccip_embedding": models.ccip_vector(
                                crop, self.cfg.ccip_model or None
                            ),
                            "embedding_version": ccip_ev, "detector_version": dv,
                        })
                    if want_siglip:
                        pending.append((crop, {
                            "kind": "concept", "bbox": list(bbox), "frame_time": t,
                            "score": score, "detector_version": dv,
                        }))
                if not want_siglip:
                    continue
                # ANATOMY components (booru_yolo) + PANELS → concept/panel crops.
                for bbox, score, label in proposers.components(frame):
                    crop = crop_region(frame, bbox)
                    if crop is not None:
                        pending.append((crop, {
                            "kind": "concept", "bbox": list(bbox), "frame_time": t,
                            "score": score, "detector_version": f"booru:{label}",
                        }))
                for bbox, score, _label in proposers.panels(frame):
                    crop = crop_region(frame, bbox)
                    if crop is not None:
                        pending.append((crop, {
                            "kind": "panel", "bbox": list(bbox), "frame_time": t,
                            "score": score, "detector_version": "panel",
                        }))
                if pending:
                    vecs = embedder.embed_batch([c for c, _ in pending])
                    for (_c, tmpl), vec in zip(pending, vecs, strict=True):
                        tmpl["siglip_embedding"] = vec
                        tmpl["embedding_version"] = embed_version
                        regions.append(tmpl)
            self._record("gpu", time.monotonic() - _t_gpu)
            _t = time.monotonic()
            self.client.submit(job["job_id"], regions, replace_kinds)
            self._record("submit", time.monotonic() - _t)
            self._bump(processed=1)
            return True
        except requests.RequestException as exc:
            if _is_transient(exc):
                # curator down/redeploying, a 5xx, or our lease was reclaimed
                # while we worked. NOT the job's fault — hand it back (best
                # effort; no-ops if the server is still down, then the server's
                # orphan-recovery reclaims it) and signal the loop to wait.
                self._bump(transient=1)
                log.info("curator unreachable — released job %s, backing off",
                         job.get("job_id"))
                self.client.release([job["job_id"]])
                return False
            # A job-specific HTTP fault (404 image gone, 400) → fail it so it
            # doesn't re-lease forever.
            self._bump(errors=1)
            log.warning("job %s (image %s) failed: %s",
                        job.get("job_id"), job.get("image_id"), str(exc)[:200])
            self.client.fail(job["job_id"], str(exc)[:500])
            return True
        except Exception as exc:   # noqa: BLE001 — a genuine job fault: report it
            self._bump(errors=1)
            log.warning("job %s (image %s) failed: %s",
                        job.get("job_id"), job.get("image_id"), str(exc)[:200])
            self.client.fail(job["job_id"], str(exc)[:500])
            return True
        finally:
            self._bump(active=-1)