import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { mkdtempSync, rmSync } from 'fs';
import { join } from 'path';
import { tmpdir } from 'os';
import { Repository } from '../db/repository.js';
import { collectWorkerStatuses } from './dashboard-workers.js';
import type { WorkerDef } from '../config.js';

describe('collectWorkerStatuses', () => {
  let tmpDir: string;
  let repo: Repository;

  beforeEach(() => {
    tmpDir = mkdtempSync(join(tmpdir(), 'dashboard-workers-test-'));
    repo = new Repository(join(tmpDir, 'test.db'));
  });

  afterEach(() => {
    rmSync(tmpDir, { recursive: true, force: true });
  });

  it('returns idle for all workers when no running jobs exist', async () => {
    const workers: WorkerDef[] = [
      { id: 'w1', endpoint: 'x', roles: ['task'] },
      { id: 'w2', endpoint: 'y', roles: ['title'] },
    ];
    const result = await collectWorkerStatuses(repo, workers);
    expect(result).toEqual([
      { id: 'w1', name: 'w1', roles: ['task'], state: 'idle', proxy: false },
      { id: 'w2', name: 'w2', roles: ['title'], state: 'idle', proxy: false },
    ]);
  });

  it('returns running for workers with active jobs', async () => {
    // Seed a running job for w1 via Repository's public API
    const j = await repo.createJob({ repo: 'local/task-1', issueNumber: 1, instruction: 'seed' });
    await repo.updateJob(j.id, { status: 'running', workerId: 'w1' });
    const workers: WorkerDef[] = [
      { id: 'w1', endpoint: 'x', roles: ['task'] },
      { id: 'w2', endpoint: 'y', roles: ['task'] },
    ];
    const result = await collectWorkerStatuses(repo, workers);
    expect(result.find(w => w.id === 'w1')!.state).toBe('running');
    expect(result.find(w => w.id === 'w2')!.state).toBe('idle');
  });

  it('does not leak job id/title/owner in the response shape', async () => {
    const j = await repo.createJob({ repo: 'local/task-1', issueNumber: 1, instruction: 'seed' });
    await repo.updateJob(j.id, { status: 'running', workerId: 'w1' });
    const result = await collectWorkerStatuses(repo, [{ id: 'w1', endpoint: 'x' }]);
    const keys = Object.keys(result[0]!).sort();
    // proxy is a new keyed field; backends is optional and absent for
    // direct workers — that's part of the public shape contract.
    expect(keys).toEqual(['id', 'name', 'proxy', 'roles', 'state']);
  });

  it('fans out proxy workers into backends[] when a registry is supplied', async () => {
    const fakeRegistry = {
      getAll: () => [
        // Self-row that the proxy probe surfaces with nodeId === workerId —
        // must be filtered out so the proxy doesn't appear as its own child.
        {
          nodeId: 'gw', workerId: 'gw', source: 'proxy' as const,
          online: true, busy: false, busySlots: 0, totalSlots: 0,
          loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        },
        {
          nodeId: 'backend-a', workerId: 'gw', source: 'proxy' as const,
          online: true, busy: true, busySlots: 2, totalSlots: 4,
          loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        },
        {
          nodeId: 'backend-b', workerId: 'gw', source: 'proxy' as const,
          online: true, busy: false, busySlots: 0, totalSlots: 4,
          loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        },
        // Belongs to a different worker — must not leak into gw.backends
        {
          nodeId: 'other-backend', workerId: 'other', source: 'proxy' as const,
          online: true, busy: false, busySlots: 0, totalSlots: 4,
          loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        },
      ],
    };
    const workers: WorkerDef[] = [
      { id: 'gw', endpoint: 'http://gw/v1', proxy: true },
      { id: 'direct', endpoint: 'http://direct/v1' },
    ];
    const result = await collectWorkerStatuses(repo, workers, fakeRegistry);
    const gw = result.find(w => w.id === 'gw')!;
    expect(gw.proxy).toBe(true);
    expect(gw.backends?.map(b => b.id).sort()).toEqual(['backend-a', 'backend-b']);
    expect(gw.backends?.find(b => b.id === 'backend-a')).toMatchObject({
      state: 'running', busySlots: 2, totalSlots: 4, online: true,
    });
    expect(gw.backends?.find(b => b.id === 'backend-b')).toMatchObject({
      state: 'idle', busySlots: 0,
    });
    const direct = result.find(w => w.id === 'direct')!;
    expect(direct.proxy).toBe(false);
    // direct workers MUST omit `backends` (undefined, not empty) so the
    // UI can distinguish "direct" from "proxy with zero backends".
    expect(direct.backends).toBeUndefined();
  });

  it('omits backends[] for proxy workers when no registry is supplied (back-compat)', async () => {
    const result = await collectWorkerStatuses(repo, [
      { id: 'gw', endpoint: 'http://gw/v1', proxy: true },
    ], null);
    expect(result[0]!.proxy).toBe(true);
    expect(result[0]!.backends).toBeUndefined();
  });

  it('surfaces busy/total slots on the row for direct workers via the registry self-row', async () => {
    // Direct workers don't have a backends[] expansion, so the slot
    // pressure has to live at the row level — otherwise the Worker
    // widget can't show "(1/3)" for them the way proxy backends do.
    const fakeRegistry = {
      getAll: () => [
        {
          nodeId: 'gpu-1', workerId: 'gpu-1', source: 'direct' as const,
          online: true, busy: true, busySlots: 1, totalSlots: 3,
          loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        },
      ],
    };
    const workers: WorkerDef[] = [
      { id: 'gpu-1', endpoint: 'http://gpu-1:8080/v1' },
    ];
    const result = await collectWorkerStatuses(repo, workers, fakeRegistry);
    const row = result[0]!;
    expect(row.proxy).toBe(false);
    expect(row.backends).toBeUndefined();
    expect(row.busySlots).toBe(1);
    expect(row.totalSlots).toBe(3);
    expect(row.online).toBe(true);
    // Probe-derived state wins over the local jobs-table check when
    // the probe sees in-flight requests this AAO didn't dispatch.
    expect(row.state).toBe('running');
  });

  it('keeps direct worker slot fields undefined when registry has no matching row', async () => {
    const fakeRegistry = {
      // Registry knows other workers, but not the one we're asking about.
      getAll: () => [{
        nodeId: 'other', workerId: 'other', source: 'direct' as const,
        online: true, busy: false, busySlots: 0, totalSlots: 4,
        loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
      }],
    };
    const workers: WorkerDef[] = [
      { id: 'gpu-1', endpoint: 'http://gpu-1:8080/v1' },
    ];
    const result = await collectWorkerStatuses(repo, workers, fakeRegistry);
    const row = result[0]!;
    expect(row.busySlots).toBeUndefined();
    expect(row.totalSlots).toBeUndefined();
    expect(row.online).toBeUndefined();
  });

  it('marks online=false for direct workers when probe failed', async () => {
    const fakeRegistry = {
      getAll: () => [{
        nodeId: 'gpu-1', workerId: 'gpu-1', source: 'direct' as const,
        online: false, busy: false, busySlots: 0, totalSlots: 0,
        loadedModel: null, throughputTps: null, lastSeen: '2026-05-21T00:00:00Z',
        lastProbeError: 'connection refused',
      }],
    };
    const result = await collectWorkerStatuses(repo, [{ id: 'gpu-1', endpoint: 'x' }], fakeRegistry);
    expect(result[0]!.online).toBe(false);
  });
});