/**
 * `GET /health` — LiteLLM-shaped health snapshot for the configured backends.
 *
 * Auth not required: this is the same surface AAO's NodeStatusWidget /
 * BackendStatusRegistry probes, so wiring a token would force every
 * monitor to be configured with a virtual key. The information leaked
 * here (backend ids + offline reasons) is the same shape LiteLLM Proxy
 * exposes, so we match its policy for drop-in compatibility.
 *
 * Response shape (verbatim from LiteLLM /health):
 *
 *   {
 *     "healthy_endpoints":   [{"model": "<id>", "litellm_params": {"model": "<id>"}}, …],
 *     "unhealthy_endpoints": [{"model": "<id>", "litellm_params": {"model": "<id>"}, "error": "..."}, …],
 *     "healthy_count":   N,
 *     "unhealthy_count": M
 *   }
 *
 * AAO's `parseLiteLLMHealth` already understands this exact shape — no
 * special-casing required on the consumer side.
 */
import type { Request, Response } from 'express';
import type { Router } from './router.js';
import type { GatewayBackendConfig } from './config.js';
import type { NodeStatus } from '../engine/backend-status-registry.js';
import type { BackendInflightCounter } from './stream-proxy.js';

export interface HealthEndpointDeps {
  router: Router;
  /**
   * Real-time in-flight counter. When supplied, `/health` reports
   * `aao_busy_slots = max(probed busySlots, inflight counter)` so a
   * client AAO sees sub-probe-cadence usage (chat-completions running
   * RIGHT NOW, not "what /slots said 5–30 seconds ago"). Optional so
   * tests can leave it off and get the probe-only view.
   */
  inflight?: BackendInflightCounter;
}

export interface LiteLLMHealthEntry {
  model: string;
  litellm_params: { model: string };
  error?: string;
  /**
   * AAO Gateway extension (LiteLLM-compatible additive fields).
   *
   * Gateway probes every backend's llama-server /slots so its busy
   * counter is the union of in-flight requests across *all* clients
   * pointing at this gateway. Surfacing it on /health lets each
   * client AAO's BackendStatusRegistry see "backend X is busy because
   * some other AAO is using it right now" without needing a new
   * endpoint or per-client visibility into the gateway's metrics.
   *
   * The fields are namespaced under `aao_*` to make it obvious in the
   * payload which keys are LiteLLM-shape and which are extensions.
   * Vanilla LiteLLM consumers simply ignore unknown keys.
   */
  aao_busy_slots?: number;
  aao_total_slots?: number;
  /** True when busy_slots >= total_slots — saturated, still online. */
  aao_saturated?: boolean;
  /** ISO 8601 timestamp of the last successful probe. */
  aao_last_seen?: string;
}

export interface LiteLLMHealthBody {
  healthy_endpoints: LiteLLMHealthEntry[];
  unhealthy_endpoints: LiteLLMHealthEntry[];
  healthy_count: number;
  unhealthy_count: number;
}

/**
 * Pure converter — given a router snapshot, produce the LiteLLM-shaped body.
 *
 * Classification rule: a backend is "healthy" iff the latest probe
 * reports `online === true` AND it isn't fully saturated (busy < total).
 * Saturated backends still get bucketed as healthy (they're up, just
 * busy) — only `online === false` lands in unhealthy. This matches
 * LiteLLM's semantics: /health reports liveness, not capacity.
 *
 * Backends without a probe row yet (registry not warm at startup) are
 * OMITTED from both arrays. LiteLLM Proxy uses the same policy — surfacing
 * them as "unhealthy" would paint NodeStatusWidget entirely red for the
 * first ~5s after boot, which dashboards consume as "site down". They'll
 * appear in the next /health response once the first probe cycle finishes.
 */
export function buildHealthBody(
  snapshot: Array<{ config: GatewayBackendConfig; status: NodeStatus | null }>,
  inflight?: BackendInflightCounter,
): LiteLLMHealthBody {
  const healthy: LiteLLMHealthEntry[] = [];
  const unhealthy: LiteLLMHealthEntry[] = [];
  for (const s of snapshot) {
    const id = s.config.id;
    const base: LiteLLMHealthEntry = {
      model: id,
      litellm_params: { model: id },
    };
    if (!s.status) {
      // Unprobed — skip entirely so the widget shows the post-boot state,
      // not an artificial all-red spike. LiteLLM-compat.
      continue;
    }
    // Attach AAO-extension fields whenever the registry has a useful
    // figure. Gateway clients (parseLiteLLMHealth) read these to feed
    // their local BackendStatusRegistry, so the busy view stays
    // consistent across every AAO pointed at this gateway.
    //
    // The live in-flight counter (`inflight`) wins over the probed
    // value when it's higher: probes run on a 5–30s cadence and miss
    // sub-cadence chats entirely. Taking max() means a chat that
    // starts in tick N+1 is visible immediately (via inflight) and
    // remains visible after it ends until the next probe lands
    // (via the probed busySlots).
    const probedBusy = s.status.busySlots;
    const liveBusy = inflight?.get(id) ?? 0;
    const total = s.status.totalSlots > 0 ? s.status.totalSlots : s.config.maxSlots;
    const reportedBusy = Math.max(probedBusy, liveBusy);
    base.aao_busy_slots = reportedBusy;
    base.aao_total_slots = total;
    base.aao_saturated = total > 0 && reportedBusy >= total;
    if (s.status.lastSeen) base.aao_last_seen = s.status.lastSeen;
    if (s.status.online === true) {
      healthy.push(base);
    } else {
      const err = s.status.lastProbeError ?? 'offline';
      unhealthy.push({ ...base, error: err });
    }
  }
  return {
    healthy_endpoints: healthy,
    unhealthy_endpoints: unhealthy,
    healthy_count: healthy.length,
    unhealthy_count: unhealthy.length,
  };
}

export function buildHealthHandler(deps: HealthEndpointDeps) {
  return function healthHandler(_req: Request, res: Response): void {
    const snapshot = deps.router.snapshot();
    res.json(buildHealthBody(snapshot, deps.inflight));
  };
}

/**
 * Liveness probe — bare 200 OK. No auth, no upstream check, nothing
 * that can block (so k8s liveness can fire it every second without
 * load). For backend-aware health use `/health` above.
 */
export function buildLivenessHandler() {
  return function livenessHandler(_req: Request, res: Response): void {
    res.status(200).json({ status: 'ok' });
  };
}
