/**
 * Stream proxy tests using a fetch-impl mock and a fake Router.
 *
 * What we verify:
 *  - 400 on missing model field
 *  - 503 when router returns null (no available backend)
 *  - 502 when upstream fetch rejects
 *  - non-streaming pass-through preserves headers (x-aao-backend-id,
 *    x-litellm-model-id, status code)
 *  - SSE pass-through: chunks reach the client + abort propagates on
 *    client close
 *  - Authorization header forwarded when backend.api_key set
 */
import { describe, it, expect, vi } from 'vitest';
import { buildChatCompletionsHandler, createStreamRegistry, createBackendInflightCounter } from './stream-proxy.js';
import { createRouter } from './router.js';
import type { Router } from './router.js';
import type { GatewayBackendConfig } from './config.js';
import type { NodeStatus, BackendStatusRegistry } from '../engine/backend-status-registry.js';
import { EventEmitter } from 'events';

function makeBackend(overrides: Partial<GatewayBackendConfig> = {}): GatewayBackendConfig {
  return {
    id: 'gpu-a',
    endpoint: 'http://gpu-a/v1',
    model: 'qwen3:8b',
    maxSlots: 4,
    ...overrides,
  };
}

function makeRouter(backend: GatewayBackendConfig | null): Router {
  return {
    pick: () => backend,
    snapshot: () => [],
  };
}

function makeReq(body: unknown, extras: Record<string, unknown> = {}): import('express').Request {
  // EventEmitter so req.on('close') / req.off work.
  const e = new EventEmitter();
  return Object.assign(e, {
    body,
    headers: {},
    ip: '127.0.0.1',
    path: '/v1/chat/completions',
    method: 'POST',
    ...extras,
  }) as unknown as import('express').Request;
}

function makeRes(): {
  res: import('express').Response;
  code: () => number | undefined;
  body: () => unknown;
  written: () => string;
  headers: Record<string, string>;
  ended: () => boolean;
  endedFlag: { v: boolean };
} {
  const chunks: Buffer[] = [];
  const headers: Record<string, string> = {};
  const endedFlag = { v: false };
  let code: number | undefined;
  let jsonBody: unknown;
  let sentText: string | undefined;
  // Extend an EventEmitter so the handler's `res.on('close', ...)`
  // wiring (added after the Node-18+ req.close-on-drain regression
  // fix) can fire close events from tests.
  const emitter = new EventEmitter();
  const res = Object.assign(emitter, {
    status(c: number) {
      code = c;
      return this;
    },
    setHeader(k: string, v: string) {
      headers[k.toLowerCase()] = v;
      return this;
    },
    json(b: unknown) {
      jsonBody = b;
      endedFlag.v = true;
      return this;
    },
    send(b: string) {
      sentText = b;
      endedFlag.v = true;
      return this;
    },
    write(b: Buffer | string) {
      chunks.push(Buffer.from(b as Buffer));
      return true;
    },
    end() {
      endedFlag.v = true;
      return this;
    },
    get writableEnded() {
      return endedFlag.v;
    },
    headersSent: false,
  }) as unknown as import('express').Response;
  return {
    res,
    code: () => code,
    body: () => jsonBody ?? sentText,
    written: () => Buffer.concat(chunks).toString('utf8'),
    headers,
    ended: () => endedFlag.v,
    endedFlag,
  };
}

describe('chatCompletionsHandler', () => {
  it('returns 400 when model is missing', async () => {
    const router = makeRouter(makeBackend());
    const handler = buildChatCompletionsHandler({
      router, upstreamTimeoutSec: 5, requestTimeoutSec: 60,
    });
    const req = makeReq({});
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(400);
  });

  it('returns 503 when router has no backend', async () => {
    const router = makeRouter(null);
    const handler = buildChatCompletionsHandler({
      router, upstreamTimeoutSec: 5, requestTimeoutSec: 60,
    });
    const req = makeReq({ model: 'qwen3:8b' });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(503);
    expect(r.body()).toEqual({ error: 'no available backend' });
  });

  it('returns 502 when upstream fetch rejects', async () => {
    const fetchImpl = vi.fn(async () => { throw new Error('ECONNREFUSED'); });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(502);
  });

  it('passes non-streaming response through with backend headers injected', async () => {
    const fakeBody = JSON.stringify({ id: 'cmpl-1', choices: [] });
    const fetchImpl = vi.fn(async () => new Response(fakeBody, {
      status: 200,
      headers: { 'content-type': 'application/json' },
    }));
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend({ id: 'gpu-x' })),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(200);
    expect(r.headers['x-aao-backend-id']).toBe('gpu-x');
    expect(r.headers['x-litellm-model-id']).toBe('gpu-x');
    expect(r.body()).toBe(fakeBody);
  });

  it('does not forward transfer-encoding for non-streaming responses', async () => {
    // Regression: forwarding `transfer-encoding: chunked` together with
    // res.send() (which sets Content-Length) violates HTTP/1.1 and
    // makes nginx/HAProxy drop the connection. Streaming responses
    // still pass it through (separate test below).
    const fakeBody = JSON.stringify({ id: 'cmpl-1', choices: [] });
    const fetchImpl = vi.fn(async () => new Response(fakeBody, {
      status: 200,
      headers: {
        'content-type': 'application/json',
        'transfer-encoding': 'chunked',
      },
    }));
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(200);
    // The non-streaming branch must NOT set transfer-encoding — Express
    // will compute Content-Length on res.send() and the two headers
    // together are illegal per RFC 9112 §6.2.
    expect(r.headers['transfer-encoding']).toBeUndefined();
    // Content-Type is still passed through.
    expect(r.headers['content-type']).toBe('application/json');
  });

  it('still forwards transfer-encoding for streaming responses', async () => {
    const enc = new TextEncoder();
    const stream = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(enc.encode('data: {"x":1}\n\n'));
        controller.close();
      },
    });
    const fetchImpl = vi.fn(async () => new Response(stream, {
      status: 200,
      headers: {
        'content-type': 'text/event-stream',
        'transfer-encoding': 'chunked',
      },
    }));
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: true });
    const r = makeRes();
    await handler(req, r.res);
    // Streaming preserves chunked encoding — the client / nginx parses
    // the framing themselves.
    expect(r.headers['transfer-encoding']).toBe('chunked');
  });

  it('forwards Authorization header when backend.api_key is set', async () => {
    const calls: Array<{ url: string; init: RequestInit }> = [];
    const fetchImpl = vi.fn(async (url: string, init: RequestInit) => {
      calls.push({ url, init });
      return new Response('{}', { status: 200, headers: { 'content-type': 'application/json' } });
    });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend({ apiKey: 'sk-upstream' })),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    await handler(req, r.res);
    const hdrs = calls[0]!.init.headers as Record<string, string>;
    expect(hdrs['Authorization']).toBe('Bearer sk-upstream');
  });

  it('relays SSE chunks to the client', async () => {
    // Build an SSE body as a ReadableStream of UTF-8 Uint8Arrays.
    const enc = new TextEncoder();
    const sse = enc.encode('data: {"choices":[{"delta":{"content":"hi"}}]}\n\n');
    const stream = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(sse);
        controller.close();
      },
    });
    const fetchImpl = vi.fn(async () => new Response(stream, {
      status: 200,
      headers: { 'content-type': 'text/event-stream' },
    }));
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: true });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.code()).toBe(200);
    expect(r.headers['content-type']).toBe('text/event-stream');
    expect(r.written()).toContain('"hi"');
  });

  it('writes an SSE error event when streaming is aborted mid-flight (request timeout)', async () => {
    // Build a stream we control. We never enqueue real data — we just
    // wait for the handler's AbortController to fire, then surface the
    // abort as a stream error so the for-await unblocks.
    let controllerRef!: ReadableStreamDefaultController<Uint8Array>;
    const stream = new ReadableStream<Uint8Array>({
      start(controller) {
        controllerRef = controller;
      },
    });
    const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
      const signal = init.signal as AbortSignal;
      signal.addEventListener('abort', () => {
        try { controllerRef.error(new Error('aborted')); } catch { /* ignore */ }
      });
      return new Response(stream, {
        status: 200,
        headers: { 'content-type': 'text/event-stream' },
      });
    });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      // Tiny request timeout — fires before any chunk arrives.
      requestTimeoutSec: 0.01,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: true });
    const r = makeRes();
    await handler(req, r.res);
    expect(r.ended()).toBe(true);
    // The client must see an SSE error event, not just a clean close.
    const written = r.written();
    expect(written).toContain('"error"');
    expect(written).toContain('gateway_timeout');
  });

  it('returns 504 JSON when non-streaming TTFB timeout fires before headers', async () => {
    // Simulate a fetch that the handler will abort before the body read
    // completes. We resolve the Response, but reading text() throws
    // because the abort fires first.
    const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
      const signal = init.signal as AbortSignal;
      // Build a response whose body read throws on abort.
      const stream = new ReadableStream<Uint8Array>({
        start(controller) {
          signal.addEventListener('abort', () => {
            try { controller.error(new Error('aborted during body read')); } catch { /* ignore */ }
          });
        },
      });
      return new Response(stream, {
        status: 200,
        headers: { 'content-type': 'application/json' },
      });
    });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 0.01,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    // The non-streaming branch calls res.status(upstream.status) before
    // attempting text(); to simulate "headers not yet sent" we override
    // status() to be a no-op and rely on headersSent staying false.
    // But our makeRes() doesn't set headersSent, so this path still
    // exercises the !res.headersSent branch correctly.
    await handler(req, r.res);
    // status() was called twice: once mirroring upstream (200) and once
    // overriding to 504. The last call wins in our fake.
    expect(r.code()).toBe(504);
    const body = r.body();
    expect(body).toMatchObject({ error: expect.stringContaining('gateway timeout') });
  });

  // Phase 3a F8: graceful shutdown drains in-flight SSE with a
  // distinguishable event so AAO clients know to retry against another
  // instance instead of treating the close as a backend timeout.
  describe('gateway_shutdown signaling (F8)', () => {
    it('emits a gateway_shutdown SSE event when the registry signals shutdown mid-stream', async () => {
      let controllerRef!: ReadableStreamDefaultController<Uint8Array>;
      const stream = new ReadableStream<Uint8Array>({
        start(controller) { controllerRef = controller; },
      });
      const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
        const signal = init.signal as AbortSignal;
        signal.addEventListener('abort', () => {
          try { controllerRef.error(new Error('aborted')); } catch { /* ignore */ }
        });
        return new Response(stream, {
          status: 200,
          headers: { 'content-type': 'text/event-stream' },
        });
      });
      const registry = createStreamRegistry();
      const handler = buildChatCompletionsHandler({
        router: makeRouter(makeBackend()),
        upstreamTimeoutSec: 5,
        requestTimeoutSec: 60,
        fetchImpl: fetchImpl as unknown as typeof fetch,
        streamRegistry: registry,
      });
      const req = makeReq({ model: 'qwen3:8b', stream: true });
      const r = makeRes();
      const p = handler(req, r.res);
      // Wait for registration to happen (handler must reach the
      // streaming branch), then fire shutdown.
      await new Promise(resolve => setImmediate(resolve));
      expect(registry.size()).toBe(1);
      registry.signalShutdown();
      await p;
      expect(r.ended()).toBe(true);
      const written = r.written();
      expect(written).toContain('gateway_shutdown');
      expect(written).not.toContain('gateway_timeout');
      // Cleanup: handler must unregister on finally even though the
      // abort came from shutdown.
      expect(registry.size()).toBe(0);
    });

    it('keeps emitting gateway_timeout for non-shutdown aborts when a registry is wired', async () => {
      let controllerRef!: ReadableStreamDefaultController<Uint8Array>;
      const stream = new ReadableStream<Uint8Array>({
        start(controller) { controllerRef = controller; },
      });
      const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
        const signal = init.signal as AbortSignal;
        signal.addEventListener('abort', () => {
          try { controllerRef.error(new Error('aborted')); } catch { /* ignore */ }
        });
        return new Response(stream, {
          status: 200,
          headers: { 'content-type': 'text/event-stream' },
        });
      });
      const registry = createStreamRegistry();
      const handler = buildChatCompletionsHandler({
        router: makeRouter(makeBackend()),
        upstreamTimeoutSec: 5,
        requestTimeoutSec: 0.01, // forces request-timeout abort
        fetchImpl: fetchImpl as unknown as typeof fetch,
        streamRegistry: registry,
      });
      const req = makeReq({ model: 'qwen3:8b', stream: true });
      const r = makeRes();
      await handler(req, r.res);
      const written = r.written();
      expect(written).toContain('gateway_timeout');
      expect(written).not.toContain('gateway_shutdown');
    });

    it('registry unregister fires on the happy path too (no leaks)', async () => {
      const stream = new ReadableStream<Uint8Array>({
        start(c) {
          c.enqueue(new TextEncoder().encode('data: "x"\n\n'));
          c.close();
        },
      });
      const fetchImpl = vi.fn(async () => new Response(stream, {
        status: 200,
        headers: { 'content-type': 'text/event-stream' },
      }));
      const registry = createStreamRegistry();
      const handler = buildChatCompletionsHandler({
        router: makeRouter(makeBackend()),
        upstreamTimeoutSec: 5,
        requestTimeoutSec: 60,
        fetchImpl: fetchImpl as unknown as typeof fetch,
        streamRegistry: registry,
      });
      const req = makeReq({ model: 'qwen3:8b', stream: true });
      const r = makeRes();
      await handler(req, r.res);
      expect(registry.size()).toBe(0);
    });
  });

  it('aborts upstream when client disconnects mid-stream', async () => {
    // Build a stream we can hold open while we fire 'close' on the res.
    let controllerRef!: ReadableStreamDefaultController<Uint8Array>;
    const stream = new ReadableStream<Uint8Array>({
      start(controller) {
        controllerRef = controller;
      },
    });
    const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
      // The handler should pass an AbortSignal. We expose it so we can
      // assert the abort happens.
      const signal = init.signal as AbortSignal;
      // Forward the abort to the stream so the for-await loop unblocks
      // and our handler reaches the finally clause.
      signal.addEventListener('abort', () => {
        try { controllerRef.error(new Error('aborted')); } catch { /* ignore */ }
      });
      return new Response(stream, {
        status: 200,
        headers: { 'content-type': 'text/event-stream' },
      });
    });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: true });
    const r = makeRes();
    // Kick the handler off, then fire response close (= client gave up
    // mid-stream while writableEnded is still false).
    const p = handler(req, r.res);
    setImmediate(() => (r.res as unknown as EventEmitter).emit('close'));
    await p;
    // The handler must have ended the response (no leaked socket).
    expect(r.ended()).toBe(true);
  });

  it('does NOT abort upstream when req emits close after body drain (Node 18+ regression)', async () => {
    // Regression for the Phase 1 (PR #326) bug discovered during the
    // 2026-05-20 dogfooding pass: 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 a response. The previous wiring
    // `req.on('close', () => ctrl.abort())` therefore aborted ctrl
    // synchronously on every real request, turning the gateway into a
    // perma-502. Only supertest's mock pipeline kept the unit tests
    // green.
    //
    // Fix: listen on res.on('close') instead. This test fires the
    // spurious req-close mid-fetch and asserts the response still
    // completes successfully (no abort propagated to upstream).
    let upstreamSignalled = false;
    const fakeBody = JSON.stringify({ id: 'cmpl-1', choices: [{ message: { content: 'ok' } }] });
    const fetchImpl = vi.fn(async (_url: string, init: RequestInit) => {
      const signal = init.signal as AbortSignal;
      signal.addEventListener('abort', () => { upstreamSignalled = true; });
      // Defer one tick so the test can fire req.close before fetch
      // resolves — mirrors the production timing.
      await new Promise(resolve => setImmediate(resolve));
      return new Response(fakeBody, {
        status: 200,
        headers: { 'content-type': 'application/json' },
      });
    });
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false });
    const r = makeRes();
    const p = handler(req, r.res);
    // Fire the spurious req-close that Node 18+ emits on body drain.
    // The handler MUST ignore it because the response stream is still
    // open and writableEnded is false (response not yet sent).
    req.emit('close');
    await p;
    expect(r.code()).toBe(200);
    expect(upstreamSignalled).toBe(false);
    expect(r.body()).toBe(fakeBody);
  });
});

describe('chatCompletionsHandler — per-host concurrency (router reservation + finally release)', () => {
  function warmStatus(id: string, busy: number, total: number): NodeStatus {
    return {
      nodeId: id, workerId: id, source: 'direct', online: true,
      busy: busy > 0, busySlots: busy, totalSlots: total,
      loadedModel: null, throughputTps: null, lastSeen: new Date().toISOString(),
    };
  }
  function staleRegistry(rows: NodeStatus[]): Pick<BackendStatusRegistry, 'getAll' | 'noteSubscriberActivity'> {
    // Returns a FIXED snapshot — models the real 5–30s-stale probe that
    // never updates within the burst window. Before the fix this let
    // every request read busySlots=0 and pile onto one host.
    return { getAll: () => rows.slice(), noteSubscriberActivity: () => undefined };
  }

  it('rejects a concurrent request past max_slots with 503, then frees the slot when the first finishes', async () => {
    const backend = makeBackend({ id: 'gpu-a', maxSlots: 1 });
    const inflight = createBackendInflightCounter();
    // Real router, fed a stale "idle" probe (busy 0 / total 1) that never
    // changes — exactly the condition that produced the over-admission bug.
    const router = createRouter({
      getBackends: () => [backend],
      registry: staleRegistry([warmStatus('gpu-a', 0, 1)]),
      inflight,
    });

    // First request's upstream hangs on a gate so it stays in-flight,
    // holding its reservation, while the second request arrives.
    let releaseUpstream!: () => void;
    const gate = new Promise<void>(res => { releaseUpstream = res; });
    const fakeBody = JSON.stringify({ id: 'cmpl-1', choices: [] });
    const fetchImpl = vi.fn(async () => {
      await gate;
      return new Response(fakeBody, { status: 200, headers: { 'content-type': 'application/json' } });
    });

    const handler = buildChatCompletionsHandler({
      router, upstreamTimeoutSec: 30, requestTimeoutSec: 60,
      fetchImpl: fetchImpl as unknown as typeof fetch,
      inflight,
    });

    // Fire request 1 (do NOT await — it parks at the gated fetch). pick()
    // runs synchronously before the first await, so the reservation is
    // already in place by the time this call returns its promise.
    const r1 = makeRes();
    const p1 = handler(makeReq({ model: 'qwen3:8b', stream: false }), r1.res);
    expect(inflight.get('gpu-a')).toBe(1); // reserved on select, not after probe

    // Request 2 arrives mid-flight. The stale probe STILL says busy=0,
    // but the live reservation makes the host saturated → 503.
    const r2 = makeRes();
    await handler(makeReq({ model: 'qwen3:8b', stream: false }), r2.res);
    expect(r2.code()).toBe(503);
    expect(r2.body()).toEqual({ error: 'no available backend' });

    // Let request 1 complete → finally{} releases the reservation.
    releaseUpstream();
    await p1;
    expect(r1.code()).toBe(200);
    expect(inflight.get('gpu-a')).toBe(0); // released exactly once

    // A later request now routes again — the cap is a live gate, not a latch.
    const r3 = makeRes();
    await handler(makeReq({ model: 'qwen3:8b', stream: false }), r3.res);
    expect(r3.code()).toBe(200);
    expect(inflight.get('gpu-a')).toBe(0);
  });
});
