/**
 * Phase 2b — usage extraction tests for stream-proxy.
 *
 * Verifies:
 *  - extractUsageFromJson on a non-streaming response body
 *  - extractUsageFromSseChunk on a streaming chunk (single + multi-frame)
 *  - extraction failures fall back to {0, 0} but still record the request
 *  - non-DB auth source (config) skips the usage write entirely
 *  - DB-backed auth source records {tokensIn=0, tokensOut=0, requests=1}
 *    when usage cannot be extracted
 */
import { describe, it, expect, vi } from 'vitest';
import { EventEmitter } from 'events';
import { Readable } from 'stream';
import { buildChatCompletionsHandler, extractUsageFromJson, extractUsageFromSseChunk } from './stream-proxy.js';
import type { Router } from './router.js';
import type { GatewayBackendConfig } from './config.js';

function makeBackend(): GatewayBackendConfig {
  return { id: 'gpu-a', endpoint: 'http://gpu-a/v1', model: 'qwen3:8b', maxSlots: 4 };
}
function makeRouter(b: GatewayBackendConfig | null): Router {
  return { pick: () => b, snapshot: () => [] };
}
function makeReq(body: unknown, extras: Record<string, unknown> = {}): import('express').Request {
  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; ended: () => boolean } {
  const endedFlag = { v: false };
  // Extend an EventEmitter so the handler's res.on('close', ...) wiring
  // (added after the Node-18+ req.close-on-drain regression fix) finds
  // a valid emitter to register on.
  const emitter = new EventEmitter();
  const res = Object.assign(emitter, {
    status() { return this; },
    setHeader() { return this; },
    json() { endedFlag.v = true; return this; },
    send() { endedFlag.v = true; return this; },
    write() { return true; },
    end() { endedFlag.v = true; return this; },
    get writableEnded() { return endedFlag.v; },
    headersSent: false,
  }) as unknown as import('express').Response;
  return { res, ended: () => endedFlag.v };
}

describe('extractUsageFromJson (pure)', () => {
  it('returns usage from a well-formed body', () => {
    const text = JSON.stringify({
      id: 'c1', choices: [{}], usage: { prompt_tokens: 100, completion_tokens: 50 },
    });
    expect(extractUsageFromJson(text)).toEqual({ promptTokens: 100, completionTokens: 50 });
  });

  it('handles missing completion_tokens by defaulting it to 0', () => {
    const text = JSON.stringify({ usage: { prompt_tokens: 100 } });
    expect(extractUsageFromJson(text)).toEqual({ promptTokens: 100, completionTokens: 0 });
  });

  it('returns null for malformed JSON', () => {
    expect(extractUsageFromJson('not json')).toBeNull();
  });

  it('returns null when usage field is missing', () => {
    expect(extractUsageFromJson(JSON.stringify({ id: 'x' }))).toBeNull();
  });

  it('returns null when prompt + completion are both non-numeric', () => {
    expect(extractUsageFromJson(JSON.stringify({ usage: { prompt_tokens: 'a' } }))).toBeNull();
  });

  it('floors fractional token counts and rejects negatives', () => {
    expect(extractUsageFromJson(JSON.stringify({ usage: { prompt_tokens: 50.7, completion_tokens: -5 } })))
      .toEqual({ promptTokens: 50, completionTokens: 0 });
  });
});

describe('extractUsageFromSseChunk (pure)', () => {
  it('extracts from a single SSE frame containing usage', () => {
    const chunk = `data: {"id":"c","usage":{"prompt_tokens":120,"completion_tokens":80}}\n\n`;
    expect(extractUsageFromSseChunk(chunk)).toEqual({ promptTokens: 120, completionTokens: 80 });
  });

  it('takes the latest usage value when multiple frames carry one', () => {
    const chunk =
      `data: {"usage":{"prompt_tokens":10,"completion_tokens":5}}\n\n` +
      `data: {"usage":{"prompt_tokens":20,"completion_tokens":15}}\n\n` +
      `data: [DONE]\n\n`;
    expect(extractUsageFromSseChunk(chunk)).toEqual({ promptTokens: 20, completionTokens: 15 });
  });

  it('returns null when no frame carries usage', () => {
    const chunk = `data: {"id":"c","choices":[{"delta":{"content":"hi"}}]}\n\n`;
    expect(extractUsageFromSseChunk(chunk)).toBeNull();
  });

  it('skips [DONE] and empty data lines without throwing', () => {
    const chunk = `data: \n\ndata: [DONE]\n\n`;
    expect(extractUsageFromSseChunk(chunk)).toBeNull();
  });

  it('skips partial JSON without throwing (TCP boundary truncation)', () => {
    const chunk = `data: {"usage":{"prompt_tokens":\n\n`;
    expect(extractUsageFromSseChunk(chunk)).toBeNull();
  });
});

describe('chatCompletionsHandler — usage recording integration', () => {
  it('records usage to the recorder when DB-backed auth handles a non-streaming response', async () => {
    const recorder = { recordUsage: vi.fn() };
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map([['content-type', 'application/json']]),
      text: () => Promise.resolve(JSON.stringify({ usage: { prompt_tokens: 100, completion_tokens: 50 } })),
      get body() { return null; },
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k1' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).toHaveBeenCalledTimes(1);
    expect(recorder.recordUsage).toHaveBeenCalledWith(
      expect.objectContaining({
        keyId: 'k1',
        tokensIn: 100,
        tokensOut: 50,
        requests: 1,
      }),
    );
  });

  it('does NOT record usage when matched key came from config (source=config)', async () => {
    const recorder = { recordUsage: vi.fn() };
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map(),
      text: () => Promise.resolve(JSON.stringify({ usage: { prompt_tokens: 5, completion_tokens: 3 } })),
      get body() { return null; },
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false }, {
      gatewayAuth: { team: 'alpha', source: 'config' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).not.toHaveBeenCalled();
  });

  it('records {tokensIn=0, tokensOut=0, requests=1} when usage extraction fails', async () => {
    const recorder = { recordUsage: vi.fn() };
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map(),
      text: () => Promise.resolve('not json'),
      get body() { return null; },
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k2' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).toHaveBeenCalledWith(
      expect.objectContaining({ keyId: 'k2', tokensIn: 0, tokensOut: 0, requests: 1 }),
    );
  });

  it('records streaming usage from chunks carrying OpenAI-style usage frames', async () => {
    const recorder = { recordUsage: vi.fn() };
    // Build a Readable that emits two SSE frames: a content delta and a
    // final frame containing usage, then closes.
    const body = Readable.from([
      Buffer.from(`data: {"choices":[{"delta":{"content":"hi"}}]}\n\n`),
      Buffer.from(`data: {"usage":{"prompt_tokens":42,"completion_tokens":7}}\n\ndata: [DONE]\n\n`),
    ]);
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map([['content-type', 'text/event-stream']]),
      text: () => Promise.resolve(''),
      body,
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b' /* stream defaults to true */ }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k3' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).toHaveBeenCalledTimes(1);
    expect(recorder.recordUsage).toHaveBeenCalledWith(
      expect.objectContaining({ keyId: 'k3', tokensIn: 42, tokensOut: 7, requests: 1 }),
    );
  });

  it('still extracts usage from the FINAL frame after 100+ content chunks (F8 perf, correctness preserved)', async () => {
    // Build a long stream: many content delta frames, then a usage frame
    // at the very end (the OpenAI convention). F8's scannedUpTo
    // optimization must not skip the final frame.
    const recorder = { recordUsage: vi.fn() };
    const chunks: Buffer[] = [];
    for (let i = 0; i < 120; i++) {
      chunks.push(Buffer.from(`data: {"choices":[{"delta":{"content":"tok${i}"}}]}\n\n`));
    }
    chunks.push(Buffer.from(`data: {"usage":{"prompt_tokens":2000,"completion_tokens":150}}\n\ndata: [DONE]\n\n`));
    const body = Readable.from(chunks);
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map([['content-type', 'text/event-stream']]),
      text: () => Promise.resolve(''),
      body,
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b' }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k-final' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).toHaveBeenCalledTimes(1);
    expect(recorder.recordUsage).toHaveBeenCalledWith(
      expect.objectContaining({ keyId: 'k-final', tokensIn: 2000, tokensOut: 150, requests: 1 }),
    );
  });

  it('recovers usage from the discard region when the SSE buffer exceeds its cap (F1)', async () => {
    // Build a malformed stream: a long stretch of `data:` content lines
    // with no `\n\n` boundary, large enough to blow past the 64 KiB cap.
    // The usage frame sits BEFORE the discard cut so that, without F1,
    // it would be silently truncated and never recorded.
    const recorder = { recordUsage: vi.fn() };
    // 80 KiB of padding lines with terminator missing so lastIndexOf('\n\n')
    // never matches and the discard branch triggers.
    const usageLine = `data: {"usage":{"prompt_tokens":777,"completion_tokens":333}}`;
    const filler = 'data: {"choices":[{"delta":{"content":"x"}}]}\n'.repeat(2000); // ~90 KB, no blank-line separators
    const noisyTail = 'a'.repeat(40_000); // ensure cap exceeded after usage emitted
    const body = Readable.from([
      Buffer.from(usageLine + '\n' + filler),
      Buffer.from(noisyTail),
      Buffer.from('\n\ndata: [DONE]\n\n'),
    ]);
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map([['content-type', 'text/event-stream']]),
      text: () => Promise.resolve(''),
      body,
    } as unknown as Response);

    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b' }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k-discard' },
    });
    const { res } = makeRes();
    await handler(req, res);

    expect(recorder.recordUsage).toHaveBeenCalledTimes(1);
    expect(recorder.recordUsage).toHaveBeenCalledWith(
      expect.objectContaining({ keyId: 'k-discard', tokensIn: 777, tokensOut: 333, requests: 1 }),
    );
  });

  it('swallows recordUsage throws so the response is not affected', async () => {
    const recorder = {
      recordUsage: vi.fn(() => { throw new Error('db busy'); }),
    };
    const fetchImpl = vi.fn().mockResolvedValue({
      status: 200,
      headers: new Map(),
      text: () => Promise.resolve(JSON.stringify({ usage: { prompt_tokens: 1, completion_tokens: 1 } })),
      get body() { return null; },
    } as unknown as Response);
    const handler = buildChatCompletionsHandler({
      router: makeRouter(makeBackend()),
      upstreamTimeoutSec: 5,
      requestTimeoutSec: 10,
      fetchImpl,
      usageRecorder: recorder,
    });
    const req = makeReq({ model: 'qwen3:8b', stream: false }, {
      gatewayAuth: { team: 'alpha', source: 'db', keyId: 'k4' },
    });
    const { res, ended } = makeRes();
    await expect(handler(req, res)).resolves.toBeUndefined();
    expect(ended()).toBe(true);
  });
});
