"""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 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 GPU-utilization-band controller. It grows the
# pool while the GPU has spare capacity (util below the low mark + VRAM headroom)
# and shrinks under saturation / memory pressure, then HOLDS while util sits in
# the band — so the worker count stays steady instead of flopping. Util is EWMA-
# smoothed and decisions are spaced out, so a single noisy nvidia-smi sample
# can't move it.
CONTROL_INTERVAL = 8.0     # seconds between samples
DECIDE_EVERY = 3           # only act every Nth sample (~24s) — stability
UTIL_LO = 70               # grow when smoothed util is below this (spare capacity)
UTIL_HI = 92               # shrink when above this (saturated)
VRAM_HI = 0.88             # shrink above this fraction of VRAM (memory pressure)
VRAM_GROW_MAX = 0.80       # don't grow past this VRAM
EWMA_ALPHA = 0.4           # util smoothing weight on the newest sample

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
        # 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()

    # --- 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):
        self._ctrl_stop.set()
        with self._lock:
            self._running = False
            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:
        with self._lock:
            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:
                jobs = self.client.lease(self.cfg.batch_size)
                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.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):
        """GPU-utilization-band controller (Auto mode). Hold the worker count
        steady while the GPU sits in a healthy util band; grow only when there's
        clear spare capacity (smoothed util below the low mark + VRAM headroom),
        shrink under saturation or memory pressure. Util is EWMA-smoothed and we
        only act every DECIDE_EVERY samples, so a noisy nvidia-smi reading can't
        make the pool flop — load stays consistent instead of probe/reverting
        every cycle (the old hill-climb's failure mode)."""
        from . import gpu as gpumod

        util_ewma = None     # smoothed GPU util%
        tick = 0             # samples since the last decision
        while not self._ctrl_stop.wait(CONTROL_INTERVAL):
            if not (self._running and self._auto):
                util_ewma, tick = None, 0
                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 (
                EWMA_ALPHA * util + (1 - EWMA_ALPHA) * 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 = 0
                continue

            tick += 1
            if tick < DECIDE_EVERY:                  # hold between decisions
                continue
            tick = 0

            t0 = self._target
            if util_ewma > UTIL_HI:                  # saturated → ease off
                self._apply_step(-1)
            elif util_ewma < UTIL_LO and vram < VRAM_GROW_MAX:
                self._apply_step(+1)                 # spare capacity → grow
            # else: util is in the band → HOLD (steady load, no flopping)

            if self._target != t0:
                log.info(
                    "autoscale: %d→%d workers (util~%d%% · vram %d%%)",
                    t0, self._target, round(util_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) -> 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:
            data = self.client.fetch_image(job["image_url"])
            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))]

            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)
                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.client.submit_embedding(job["job_id"], vec, embed_version)
                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}"

            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):
                        tmpl["siglip_embedding"] = vec
                        tmpl["embedding_version"] = embed_version
                        regions.append(tmpl)
            self.client.submit(job["job_id"], regions, replace_kinds)
            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)