/**
 * Stream proxy for `POST /v1/chat/completions`.
 *
 * One handler that covers both SSE streaming (the default for chat
 * completions) and the non-streaming fallback (`stream: false`).
 *
 * Header behaviour
 * ────────────────
 * Once the upstream response is in hand we:
 *  - inject `x-aao-backend-id` and `x-litellm-model-id` (same value)
 *    so AAO's existing openai-compat client (which reads the LiteLLM
 *    header) keeps working unchanged
 *  - pass through the upstream's Content-Type / Transfer-Encoding /
 *    Cache-Control so the wire shape stays identical
 *  - mirror the upstream status code so 4xx/5xx surface to the caller
 *
 * Streaming
 * ─────────
 * We iterate the upstream body and write each chunk to the client.
 * Two timeouts apply:
 *  - `upstreamTimeoutSec` caps the time-to-first-byte (TTFB) — a hung
 *    upstream gets aborted before the client even sees a response.
 *  - A 60s idle timer kills the stream if the upstream stops emitting
 *    chunks. Catches the "first token fast, then dead" failure mode
 *    that bare request_timeout would miss.
 *
 * Client disconnect (`res.on('close')` with `writableEnded === false`)
 * propagates to the upstream via AbortController.abort() — no leaked
 * sockets after a browser refresh. We listen on the RESPONSE close, not
 * the request close, because Node 18+ fires `req.on('close')` on body
 * drain even while the keep-alive socket stays open.
 *
 * Why for-await over response.body?.pipe(res)
 * ──────────────────────────────────────────
 * Two reasons:
 *  1. We need to inject headers BEFORE the first byte is written, which
 *     is fiddly with pipe().
 *  2. The idle timer needs visibility into chunk arrival — pipe() hides
 *     that behind the stream.
 */
import type { Request, Response } from 'express';
import { logger } from '../logger.js';
import type { Router } from './router.js';
import type { AuthenticatedRequest } from './auth.js';
import { currentPeriodYearMonth } from './period.js';
import type { GatewayMetrics, GatewayRequestStatus } from '../metrics/gateway-metrics.js';

export const IDLE_CHUNK_TIMEOUT_MS = 60_000;

/**
 * Optional sink for usage / request counters. Wired by gateway/bootstrap.ts
 * when a Repository is open. Calls are best-effort: a throw is swallowed
 * (and logged) so a failed bookkeeping write never affects the client
 * response that already streamed cleanly.
 */
export interface UsageRecorderDeps {
  recordUsage(params: {
    keyId: string;
    period: string;
    tokensIn: number;
    tokensOut: number;
    requests: number;
  }): void;
}

/**
 * Phase 3a F8 — registry of in-flight streams so graceful shutdown can
 * signal them with a distinguishable reason (vs a normal abort) and let
 * each handler emit a `gateway_shutdown` SSE event before closing. The
 * registry is also used by tests to drive synthetic shutdowns without
 * a real signal.
 */
export interface StreamAbortHandle {
  /** Abort the in-flight stream. The flag stays set so handlers can
   * tell shutdown from a regular abort. */
  abort(): void;
}

export interface StreamRegistry {
  /** Register an in-flight stream. The returned unregister fn MUST be
   * called from the request's finally block to prevent leaks. */
  register(handle: StreamAbortHandle): () => void;
  /** Mark the gateway as shutting down + signal every registered stream. */
  signalShutdown(): void;
  /** True when signalShutdown has been called at least once. */
  isShuttingDown(): boolean;
  /** Test hook: drop all registrations + reset the shutdown flag. */
  reset(): void;
  /** Test hook: number of currently-registered streams. */
  size(): number;
}

/**
 * Construct a new in-memory stream registry. Bootstrap wires one per
 * gateway instance and hands it to the chat-completions handler so the
 * graceful-shutdown hook can drain SSE clients with a structured event.
 */
export function createStreamRegistry(): StreamRegistry {
  const handles = new Set<StreamAbortHandle>();
  let shuttingDown = false;
  return {
    register(h) {
      handles.add(h);
      return (): void => {
        handles.delete(h);
      };
    },
    signalShutdown() {
      shuttingDown = true;
      // Snapshot before iterating: abort() may unregister via finally.
      for (const h of [...handles]) {
        try {
          h.abort();
        } catch (e) {
          logger.warn(`[gateway-stream] shutdown abort threw: ${e instanceof Error ? e.message : String(e)}`);
        }
      }
    },
    isShuttingDown() {
      return shuttingDown;
    },
    reset() {
      handles.clear();
      shuttingDown = false;
    },
    size() {
      return handles.size;
    },
  };
}

export interface StreamProxyDeps {
  router: Router;
  /** Upstream fetch timeout (whole-response cap for non-streaming, TTFB cap for streaming). */
  upstreamTimeoutSec: number;
  /** Total request budget incl. streaming body. AbortController fires at this point. */
  requestTimeoutSec: number;
  /** Injected so tests can mock without monkey-patching globalThis. */
  fetchImpl?: typeof fetch;
  /**
   * Phase 2b: optional usage recorder. When present, the proxy extracts
   * `usage.{prompt_tokens, completion_tokens}` from the upstream response
   * and forwards it as a single `recordUsage` call in finally. Omit to
   * keep Phase 1 behavior (no DB writes from the proxy).
   */
  usageRecorder?: UsageRecorderDeps;
  /**
   * Phase 3a F8: optional registry so graceful shutdown can abort
   * in-flight SSE streams with a distinguishable signal. When wired,
   * the handler registers on stream start, unregisters on finally, and
   * checks `registry.isShuttingDown()` inside the abort branch so it
   * emits a `gateway_shutdown` event instead of the generic
   * `gateway_timeout` one.
   */
  streamRegistry?: StreamRegistry;
  /**
   * Phase 3b: optional Prometheus metrics handle. When set, response
   * completion emits requests_total / request_duration_seconds with
   * the observed status enum, plus tokens_total (in/out) when usage
   * was parsed from the upstream response. Failures in metric emission
   * are swallowed — metrics MUST NOT alter the client response.
   */
  metrics?: GatewayMetrics;
  /**
   * Real-time per-backend in-flight counter. Incremented immediately
   * after the router picks a backend (before fetch) and decremented
   * in finally. Lets `/health` surface live busy state without waiting
   * for the next BackendStatusRegistry probe (which runs at 5–30s
   * cadence and often misses sub-5s chats entirely). Optional so unit
   * tests can leave it off without wiring a stub.
   */
  inflight?: BackendInflightCounter;
}

/**
 * Per-backend in-flight counter, owned by the gateway. Stream-proxy
 * increments before the upstream fetch and decrements in finally;
 * health-endpoint reads `get(backendId)` and reports the live max of
 * (probed busy slots, in-flight count). The counter naturally reflects
 * the union of every client AAO talking to this gateway, so the busy
 * view propagates correctly across multi-client deployments.
 */
export interface BackendInflightCounter {
  inc(backendId: string): void;
  dec(backendId: string): void;
  get(backendId: string): number;
  /** Snapshot all currently-tracked backends + their counts. */
  snapshot(): Map<string, number>;
}

export function createBackendInflightCounter(): BackendInflightCounter {
  const counts = new Map<string, number>();
  return {
    inc(id) {
      counts.set(id, (counts.get(id) ?? 0) + 1);
    },
    dec(id) {
      const next = (counts.get(id) ?? 0) - 1;
      if (next <= 0) counts.delete(id);
      else counts.set(id, next);
    },
    get(id) {
      return counts.get(id) ?? 0;
    },
    snapshot() {
      return new Map(counts);
    },
  };
}

/**
 * Extract `{prompt_tokens, completion_tokens}` from a non-streaming
 * chat-completions JSON body. Returns null when the field is missing or
 * malformed — never throws, because a parsing failure must not affect
 * the client response.
 */
export function extractUsageFromJson(text: string): { promptTokens: number; completionTokens: number } | null {
  try {
    const parsed: unknown = JSON.parse(text);
    if (!isRecord(parsed)) return null;
    const usage = parsed['usage'];
    if (!isRecord(usage)) return null;
    return coerceUsageRecord(usage);
  } catch {
    return null;
  }
}

/**
 * Scan an SSE chunk for the OpenAI usage record. The OpenAI streaming
 * convention puts the cumulative usage on the *last* `data:` event
 * before `data: [DONE]`, but some backends (vLLM, latest Ollama) emit
 * incremental usage on earlier chunks as well. We scan every `data:`
 * payload and return the latest non-null observation, so the caller
 * sees the final value when the stream ends.
 *
 * Pure: does not mutate input. Returns null when no usage was found in
 * the chunk.
 */
export function extractUsageFromSseChunk(chunk: string): { promptTokens: number; completionTokens: number } | null {
  // SSE frames are separated by blank lines. Within a frame the `data:`
  // payload may span multiple lines per spec, though chat-completions
  // always emits a single-line payload. We split on \n and parse any
  // line starting with `data: ` whose body is JSON.
  let latest: { promptTokens: number; completionTokens: number } | null = null;
  for (const rawLine of chunk.split(/\r?\n/)) {
    if (!rawLine.startsWith('data:')) continue;
    const payload = rawLine.slice(5).trimStart();
    if (payload === '[DONE]' || payload.length === 0) continue;
    try {
      const parsed: unknown = JSON.parse(payload);
      if (!isRecord(parsed)) continue;
      const usage = parsed['usage'];
      if (!isRecord(usage)) continue;
      const coerced = coerceUsageRecord(usage);
      if (coerced) latest = coerced;
    } catch {
      // Partial chunks happen at TCP boundaries — silently skip and
      // wait for the next chunk to bring a complete frame.
      continue;
    }
  }
  return latest;
}

function isRecord(v: unknown): v is Record<string, unknown> {
  return typeof v === 'object' && v !== null && !Array.isArray(v);
}

function coerceUsageRecord(usage: Record<string, unknown>): { promptTokens: number; completionTokens: number } | null {
  const p = usage['prompt_tokens'];
  const c = usage['completion_tokens'];
  const pt = typeof p === 'number' && Number.isFinite(p) && p >= 0 ? Math.floor(p) : null;
  const ct = typeof c === 'number' && Number.isFinite(c) && c >= 0 ? Math.floor(c) : null;
  if (pt === null && ct === null) return null;
  return { promptTokens: pt ?? 0, completionTokens: ct ?? 0 };
}

/**
 * Build the POST /v1/chat/completions handler.
 */
export function buildChatCompletionsHandler(deps: StreamProxyDeps) {
  const fetchImpl = deps.fetchImpl ?? fetch;

  return async function chatCompletionsHandler(req: Request, res: Response): Promise<void> {
    const body = (req.body ?? {}) as { model?: unknown; stream?: unknown };
    const model = typeof body.model === 'string' ? body.model : '';
    const auth0 = (req as AuthenticatedRequest).gatewayAuth;
    const team0 = auth0?.team ?? 'unknown';
    if (!model) {
      // Treat malformed body as upstream_error (input side). No backend
      // resolved yet, so use sentinel labels.
      emitRequestMetric(deps.metrics, {
        team: team0, backend: 'none', model: 'none', status: 'upstream_error', durationMs: 0,
      });
      res.status(400).json({ error: 'missing model field' });
      return;
    }

    const backend = deps.router.pick(model);
    if (!backend) {
      emitRequestMetric(deps.metrics, {
        team: team0, backend: 'none', model, status: 'no_backend', durationMs: 0,
      });
      // Fast-fail rather than queue. The client (AAO worker via
      // openai-compat.ts) has its own retry/backoff and a 503 is the
      // honest answer when no slot is free.
      res.status(503).json({ error: 'no available backend' });
      return;
    }

    const isStreaming = body.stream !== false; // chat/completions defaults to true

    const auth = auth0;
    const team = team0;
    const t0 = Date.now();
    // Real-time in-flight bookkeeping. The reservation (inc) is performed
    // by router.pick() atomically with backend selection — that's what
    // enforces the per-host max_slots cap and makes the dispatch visible
    // to the very next pick() in the same probe interval (and to /health).
    // Here we own only the matching RELEASE: a single dec() guaranteed in
    // every terminal path (upstream-error / abort / streaming /
    // non-streaming), guarded so it runs exactly once. Note: this pairs
    // with the inc in pick(); a non-null pick() must always reach one of
    // the decInflightOnce() call sites below.
    let inflightDecremented = false;
    const decInflightOnce = (): void => {
      if (inflightDecremented) return;
      inflightDecremented = true;
      try { deps.inflight?.dec(backend.id); } catch { /* never throw from bookkeeping */ }
    };
    // Track usage across the request — populated by the non-streaming
    // body read or by the streaming chunk peek, then flushed via
    // recordUsageIfDb() in finally. Default 0/0 means "no usage seen";
    // we still emit a request count when the auth is DB-backed so the
    // counter exists for the admin GET /usage history view.
    let observedUsage: { promptTokens: number; completionTokens: number } | null = null;
    const recordUsageIfDb = (): void => {
      // Skip when:
      //  - no usage recorder wired (Phase 1 stateless deploy)
      //  - the matched key wasn't DB-backed (config fallback; no row id)
      //  - the auth ran before the deps were configured (defensive)
      if (!deps.usageRecorder || !auth?.keyId || auth.source !== 'db') return;
      try {
        deps.usageRecorder.recordUsage({
          keyId: auth.keyId,
          period: currentPeriodYearMonth(),
          tokensIn: observedUsage?.promptTokens ?? 0,
          tokensOut: observedUsage?.completionTokens ?? 0,
          requests: 1,
        });
      } catch (e) {
        // Swallow so bookkeeping never fails the user-visible response.
        logger.warn(`[gateway] recordUsage failed team=${team} err=${e instanceof Error ? e.message : String(e)}`);
      }
    };

    // One AbortController gates the entire upstream call. Chained signals:
    //  - request-wide timeout (requestTimeoutSec)
    //  - upstream TTFB / non-streaming whole-response timeout
    //  - client disconnect (res.on('close') — see note below)
    //  - idle-chunk timer (streaming only, reset per chunk)
    const ctrl = new AbortController();
    const requestTimeoutMs = Math.max(1, Math.floor(deps.requestTimeoutSec * 1000));
    const requestTimer = setTimeout(() => {
      logger.warn(`[gateway] request timeout team=${team} backend=${backend.id} model=${model}`);
      ctrl.abort();
    }, requestTimeoutMs);
    if (typeof requestTimer.unref === 'function') requestTimer.unref();

    // Detect client disconnect via the RESPONSE stream's close event, not
    // the request's. Node 18+ fires `req.on('close')` as soon as the
    // request body has been fully consumed, even when the underlying
    // HTTP/1.1 keep-alive socket stays open and the client is still
    // waiting for the response. The previous req.on('close') wiring
    // (Phase 1, PR #326) therefore aborted ctrl synchronously on every
    // request, turning the gateway into a perma-502 against any real
    // client (only supertest's mock pipeline kept the tests green).
    //
    // `res.on('close')` only fires once the response stream is closed —
    // either we sent it all (res.writableEnded === true) or the client
    // gave up mid-response (res.writableEnded === false). Only the
    // second case warrants aborting the upstream fetch.
    const onClientClose = (): void => {
      if (!res.writableEnded) {
        ctrl.abort();
      }
    };
    res.on('close', onClientClose);

    const upstreamUrl = `${stripTrailingSlash(backend.endpoint)}/chat/completions`;
    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      'Accept': isStreaming ? 'text/event-stream' : 'application/json',
    };
    if (backend.apiKey) headers['Authorization'] = `Bearer ${backend.apiKey}`;

    // The upstream-fetch timeout caps TTFB only; once headers arrive we
    // switch to the per-chunk idle timer for streaming, or to the body
    // read for non-streaming.
    const ttfbTimer = setTimeout(() => {
      logger.warn(`[gateway] upstream TTFB timeout team=${team} backend=${backend.id}`);
      ctrl.abort();
    }, Math.max(1, Math.floor(deps.upstreamTimeoutSec * 1000)));
    if (typeof ttfbTimer.unref === 'function') ttfbTimer.unref();

    let upstreamRes: globalThis.Response;
    try {
      upstreamRes = await fetchImpl(upstreamUrl, {
        method: 'POST',
        headers,
        body: JSON.stringify(req.body ?? {}),
        signal: ctrl.signal,
      });
    } catch (err) {
      clearTimeout(ttfbTimer);
      clearTimeout(requestTimer);
      res.off('close', onClientClose);
      decInflightOnce();
      // Classify the failure based on the error shape rather than
      // ctrl.signal.aborted alone: supertest reliably races a request
      // 'close' against our finally block, so the abort flag is
      // often true here even when the upstream rejected with a
      // network error. AbortError on `err` is the canonical "we aborted
      // the fetch" signal.
      const isAbortErr = (err as { name?: string } | null)?.name === 'AbortError';
      // Aborted (client gone, request timeout, etc.) → no response to write.
      if (ctrl.signal.aborted && isAbortErr && res.writableEnded) {
        emitRequestMetric(deps.metrics, {
          team, backend: backend.id, model, status: 'client_aborted', durationMs: Date.now() - t0,
        });
        return;
      }
      const msg = err instanceof Error ? err.message : String(err);
      logger.warn(`[gateway] upstream fetch failed team=${team} backend=${backend.id} model=${model} err=${msg}`);
      if (!res.headersSent) {
        res.status(502).json({ error: 'upstream unreachable' });
      } else {
        try { res.end(); } catch { /* ignore */ }
      }
      emitRequestMetric(deps.metrics, {
        team, backend: backend.id, model,
        status: isAbortErr ? 'gateway_timeout' : 'upstream_error',
        durationMs: Date.now() - t0,
      });
      return;
    }
    clearTimeout(ttfbTimer);

    // Mirror upstream status. Preserve the upstream Content-Type so SSE
    // stays SSE.
    //
    // Transfer-Encoding is only forwarded in the streaming branch.
    // For non-streaming we use res.send(text), which sets Content-Length
    // automatically — passing through `transfer-encoding: chunked` on top
    // produces both headers in the response, which violates RFC 9112 §6.2
    // and triggers connection drops at nginx / HAProxy.
    const passthroughHeaders = isStreaming
      ? ['content-type', 'transfer-encoding', 'cache-control']
      : ['content-type', 'cache-control'];
    for (const h of passthroughHeaders) {
      const v = upstreamRes.headers.get(h);
      if (v) res.setHeader(h, v);
    }
    res.setHeader('x-aao-backend-id', backend.id);
    // LiteLLM compatibility header — AAO's existing openai-compat client
    // reads this to map "which physical backend handled the call".
    res.setHeader('x-litellm-model-id', backend.id);
    res.status(upstreamRes.status);

    if (!isStreaming) {
      // Non-streaming: drain to JSON / text. We don't reparse the JSON —
      // pass the bytes through verbatim to avoid corrupting model output.
      let bodyReadOk = false;
      try {
        const text = await upstreamRes.text();
        res.send(text);
        bodyReadOk = true;
        // Side-effect: peek into the JSON to extract usage for the
        // bookkeeping write. Wrap in its own try/catch (separate from
        // the body-read failure path) so a JSON parse error never alters
        // the response we already sent.
        try {
          const usage = extractUsageFromJson(text);
          if (usage) observedUsage = usage;
        } catch { /* extractUsageFromJson is itself try/catch — defensive */ }
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        // If the abort fired before we managed to start sending headers,
        // the client never saw a status — surface a 504 instead of an
        // empty close. Without this the caller sees "socket hang up" with
        // no way to tell it was an upstream timeout.
        if (ctrl.signal.aborted && !res.headersSent) {
          res.status(504).json({ error: 'gateway timeout (upstream took too long)' });
        } else {
          logger.warn(`[gateway] non-streaming body read failed team=${team} backend=${backend.id} err=${msg}`);
          if (!res.headersSent) res.status(502);
          try { res.end(); } catch { /* ignore */ }
        }
      } finally {
        clearTimeout(requestTimer);
        res.off('close', onClientClose);
        decInflightOnce();
        recordUsageIfDb();
        const ms = Date.now() - t0;
        const usageStr = observedUsage
          ? ` tokens_in=${observedUsage.promptTokens} tokens_out=${observedUsage.completionTokens}`
          : '';
        logger.info(`[gateway] team=${team} backend=${backend.id} model=${model} streaming=false status=${upstreamRes.status} duration_ms=${ms}${usageStr}`);
        emitRequestMetric(deps.metrics, {
          team, backend: backend.id, model,
          status: classifyResponseStatus(upstreamRes.status, ctrl, deps.streamRegistry, bodyReadOk),
          durationMs: ms,
          tokensIn: observedUsage?.promptTokens ?? 0,
          tokensOut: observedUsage?.completionTokens ?? 0,
        });
      }
      return;
    }

    // Streaming: pump chunks while watching for idle.
    const body2 = upstreamRes.body;
    if (!body2) {
      // No body to pump — close cleanly.
      clearTimeout(requestTimer);
      res.off('close', onClientClose);
      decInflightOnce();
      res.end();
      return;
    }

    // F8: register this stream with the shutdown registry so a graceful
    // shutdown can flip the shutting-down flag and abort the controller
    // from outside. Stream unregisters via `unregisterStream()` in the
    // finally block below. No-op when no registry was wired.
    const unregisterStream = deps.streamRegistry
      ? deps.streamRegistry.register({ abort: () => ctrl.abort() })
      : null;

    let lastChunkAt = Date.now();
    const idleTimer = setInterval(() => {
      if (Date.now() - lastChunkAt > IDLE_CHUNK_TIMEOUT_MS) {
        logger.warn(`[gateway] upstream idle timeout team=${team} backend=${backend.id} since_ms=${Date.now() - lastChunkAt}`);
        ctrl.abort();
        clearInterval(idleTimer);
      }
    }, Math.min(IDLE_CHUNK_TIMEOUT_MS, 5000));
    if (typeof idleTimer.unref === 'function') idleTimer.unref();

    // SSE usage extraction state. We accumulate a small text buffer
    // across chunks because OpenAI's `data:` frames can be split at
    // arbitrary byte boundaries by the TCP layer — extracting usage
    // requires whole JSON payloads. The buffer is trimmed to the last
    // partial frame after each scan to bound memory: even at 100k
    // chunks/sec the buffer stays under a few KB because we drop
    // everything before the last `\n\n` frame separator.
    let sseBuf = '';
    const SSE_BUF_MAX = 64 * 1024; // 64 KiB hard cap as defense
    const textDecoder = new TextDecoder('utf-8', { fatal: false });
    // F8 perf: track how far into sseBuf we've already scanned for usage
    // so each new chunk only re-parses the newly-appended bytes. Without
    // this, a malformed stream (no `\n\n` until buffer-cap) parses the
    // entire growing buffer every iteration → O(n²) over the run.
    let sseBufScannedUpTo = 0;
    let streamCompletedOk = false;
    try {
      for await (const chunk of body2 as unknown as AsyncIterable<Uint8Array>) {
        if (res.writableEnded) break;
        lastChunkAt = Date.now();
        res.write(chunk);
        // Side-effect: peek into the chunk for usage. We never modify
        // `chunk` — the write above already passed the original bytes
        // through unchanged. A bug here cannot corrupt the stream.
        try {
          const piece = textDecoder.decode(chunk, { stream: true });
          sseBuf += piece;
          // Process complete frames (separated by \n\n) and keep the
          // trailing partial fragment for the next iteration. We only
          // scan from `sseBufScannedUpTo` onward (minus a small overlap
          // so a `data:` frame that straddles the previous chunk
          // boundary still parses).
          const lastBoundary = sseBuf.lastIndexOf('\n\n');
          if (lastBoundary >= 0) {
            // Re-scan a small overlap so a frame whose `\n\n` arrives in
            // the current chunk but whose `data:` line started in the
            // previous chunk is still found.
            const scanStart = Math.max(0, Math.min(sseBufScannedUpTo, lastBoundary + 2) - 256);
            const complete = sseBuf.slice(scanStart, lastBoundary + 2);
            const usage = extractUsageFromSseChunk(complete);
            if (usage) observedUsage = usage;
            sseBuf = sseBuf.slice(lastBoundary + 2);
            sseBufScannedUpTo = 0; // buffer was sliced; reset offset
          } else {
            sseBufScannedUpTo = sseBuf.length;
          }
          // Bound buffer growth in case the upstream never emits \n\n
          // (malformed streaming response). Before dropping the oldest
          // half, scan the discard region for any `data:` frames that
          // might contain a usage payload — otherwise an upstream that
          // emits a giant `[DONE]`-less usage frame would silently lose
          // its bookkeeping.
          if (sseBuf.length > SSE_BUF_MAX) {
            const discardEnd = sseBuf.length - SSE_BUF_MAX / 2;
            const discardPart = sseBuf.slice(0, discardEnd);
            try {
              const usageInDiscard = extractUsageFromSseChunk(discardPart);
              if (usageInDiscard) observedUsage = usageInDiscard;
            } catch { /* extraction is best-effort */ }
            logger.warn(
              `[gateway-stream] SSE buffer exceeded ${SSE_BUF_MAX} team=${team} backend=${backend.id} discarded_bytes=${discardPart.length}`,
            );
            sseBuf = sseBuf.slice(discardEnd);
            sseBufScannedUpTo = sseBuf.length;
          }
        } catch {
          /* extraction is best-effort */
        }
      }
      // Final flush of any leftover bytes — most upstreams send a final
      // `[DONE]` frame after the usage payload, so the usage was already
      // captured above; this is a defensive sweep.
      try {
        const tail = textDecoder.decode();
        if (tail) sseBuf += tail;
        if (sseBuf.length > 0) {
          const usage = extractUsageFromSseChunk(sseBuf);
          if (usage) observedUsage = usage;
        }
      } catch { /* ignore */ }
      if (!res.writableEnded) res.end();
      streamCompletedOk = true;
    } catch (err) {
      const msg = err instanceof Error ? err.message : String(err);
      if (ctrl.signal.aborted) {
        // The abort came from us (request-wide timeout, TTFB timeout,
        // idle timeout, client close, or — Phase 3a F8 — graceful
        // shutdown). Push an SSE-formatted error event so the client
        // sees a structured failure instead of a silently truncated
        // stream followed by a clean close. The shutdown branch uses a
        // distinct event type so the client knows the gateway is
        // intentionally draining and can retry against another
        // instance instead of treating it like a timeout (which would
        // imply the backend is sick).
        if (!res.writableEnded) {
          const shuttingDown = deps.streamRegistry?.isShuttingDown() === true;
          const payload = shuttingDown
            ? `data: {"error":{"type":"gateway_shutdown","message":"gateway draining, client should retry"}}\n\n`
            : `data: {"error":{"type":"gateway_timeout","message":"upstream aborted (request/idle/client-close)"}}\n\n`;
          try {
            res.write(payload);
          } catch { /* ignore — connection may already be torn down */ }
        }
      } else {
        logger.warn(`[gateway] stream relay error team=${team} backend=${backend.id} err=${msg}`);
      }
      if (!res.writableEnded) {
        try { res.end(); } catch { /* ignore */ }
      }
    } finally {
      clearInterval(idleTimer);
      clearTimeout(requestTimer);
      res.off('close', onClientClose);
      decInflightOnce();
      // F8: drop the stream from the shutdown registry — even on the
      // happy path. Without this we'd leak handle references at the
      // registry forever, defeating its memory bound.
      if (unregisterStream) {
        try { unregisterStream(); } catch { /* ignore */ }
      }
      recordUsageIfDb();
      const ms = Date.now() - t0;
      const usageStr = observedUsage
        ? ` tokens_in=${observedUsage.promptTokens} tokens_out=${observedUsage.completionTokens}`
        : '';
      logger.info(`[gateway] team=${team} backend=${backend.id} model=${model} streaming=true status=${upstreamRes.status} duration_ms=${ms}${usageStr}`);
      emitRequestMetric(deps.metrics, {
        team, backend: backend.id, model,
        // Phase 3b post-review: pass resWritableEnded too — a clean
        // SSE close where for-await exited normally AND res ended
        // must classify as success, even when a late client-close
        // event aborted ctrl after the fact.
        status: classifyResponseStatus({
          upstreamStatus: upstreamRes.status,
          ctrl,
          streamRegistry: deps.streamRegistry,
          bodyReadOk: streamCompletedOk,
          resWritableEnded: res.writableEnded,
        }),
        durationMs: ms,
        tokensIn: observedUsage?.promptTokens ?? 0,
        tokensOut: observedUsage?.completionTokens ?? 0,
      });
    }
  };
}

/**
 * Map the upstream HTTP status + abort state to the closed status enum
 * used by `requests_total`. Streaming and non-streaming branches share
 * this helper so labels stay consistent across both code paths.
 *
 * Priority order:
 *   1. `bodyReadOk=true` AND 2xx upstream → success. We trust the body
 *      completion signal over a late-firing abort (Express emits
 *      `'close'` on the request after a successful send too, which can
 *      trigger ctrl.abort in the onClientClose handler before our
 *      finally block reads the flag — supertest reliably hits this).
 *      Phase 3b post-review: also accept `res.writableEnded` as a
 *      stronger success signal when streaming — once the response
 *      socket is fully flushed there is no "abort" to record, even if
 *      the underlying fetch's signal was tripped by a late close
 *      event.
 *   2. Abort + shutdown registry says shutting down → gateway_shutdown.
 *   3. Abort otherwise → gateway_timeout.
 *   4. Non-2xx upstream → upstream_error.
 */
export interface ClassifyResponseStatusInput {
  upstreamStatus: number;
  ctrl: AbortController;
  streamRegistry?: StreamRegistry | undefined;
  /**
   * True when the response body was fully drained / streamed without
   * a thrown error. Streaming branch: set after the for-await loop
   * completes; non-streaming branch: set after upstreamRes.text()
   * resolved and res.send() ran.
   */
  bodyReadOk?: boolean;
  /**
   * Streaming-only: set when `res.writableEnded` after the loop. A
   * late-firing client `close` event with res already ended must NOT
   * be classified as an abort.
   */
  resWritableEnded?: boolean;
}

export function classifyResponseStatus(
  arg1: number | ClassifyResponseStatusInput,
  ctrl?: AbortController,
  streamRegistry?: StreamRegistry | undefined,
  bodyReadOk: boolean = false,
): GatewayRequestStatus {
  // Backwards-compatible positional signature (used by tests). New
  // call sites should pass an object so we can include the
  // resWritableEnded signal.
  const input: ClassifyResponseStatusInput =
    typeof arg1 === 'number'
      ? {
          upstreamStatus: arg1,
          ctrl: ctrl!,
          streamRegistry,
          bodyReadOk,
        }
      : arg1;
  const { upstreamStatus, ctrl: c, streamRegistry: reg } = input;
  const ok = upstreamStatus >= 200 && upstreamStatus < 400;
  const successSignal = (input.bodyReadOk ?? false) || (input.resWritableEnded ?? false);
  if (successSignal && ok) return 'success';
  if (c.signal.aborted) {
    return reg?.isShuttingDown() === true ? 'gateway_shutdown' : 'gateway_timeout';
  }
  if (ok) return 'success';
  return 'upstream_error';
}

function stripTrailingSlash(s: string): string {
  return s.replace(/\/+$/, '');
}

/**
 * Phase 3b: best-effort emission of the `requests_total` /
 * `request_duration_seconds` counter+histogram pair. Always swallows
 * its own errors — a prom-client bug must never propagate into the
 * response.
 *
 * Tokens are emitted only when usage > 0, so a label set with both
 * directions = 0 doesn't spam the time series space.
 */
export function emitRequestMetric(
  metrics: GatewayMetrics | undefined,
  params: {
    team: string;
    backend: string;
    model: string;
    status: GatewayRequestStatus;
    durationMs: number;
    tokensIn?: number;
    tokensOut?: number;
  },
): void {
  if (!metrics) return;
  const { team, backend, model, status, durationMs, tokensIn = 0, tokensOut = 0 } = params;
  try {
    metrics.requestsTotal.labels({ team, backend, model, status }).inc();
    metrics.requestDurationSeconds.labels({ team, backend, model, status }).observe(durationMs / 1000);
    if (tokensIn > 0) {
      metrics.tokensTotal.labels({ team, backend, model, direction: 'in' }).inc(tokensIn);
    }
    if (tokensOut > 0) {
      metrics.tokensTotal.labels({ team, backend, model, direction: 'out' }).inc(tokensOut);
    }
  } catch {
    // Metrics MUST NOT affect client responses.
  }
}

