/**
 * Probe implementations for BackendStatusRegistry.
 *
 * Two upstream shapes are supported in Phase B:
 *
 * 1. **Direct workers** (llama-server compatible). We hit `<endpoint>/slots`
 *    and derive busySlots / totalSlots / loadedModel from the slot array.
 *    Recent llama-server builds let operators disable /slots via
 *    `--no-slots`; in that case the endpoint returns 404/405/501 and we
 *    fall back to `/health` to determine `online` only.
 *
 * 2. **Proxy workers** (LiteLLM Proxy). We hit `<endpoint>/health` and
 *    convert each healthy_endpoints / unhealthy_endpoints entry into a
 *    per-deployment NodeStatus. Slot / model info isn't currently
 *    surfaced by LiteLLM's /health; v1 reports `online` + cache hits and
 *    leaves the rest null. Phase C can switch to `/metrics` or
 *    `/model/info` once they're wired through the proxy auth surface.
 *
 * Both probes:
 *  - apply a 3s per-request timeout via AbortController (cluster-wide
 *    hangs would otherwise wedge the entire registry tick),
 *  - send the worker's apiKey as a Bearer header so team-scoped tokens
 *    work,
 *  - normalize the endpoint by stripping trailing slashes and stripping
 *    a single trailing `/v1` segment (the worker config typically points
 *    at the OpenAI-shaped `…/v1` base; llama-server's `/slots` lives at
 *    the server root, not under /v1).
 *
 * Long-form rationale lives in
 * docs/superpowers/specs/2026-05-18-multi-team-gpu-pool-and-node-status-design.md (Phase B).
 */

import type { WorkerDef } from '../config.js';
import type { NodeStatus, ProbeContext } from './backend-status-registry.js';

/**
 * Per-probe HTTP timeout. Default was 3s historically but llama-server
 * stalls its HTTP loop while a long-running generation is in flight on
 * its event loop, so `/slots` regularly comes back >3s on a busy GPU
 * and the registry would flip the node to offline mid-inference. 10s
 * gives the server-side enough headroom to interleave the probe with
 * an active generation without losing real liveness sensitivity (a
 * truly dead node still trips ECONNREFUSED / DNS failure inside the
 * window). 2026-05-21 dogfooding feedback.
 */
export const DEFAULT_PROBE_TIMEOUT_MS = 10_000;

/** Trim trailing slashes and one optional `/v1` suffix so we can build sibling paths. */
export function normalizeWorkerBase(endpoint: string): string {
  const trimmed = endpoint.replace(/\/+$/, '');
  return trimmed.endsWith('/v1') ? trimmed.slice(0, -3) : trimmed;
}

function authHeaders(apiKey: string | undefined): Record<string, string> {
  const h: Record<string, string> = { Accept: 'application/json' };
  if (apiKey) h['Authorization'] = `Bearer ${apiKey}`;
  return h;
}

async function fetchWithTimeout(
  url: string,
  init: RequestInit,
  timeoutMs: number,
  fetchImpl: typeof fetch,
  externalSignal?: AbortSignal,
): Promise<Response> {
  const ctrl = new AbortController();
  const t = setTimeout(() => ctrl.abort(), timeoutMs);
  // Chain the registry-wide cycle signal so shutdown can cancel pending
  // probes immediately instead of waiting for timeoutMs.
  let onExternalAbort: (() => void) | undefined;
  if (externalSignal) {
    if (externalSignal.aborted) {
      ctrl.abort();
    } else {
      onExternalAbort = () => ctrl.abort();
      externalSignal.addEventListener('abort', onExternalAbort, { once: true });
    }
  }
  try {
    return await fetchImpl(url, { ...init, signal: ctrl.signal });
  } finally {
    clearTimeout(t);
    if (externalSignal && onExternalAbort) {
      externalSignal.removeEventListener('abort', onExternalAbort);
    }
  }
}

export interface ProbeDeps {
  fetchImpl?: typeof fetch;
  timeoutMs?: number;
  now?: () => string;
}

// ── llama-server (direct worker) ───────────────────────────────────────────

/**
 * Parse a Prometheus text-format /metrics payload from llama-server and
 * extract a current decode throughput in tokens/sec.
 *
 * llama-server exposes several Prometheus metrics relevant to throughput:
 *  - `llamacpp:n_decode_total` (counter)
 *  - `llamacpp:n_prompt_tokens_total` (counter)
 *  - `llamacpp:requests_processing` (gauge)
 *  - `llamacpp:tokens_per_second` (gauge) — newer builds
 *  - `llamacpp:prompt_tokens_seconds` (gauge) — older builds
 *
 * We prefer a directly-named tokens/sec gauge when present; otherwise we
 * fall back to `prompt_tokens_seconds` (also a gauge in /metrics).
 * Counter-based rate computation requires two scrapes and a per-scrape
 * delta — that's a Phase C+ enhancement; Phase C v1 only surfaces gauges.
 *
 * Exported for unit testing. Returns null when no recognised gauge is
 * present in the body.
 */
export function parseLlamaMetricsThroughput(text: string): number | null {
  if (typeof text !== 'string' || text.length === 0) return null;
  // Match `metric_name <number>` ignoring optional `{labels}` and
  // trailing timestamp. Anchored per-line, comment lines (`# HELP`/
  // `# TYPE`) skipped.
  //
  // Order across candidates matters: the first recognised metric family
  // wins. tokens_per_second is the cleanest signal; prompt_tokens_seconds
  // is a fallback. Within a family we use the `g` flag so multi-label
  // exports (one llama-server process serving multiple slots/models with
  // distinct `{model="..."}` labels) are all collected — first-match
  // would silently drop sibling series and misreport the cluster's
  // throughput. We surface the max across labels so the widget shows
  // the fastest currently-decoding model instead of an arbitrary one.
  const candidates = [
    /^llamacpp:tokens_per_second(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
    /^llamacpp:prompt_tokens_seconds(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
    /^llamacpp:n_decode_tokens_per_second(?:\{[^}]*\})?\s+([0-9eE.+-]+)/gm,
  ];
  for (const re of candidates) {
    const values: number[] = [];
    for (const m of text.matchAll(re)) {
      const v = Number(m[1]);
      if (Number.isFinite(v) && v >= 0) values.push(v);
    }
    if (values.length > 0) return Math.max(...values);
  }
  return null;
}

interface LlamaSlot {
  id?: number;
  // `is_processing` is the canonical field on recent llama-server builds.
  // Older builds called it `state` (0 = idle, 1 = processing).
  is_processing?: boolean;
  state?: number;
  model?: string;
  task_id?: number | string;
}

/**
 * Parse a llama-server `/slots` JSON payload into busy / total / model.
 *
 * Tolerates both the modern object-array shape and the older shape where
 * the response was already an envelope (`{ slots: [...] }`). Returns
 * conservative zeros if the payload doesn't look like a slot array.
 */
export function parseLlamaSlots(payload: unknown): {
  busySlots: number;
  totalSlots: number;
  loadedModel: string | null;
} {
  let slots: LlamaSlot[] = [];
  if (Array.isArray(payload)) {
    slots = payload as LlamaSlot[];
  } else if (payload && typeof payload === 'object') {
    const inner = (payload as { slots?: unknown }).slots;
    if (Array.isArray(inner)) slots = inner as LlamaSlot[];
  }
  let busy = 0;
  let model: string | null = null;
  for (const s of slots) {
    const isBusy = s.is_processing === true || (typeof s.state === 'number' && s.state !== 0);
    if (isBusy) busy++;
    if (!model && typeof s.model === 'string' && s.model.length > 0) model = s.model;
  }
  return { busySlots: busy, totalSlots: slots.length, loadedModel: model };
}

/**
 * Probe a direct (llama-server) worker.
 *
 * Tries `/slots` first; if the server rejects it (--no-slots or older
 * builds), falls back to `/health` so we can still report online/offline.
 */
export function buildDirectProbe(deps: ProbeDeps = {}): (worker: WorkerDef, ctx?: ProbeContext) => Promise<NodeStatus> {
  const fetchImpl = deps.fetchImpl ?? fetch;
  const timeoutMs = deps.timeoutMs ?? DEFAULT_PROBE_TIMEOUT_MS;
  const now = deps.now ?? (() => new Date().toISOString());

  return async function probeDirect(worker: WorkerDef, ctx?: ProbeContext): Promise<NodeStatus> {
    const base = normalizeWorkerBase(worker.endpoint);
    // Security: direct probes (llama-server compatible) must NOT receive the
    // worker.apiKey. That key is a LiteLLM virtual key intended for the
    // proxy auth surface — forwarding it to a direct llama-server would
    // leak the secret to upstream access logs in plaintext. Only the proxy
    // probe path (buildProxyProbe) sends Authorization.
    const headers = authHeaders(undefined);
    const slotsUrl = `${base}/slots`;

    let slotsRes: Response;
    try {
      slotsRes = await fetchWithTimeout(slotsUrl, { method: 'GET', headers }, timeoutMs, fetchImpl, ctx?.signal);
    } catch (err) {
      // Network-level failure (DNS, refused, timeout). The server may
      // still be alive but we can't talk to it; record online=false.
      return errorStatus(worker, 'direct', err, now());
    }

    if (slotsRes.ok) {
      let body: unknown;
      try {
        body = await slotsRes.json();
      } catch (err) {
        return errorStatus(worker, 'direct', err, now());
      }
      const parsed = parseLlamaSlots(body);
      // /metrics is a best-effort enrichment: a failure here must NOT
      // demote the status to offline. Some llama-server builds disable
      // /metrics by default (`--metrics` opt-in) or run behind a proxy
      // that strips Prometheus endpoints.
      const throughputTps = await tryFetchThroughput(
        base, headers, timeoutMs, fetchImpl, ctx?.signal,
      );
      return {
        nodeId: worker.id,
        workerId: worker.id,
        source: 'direct',
        online: true,
        busy: parsed.busySlots > 0,
        busySlots: parsed.busySlots,
        totalSlots: parsed.totalSlots,
        loadedModel: parsed.loadedModel ?? worker.model ?? null,
        throughputTps,
        lastSeen: now(),
      };
    }

    // /slots disabled (--no-slots) returns 404 on most builds and 501 on
    // some forks. /health is the documented liveness endpoint and is
    // always available.
    if (slotsRes.status === 404 || slotsRes.status === 405 || slotsRes.status === 501) {
      return probeHealthFallback(base, headers, worker, timeoutMs, fetchImpl, now, ctx?.signal);
    }

    return errorStatus(
      worker,
      'direct',
      new Error(`/slots returned HTTP ${slotsRes.status}`),
      now(),
    );
  };
}

/**
 * Best-effort throughput fetch. Returns null on any failure (404 from
 * builds without --metrics, network error, parse failure) so the
 * primary /slots-derived status is unaffected.
 *
 * Uses the same per-request timeout cap as /slots so a hung /metrics
 * endpoint can't double the worst-case probe latency.
 */
async function tryFetchThroughput(
  base: string,
  headers: Record<string, string>,
  timeoutMs: number,
  fetchImpl: typeof fetch,
  externalSignal?: AbortSignal,
): Promise<number | null> {
  try {
    const res = await fetchWithTimeout(
      `${base}/metrics`,
      { method: 'GET', headers: { ...headers, Accept: 'text/plain' } },
      timeoutMs,
      fetchImpl,
      externalSignal,
    );
    if (!res.ok) return null;
    const text = await res.text();
    return parseLlamaMetricsThroughput(text);
  } catch {
    return null;
  }
}

async function probeHealthFallback(
  base: string,
  headers: Record<string, string>,
  worker: WorkerDef,
  timeoutMs: number,
  fetchImpl: typeof fetch,
  now: () => string,
  externalSignal?: AbortSignal,
): Promise<NodeStatus> {
  try {
    const res = await fetchWithTimeout(`${base}/health`, { method: 'GET', headers }, timeoutMs, fetchImpl, externalSignal);
    if (!res.ok) {
      return errorStatus(worker, 'direct', new Error(`/health returned HTTP ${res.status}`), now());
    }
    return {
      nodeId: worker.id,
      workerId: worker.id,
      source: 'direct',
      online: true,
      busy: false,
      busySlots: 0,
      totalSlots: 0,
      loadedModel: worker.model ?? null,
      throughputTps: null,
      lastSeen: now(),
    };
  } catch (err) {
    return errorStatus(worker, 'direct', err, now());
  }
}

// ── LiteLLM (proxy worker) ─────────────────────────────────────────────────

/**
 * Single LiteLLM `/health` entry. The exact field names depend on the
 * LiteLLM version; we treat all fields as optional and only consume the
 * ones we recognise. `model` is the deployment alias; some versions
 * nest it under `litellm_params.model`.
 */
interface LiteLLMHealthEntry {
  model?: unknown;
  litellm_params?: { model?: unknown } | null;
  cache?: unknown;
  error?: unknown;
  /**
   * AAO Gateway extension fields (additive — vanilla LiteLLM omits them
   * and we treat absence as "unknown busy state", same as before this
   * extension existed). When present, parseLiteLLMHealth populates the
   * NodeStatus with the gateway-aggregated busy view so every AAO
   * client pointed at the same gateway sees consistent backend usage.
   */
  aao_busy_slots?: unknown;
  aao_total_slots?: unknown;
  aao_saturated?: unknown;
  aao_last_seen?: unknown;
}

interface LiteLLMHealthBody {
  healthy_endpoints?: LiteLLMHealthEntry[];
  unhealthy_endpoints?: LiteLLMHealthEntry[];
  // Some LiteLLM versions return `healthy_count` etc.; we don't need them
  // since we walk the arrays directly.
  [key: string]: unknown;
}

/**
 * Coerce an arbitrary `aao_*_slots` value into a non-negative integer.
 * Floors fractional inputs, drops anything non-finite or negative to
 * zero so a malformed entry never corrupts the registry.
 */
function coerceNonNegInt(v: unknown): number {
  if (typeof v !== 'number' || !Number.isFinite(v) || v < 0) return 0;
  return Math.floor(v);
}

function extractDeploymentId(entry: LiteLLMHealthEntry): string | null {
  if (typeof entry.model === 'string' && entry.model.length > 0) return entry.model.trim();
  const nested = entry.litellm_params?.model;
  if (typeof nested === 'string' && nested.length > 0) return nested.trim();
  return null;
}

/**
 * Coerce a LiteLLM `error` field to a human-readable string.
 *
 * LiteLLM pre-1.40 surfaced errors as plain strings, but recent
 * versions (post-1.40) wrap them in `{ message: string, type: string }`
 * objects. The previous string-only guard silently dropped the object
 * form, leaving the widget with "unhealthy but reason unknown" — a
 * red icon with nothing to action.
 *
 * Strategy:
 *  - string → return as-is
 *  - object with `.message` → return that
 *  - other object → JSON.stringify so something useful surfaces in
 *    logs / UI (better than the cryptic "[object Object]")
 *  - null/undefined → undefined (caller skips the field)
 */
function extractErrorMessage(err: unknown): string | undefined {
  if (typeof err === 'string') return err;
  if (err && typeof err === 'object') {
    const obj = err as { message?: unknown };
    if (typeof obj.message === 'string') return obj.message;
    try {
      return JSON.stringify(err);
    } catch {
      return String(err);
    }
  }
  return undefined;
}

/**
 * Convert a LiteLLM /health response into one NodeStatus per deployment.
 * Exported for unit-testing without a live proxy.
 *
 * Dedup policy (Phase C): **unhealthy wins precedence**. When the same
 * deployment id appears in both `healthy_endpoints` and
 * `unhealthy_endpoints` (a flapping backend), we surface the unhealthy
 * row so the widget doesn't mask a degraded state behind a green icon.
 * The error message from the unhealthy entry is preserved in
 * `lastProbeError`. Phase B used "first appearance wins" which silently
 * hid flap conditions from operators.
 */
export function parseLiteLLMHealth(
  body: unknown,
  workerId: string,
  ts: string,
): NodeStatus[] {
  const obj = (body && typeof body === 'object' ? body : {}) as LiteLLMHealthBody;
  const out: NodeStatus[] = [];
  const seen = new Set<string>();
  const push = (entry: LiteLLMHealthEntry, online: boolean): void => {
    const id = extractDeploymentId(entry);
    if (!id || seen.has(id)) return;
    seen.add(id);
    // AAO Gateway extension: when the gateway annotates each entry
    // with aao_busy_slots / aao_total_slots, we inherit its
    // multi-client-aggregated view. Without these fields (vanilla
    // LiteLLM or older gateway) busy stays 0 — same as before the
    // extension existed.
    const busySlots = coerceNonNegInt(entry.aao_busy_slots);
    const totalSlots = coerceNonNegInt(entry.aao_total_slots);
    out.push({
      nodeId: id,
      workerId,
      source: 'proxy',
      online,
      busy: busySlots > 0,
      busySlots,
      totalSlots,
      loadedModel: id,
      throughputTps: null,
      lastSeen: ts,
      lastProbeError: online ? undefined : extractErrorMessage(entry.error),
    });
  };
  // Iterate unhealthy first so a flapping deployment registers as
  // unhealthy and the healthy entry's subsequent push is skipped by the
  // `seen` guard.
  for (const entry of obj.unhealthy_endpoints ?? []) push(entry, false);
  for (const entry of obj.healthy_endpoints ?? []) push(entry, true);
  return out;
}

/**
 * Probe a LiteLLM proxy worker.
 *
 * Returns one NodeStatus per deployment the proxy reports. If the proxy
 * itself is unreachable we synthesize a single offline NodeStatus keyed
 * to the workerId so the widget can render a "proxy down" row.
 */
export function buildProxyProbe(deps: ProbeDeps = {}): (worker: WorkerDef, ctx?: ProbeContext) => Promise<NodeStatus[]> {
  const fetchImpl = deps.fetchImpl ?? fetch;
  const timeoutMs = deps.timeoutMs ?? DEFAULT_PROBE_TIMEOUT_MS;
  const now = deps.now ?? (() => new Date().toISOString());

  return async function probeProxy(worker: WorkerDef, ctx?: ProbeContext): Promise<NodeStatus[]> {
    const base = normalizeWorkerBase(worker.endpoint);
    const headers = authHeaders(worker.apiKey);
    const url = `${base}/health`;
    let res: Response;
    try {
      res = await fetchWithTimeout(url, { method: 'GET', headers }, timeoutMs, fetchImpl, ctx?.signal);
    } catch (err) {
      return [errorStatus(worker, 'proxy', err, now())];
    }
    if (!res.ok) {
      return [errorStatus(worker, 'proxy', new Error(`/health returned HTTP ${res.status}`), now())];
    }
    let body: unknown;
    try {
      body = await res.json();
    } catch (err) {
      return [errorStatus(worker, 'proxy', err, now())];
    }
    const ts = now();
    const parsed = parseLiteLLMHealth(body, worker.id, ts);
    if (parsed.length === 0) {
      // Proxy is alive but reports zero deployments: surface a single
      // synthetic row so the widget shows the proxy itself rather than
      // a blank panel.
      return [{
        nodeId: worker.id,
        workerId: worker.id,
        source: 'proxy',
        online: true,
        busy: false,
        busySlots: 0,
        totalSlots: 0,
        loadedModel: null,
        throughputTps: null,
        lastSeen: ts,
      }];
    }
    return parsed;
  };
}

function errorStatus(
  worker: WorkerDef,
  source: 'direct' | 'proxy',
  err: unknown,
  ts: string,
): NodeStatus {
  return {
    nodeId: worker.id,
    workerId: worker.id,
    source,
    online: false,
    busy: false,
    busySlots: 0,
    totalSlots: 0,
    loadedModel: null,
    throughputTps: null,
    lastSeen: ts,
    lastProbeError: err instanceof Error ? err.message : String(err),
  };
}
