/**
 * Router selection tests.
 *
 * Covered:
 *  - strict model match (Phase 1 has no alias)
 *  - drops offline / saturated backends
 *  - picks lowest busy ratio
 *  - round-robin on tied ratios
 *  - null when no candidate
 *  - first-request-before-warm behaviour (no status row yet)
 */
import { describe, it, expect } from 'vitest';
import { createRouter } from './router.js';
import { createBackendInflightCounter } from './stream-proxy.js';
import type { GatewayBackendConfig } from './config.js';
import type { NodeStatus, BackendStatusRegistry } from '../engine/backend-status-registry.js';

function status(id: string, online: boolean, busy: number, total: number): NodeStatus {
  return {
    nodeId: id,
    workerId: id,
    source: 'direct',
    online,
    busy: busy > 0,
    busySlots: busy,
    totalSlots: total,
    loadedModel: null,
    throughputTps: null,
    lastSeen: new Date().toISOString(),
  };
}

function makeRegistry(rows: NodeStatus[]): Pick<BackendStatusRegistry, 'getAll' | 'noteSubscriberActivity'> {
  return {
    getAll: () => rows.slice(),
    noteSubscriberActivity: () => undefined,
  };
}

function bk(id: string, model: string, maxSlots = 4): GatewayBackendConfig {
  return { id, endpoint: `http://${id}/v1`, model, maxSlots };
}

function bkRoles(id: string, model: string, roles: string[], maxSlots = 4): GatewayBackendConfig {
  return { id, endpoint: `http://${id}/v1`, model, maxSlots, roles };
}

describe('createRouter.pick', () => {
  it('returns null when model has no configured backend', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b')],
      registry: makeRegistry([status('a', true, 0, 4)]),
    });
    expect(r.pick('llama:70b')).toBeNull();
  });

  it('returns null when only candidate is offline', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b')],
      registry: makeRegistry([status('a', false, 0, 4)]),
    });
    expect(r.pick('qwen3:8b')).toBeNull();
  });

  it('returns null when all candidates saturated', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', true, 4, 4),
        status('b', true, 4, 4),
      ]),
    });
    expect(r.pick('qwen3:8b')).toBeNull();
  });

  it('picks lowest ratio first', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', true, 3, 4), // ratio 0.75
        status('b', true, 1, 4), // ratio 0.25 (winner)
      ]),
    });
    expect(r.pick('qwen3:8b')?.id).toBe('b');
  });

  it('round-robins on ties', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b'), bk('c', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', true, 1, 4),
        status('b', true, 1, 4),
        status('c', true, 1, 4),
      ]),
    });
    // Walk through twice the tie group so we see the cursor wrap.
    const picks: string[] = [];
    for (let i = 0; i < 6; i++) picks.push(r.pick('qwen3:8b')!.id);
    // Each backend should appear exactly twice over 6 picks.
    const counts = picks.reduce<Record<string, number>>((acc, id) => {
      acc[id] = (acc[id] ?? 0) + 1;
      return acc;
    }, {});
    expect(counts).toEqual({ a: 2, b: 2, c: 2 });
  });

  // ── Sticky routing (x-aao-preferred-backend) ────────────────────────────

  it('honors the preferred backend even when another is idler', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', true, 3, 4), // busier, but preferred (KV cache lives here)
        status('b', true, 0, 4), // idler
      ]),
    });
    expect(r.pick('qwen3:8b', 'a')?.id).toBe('a');
  });

  it('falls back to normal ranking when the preferred backend is saturated', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', true, 4, 4), // preferred but full
        status('b', true, 1, 4),
      ]),
    });
    expect(r.pick('qwen3:8b', 'a')?.id).toBe('b');
  });

  it('falls back when the preferred backend is offline', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([
        status('a', false, 0, 4),
        status('b', true, 1, 4),
      ]),
    });
    expect(r.pick('qwen3:8b', 'a')?.id).toBe('b');
  });

  it('ignores a preferred backend that does not serve the requested model/role', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('x', 'llama:70b')],
      registry: makeRegistry([
        status('a', true, 2, 4),
        status('x', true, 0, 4), // idle but wrong model — must not leak across pools
      ]),
    });
    expect(r.pick('qwen3:8b', 'x')?.id).toBe('a');
  });

  it('preferred pick still reserves an inflight slot', () => {
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 2)],
      registry: makeRegistry([status('a', true, 0, 2)]),
      inflight,
    });
    expect(r.pick('qwen3:8b', 'a')?.id).toBe('a');
    expect(inflight.get('a')).toBe(1);
    expect(r.pick('qwen3:8b', 'a')?.id).toBe('a');
    expect(inflight.get('a')).toBe(2);
    // Saturated by reservations → preferred no longer admitted → null
    expect(r.pick('qwen3:8b', 'a')).toBeNull();
  });

  it('treats registry-cold backends as idle (avoids startup outage)', () => {
    // Registry has nothing yet — first request still routes.
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b')],
      registry: makeRegistry([]),
    });
    expect(r.pick('qwen3:8b')?.id).toBe('a');
  });

  it('prefers a warm (probed-alive) backend over a cold (unprobed) one', () => {
    // Before the fix: cold (ratio 0) beats warm (ratio 0.25) because
    // both go in one list and cold sorts first. After: cold is only a
    // fallback when no warm candidate exists.
    const r = createRouter({
      getBackends: () => [bk('warm', 'qwen3:8b'), bk('cold', 'qwen3:8b')],
      registry: makeRegistry([status('warm', true, 1, 4)]),
    });
    // 10 picks — never pick cold while warm has capacity.
    for (let i = 0; i < 10; i++) {
      expect(r.pick('qwen3:8b')?.id).toBe('warm');
    }
  });

  it('falls back to cold round-robin when all candidates are cold', () => {
    // Two unprobed backends — round-robin across them rather than
    // always returning the first by id.
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:8b')],
      registry: makeRegistry([]),
    });
    const picks: string[] = [];
    for (let i = 0; i < 4; i++) picks.push(r.pick('qwen3:8b')!.id);
    const counts = picks.reduce<Record<string, number>>((acc, id) => {
      acc[id] = (acc[id] ?? 0) + 1;
      return acc;
    }, {});
    expect(counts).toEqual({ a: 2, b: 2 });
  });

  it('calls registry.getAll() at most once per pick (no O(N²) lookup)', () => {
    let getAllCalls = 0;
    const rows = [
      status('a', true, 1, 4),
      status('b', true, 2, 4),
      status('c', true, 3, 4),
    ];
    const registry: Pick<BackendStatusRegistry, 'getAll' | 'noteSubscriberActivity'> = {
      getAll: () => { getAllCalls++; return rows.slice(); },
      noteSubscriberActivity: () => undefined,
    };
    const r = createRouter({
      getBackends: () => [
        bk('a', 'qwen3:8b'),
        bk('b', 'qwen3:8b'),
        bk('c', 'qwen3:8b'),
      ],
      registry,
    });
    r.pick('qwen3:8b');
    // Before the fix this would be 3 (one per backend in the loop).
    expect(getAllCalls).toBe(1);
  });

  it('snapshot returns one entry per backend with attached status', () => {
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b'), bk('b', 'qwen3:14b')],
      registry: makeRegistry([status('a', true, 2, 4)]),
    });
    const snap = r.snapshot();
    expect(snap).toHaveLength(2);
    expect(snap[0]!.status?.online).toBe(true);
    expect(snap[1]!.status).toBeNull();
  });

  it('matches request.model against backend.id (2026-05-21 dogfood regression)', () => {
    // Reproduces the production 503 found during Phase F dogfooding:
    // /v1/models + /health expose `id: backend.id`, the worker echoes
    // that id back as request.model, and the previous router (which
    // filtered only on b.model) couldn't find a match.
    const r = createRouter({
      getBackends: () => [bk('dogfood-gpu-1', 'Qwen3.6-35B-A3B-UD-IQ2_M.gguf')],
      registry: makeRegistry([status('dogfood-gpu-1', true, 0, 4)]),
    });
    expect(r.pick('dogfood-gpu-1')?.id).toBe('dogfood-gpu-1');
  });

  it('still matches request.model against backend.model (legacy callers)', () => {
    const r = createRouter({
      getBackends: () => [bk('dogfood-gpu-1', 'Qwen3.6-35B-A3B-UD-IQ2_M.gguf')],
      registry: makeRegistry([status('dogfood-gpu-1', true, 0, 4)]),
    });
    expect(r.pick('Qwen3.6-35B-A3B-UD-IQ2_M.gguf')?.id).toBe('dogfood-gpu-1');
  });
});

describe('createRouter.pick — role-based routing (2026-06-09)', () => {
  it('routes a role key only to backends that declare the role', () => {
    // Heterogeneous pool: two GPUs serve the quality tier with DIFFERENT
    // model names; one serves fast. A request for "quality" must never
    // land on the fast-only backend even though it is the most idle.
    const r = createRouter({
      getBackends: () => [
        bkRoles('q1', 'big-model-a.gguf', ['quality']),
        bkRoles('q2', 'big-model-b.gguf', ['quality']),
        bkRoles('f1', 'small-model.gguf', ['fast']),
      ],
      registry: makeRegistry([
        status('q1', true, 3, 4), // ratio 0.75
        status('q2', true, 1, 4), // ratio 0.25 (winner within the role)
        status('f1', true, 0, 4), // idle but WRONG role
      ]),
    });
    expect(r.pick('quality')?.id).toBe('q2');
  });

  it('returns null (503) when every backend serving the role is saturated — no spill to other tiers', () => {
    const r = createRouter({
      getBackends: () => [
        bkRoles('q1', 'a.gguf', ['quality'], 2),
        bkRoles('f1', 'b.gguf', ['fast'], 2),
      ],
      registry: makeRegistry([
        status('q1', true, 2, 2), // quality full
        status('f1', true, 0, 2), // fast idle — must NOT be picked for quality
      ]),
    });
    expect(r.pick('quality')).toBeNull();
  });

  it('enforces per-id max_slots independently within a role', () => {
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [
        bkRoles('q1', 'a.gguf', ['quality'], 1),
        bkRoles('q2', 'b.gguf', ['quality'], 1),
      ],
      registry: makeRegistry([status('q1', true, 0, 1), status('q2', true, 0, 1)]),
      inflight,
    });
    const p1 = r.pick('quality')!.id;
    const p2 = r.pick('quality')!.id;
    expect(new Set([p1, p2])).toEqual(new Set(['q1', 'q2'])); // one each
    expect(r.pick('quality')).toBeNull(); // both reserved
  });

  it('treats a backend with unspecified roles as serving every role', () => {
    const r = createRouter({
      getBackends: () => [bk('generalist', 'm.gguf')], // no roles => all
      registry: makeRegistry([status('generalist', true, 0, 4)]),
    });
    expect(r.pick('auto')?.id).toBe('generalist');
    expect(r.pick('fast')?.id).toBe('generalist');
    expect(r.pick('quality')?.id).toBe('generalist');
  });

  it('falls back to id/model exact match when NO backend declares the role (backward compat)', () => {
    // A legacy caller sends a literal model name. No backend declares it
    // as a role, so the router falls back to strict id/model matching.
    const r = createRouter({
      getBackends: () => [bkRoles('q1', 'legacy-model.gguf', ['quality'])],
      registry: makeRegistry([status('q1', true, 0, 4)]),
    });
    expect(r.pick('legacy-model.gguf')?.id).toBe('q1'); // model-name fallback
    expect(r.pick('q1')?.id).toBe('q1'); // id fallback
  });

  it('prefers role match over an id/model collision (role takes precedence)', () => {
    // 'quality' is BOTH a declared role on roleBackend AND the id of a
    // different backend. Role membership wins per decision 3.
    const r = createRouter({
      getBackends: () => [
        bkRoles('roleBackend', 'm-a.gguf', ['quality']),
        bkRoles('quality', 'm-b.gguf', ['fast']), // id literally "quality" but serves fast — the trap
      ],
      registry: makeRegistry([
        status('roleBackend', true, 0, 4),
        status('quality', true, 0, 4),
      ]),
    });
    // Only the role-declaring backend is eligible; the id-collision one is ignored.
    for (let i = 0; i < 5; i++) expect(r.pick('quality')?.id).toBe('roleBackend');
  });
});

describe('createRouter.pick — live reservation (per-host max_slots enforcement)', () => {
  it('reserves a slot on pick so a maxSlots=1 backend saturates after one pick', () => {
    // Root-cause regression: before the fix, pick() judged saturation
    // only on the 5–30s-stale probe snapshot (busySlots), so a burst of
    // requests arriving inside one probe interval all read busySlots=0
    // and piled onto the same host past max_slots. The reservation makes
    // the second pick see the first request immediately.
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 1)],
      registry: makeRegistry([status('a', true, 0, 1)]), // probe says idle
      inflight,
    });
    expect(r.pick('qwen3:8b')?.id).toBe('a'); // reserves slot 1/1
    expect(inflight.get('a')).toBe(1);
    expect(r.pick('qwen3:8b')).toBeNull(); // saturated by the live reservation, NOT the stale probe
  });

  it('treats live reservations as busy via max(probed, reserved)', () => {
    // total=2, probe still reports 0 busy, but 2 requests are already
    // reserved in-flight → the host is full even though the probe lags.
    const inflight = createBackendInflightCounter();
    inflight.inc('a');
    inflight.inc('a');
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 2)],
      registry: makeRegistry([status('a', true, 0, 2)]),
      inflight,
    });
    expect(r.pick('qwen3:8b')).toBeNull();
  });

  it('spreads a burst across two backends instead of stacking one host', () => {
    // Both probed idle (busy 0/1). Without reservations the stale probe
    // would let the least-busy sort send every request to the same id.
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 1), bk('b', 'qwen3:8b', 1)],
      registry: makeRegistry([status('a', true, 0, 1), status('b', true, 0, 1)]),
      inflight,
    });
    const p1 = r.pick('qwen3:8b')!.id;
    const p2 = r.pick('qwen3:8b')!.id;
    expect(new Set([p1, p2])).toEqual(new Set(['a', 'b'])); // one each
    expect(r.pick('qwen3:8b')).toBeNull(); // both now full
  });

  it('gates registry-cold (unprobed) backends by maxSlots reservations', () => {
    // Boot warm-up window: no probe row yet. The cap must still hold off
    // the live reservation count, not wait for the first probe.
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 1)],
      registry: makeRegistry([]), // cold
      inflight,
    });
    expect(r.pick('qwen3:8b')?.id).toBe('a');
    expect(r.pick('qwen3:8b')).toBeNull();
  });

  it('releasing a reservation (dec) frees the slot for the next pick', () => {
    const inflight = createBackendInflightCounter();
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 1)],
      registry: makeRegistry([status('a', true, 0, 1)]),
      inflight,
    });
    expect(r.pick('qwen3:8b')?.id).toBe('a');
    expect(r.pick('qwen3:8b')).toBeNull();
    inflight.dec('a'); // stream-proxy finally would do this when the request ends
    expect(r.pick('qwen3:8b')?.id).toBe('a');
  });

  it('without an inflight dep, falls back to probe-only behaviour (no reservation)', () => {
    // Legacy/unit callers that do not wire inflight must keep working
    // exactly as before — no hidden reservation side effect.
    const r = createRouter({
      getBackends: () => [bk('a', 'qwen3:8b', 1)],
      registry: makeRegistry([status('a', true, 0, 1)]),
    });
    expect(r.pick('qwen3:8b')?.id).toBe('a');
    expect(r.pick('qwen3:8b')?.id).toBe('a'); // still routes — probe still says idle
  });
});
