import { randomUUID } from 'node:crypto';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { newFixture, type Fixture } from '../helpers/fixture.js';
import { prepareNativeHost } from '../helpers/native-host.js';
import type { RebalanceMove } from '../../src/lib/rebalancing.js';

vi.mock('../../src/lib/server/pipeline/worker.js', () => ({ kick: vi.fn(async () => {}), kickProvisioning: vi.fn(async () => {}) }));
vi.mock('../../src/lib/server/integrations/plutonium.js', () => ({ checkRelease: vi.fn(async () => {}), freshRelease: () => 5354, probeVersion: vi.fn(async () => 5354) }));
vi.mock('../../src/lib/server/integrations/rcon.js', () => ({
    rconStatusResult: vi.fn(async () => ({ ok: true, endpoint: { host: 'agent', port: 28960 }, rawReply: '',
        status: { players: 0, map: 'zombie_theater', gametype: 'zom', roster: [] } })),
    rconGetPlayerCount: vi.fn(async () => 0), rconPlutoniumVersion: vi.fn(async () => ({ revision: 5354, error: null }))
}));

let fixture: Fixture, sourceId: string, targetId: string, libraryId: string;
const names = Array.from({ length: 4 }, (_, index) => `t5_zombie_theater_${index + 1}`);
const image = 'example.test/plutainer:rebalance';
const gib = 1024 ** 3;
const clockStep = () => vi.setSystemTime(Date.now() + 1000);
const settings = () => ({ hostIds: [sourceId, targetId], game: 'all' as const, maxMoves: 5 });

async function heartbeat(hostId: string, cpuUsagePercent: number) {
    const { db } = await import('../../src/lib/server/db/index.js');
    const { heartbeat } = await import('../../src/lib/server/hosts/index.js');
    const instances = db().prepare('SELECT id,generation FROM native_instances WHERE host_id=? AND retired_at IS NULL')
        .all(hostId) as { id: string; generation: number }[];
    heartbeat(hostId, { protocolVersion: 1, cpuCount: 16, cpuUsagePercent, memoryTotal: 32 * gib,
        memoryAvailable: 32 * gib - (2 + instances.length * 2) * gib, dockerRunning: true,
        directoriesWritable: true, storageAvailable: 100 * gib, runtimeImages: [image],
        libraries: [{ id: libraryId, state: 'ready' }], instanceResources: instances.map(instance => ({
            instanceId: instance.id, generation: instance.generation, uptimeSeconds: 3600, cpuCores: 3, memoryBytes: 2 * gib })) });
}

beforeEach(async () => {
    vi.resetModules(); vi.clearAllMocks();
    vi.useFakeTimers({ toFake: ['Date'] }); vi.setSystemTime(new Date('2026-10-08T16:00:00.000Z'));
    fixture = await newFixture([]); vi.stubEnv('PLUTAINER_IMAGE', image); vi.stubEnv('RCON_PASSWORD', 'preserved-password');
    const { db } = await import('../../src/lib/server/db/index.js');
    const at = new Date().toISOString(), launched = new Date(Date.now() - 600_000).toISOString();
    for (const [index, name] of names.entries()) db().prepare(`INSERT INTO servers(name,game,map,game_mode,map_slug,
        config_file,game_id,image,game_port,container_state,container_checked_at,last_players,last_map,last_seen_at,last_checked_at,last_roster)
        VALUES(?,'t5','zombie_theater','zombie','theater','dedicated_sp.cfg','t5sp',?,?,'running',?,0,'zombie_theater',?,?,'[]')`)
        .run(name, image, 28960 + index, at, at, at);
    const prepared = await prepareNativeHost(names); sourceId = prepared.host.id; libraryId = prepared.versions.get('t5')!;
    const hosts = await import('../../src/lib/server/hosts/index.js');
    targetId = hosts.createHost({ name: 'Empty destination', portMin: 29000, portMax: 29020 }).id;
    hosts.enrollHost(hosts.createEnrollment(targetId).token, { protocolVersion: 1, cpuCount: 16, cpuUsagePercent: 6.25,
        memoryTotal: 32 * gib, memoryAvailable: 30 * gib, dockerRunning: true, directoriesWritable: true,
        storageAvailable: 100 * gib, runtimeImages: [image], libraries: [{ id: libraryId, state: 'ready' }] });
    const libraries = await import('../../src/lib/server/game-files/index.js');
    libraries.updateHostLibraryCaches(targetId, [{ libraryVersionId: libraryId, status: 'ready',
        manifestSha256: libraries.getLibraryVersion(libraryId)!.manifestSha256 }]);
    const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
    const { beginDeployment } = await import('../../src/lib/server/provisioning/history.js');
    const { captureRecipe, recipeStorage } = await import('../../src/lib/server/fleet/launch-recipes.js');
    const { reserveServerKey } = await import('../../src/lib/server/fleet/server-keys.js');
    const { recordScript } = await import('../../src/lib/server/fleet/inventory.js');
    for (const name of names) {
        const instance = nativeInstance(name)!;
        reserveServerKey(name, 't5', `saved-key-${name}`);
        for (const protocol of ['udp', 'tcp'] as const) hosts.reserveHostPort(sourceId, instance.id, { port: instance.port, bindAddress: '0.0.0.0', protocol });
        const id = beginDeployment(name, 'create', 'fixture');
        const operation = hosts.queueHostOperation(sourceId, 'instance.create', { name, rconPassword: 'preserved-password' },
            { instanceId: instance.id, generation: 1, idempotencyKey: randomUUID() });
        hosts.reportHostOperation(sourceId, operation.id, { state: 'completed', result: { state: 'running', healthy: true } });
        db().prepare("UPDATE server_deployments SET state='done',submitted_at=?,finished_at=? WHERE id=?").run(launched, launched, id);
        db().prepare('UPDATE native_provisioning SET operation_id=?,verified_at=? WHERE deployment_id=?').run(operation.id, launched, id);
        const recipe = captureRecipe('t5', 'theater', 'preserved-password');
        const stored = recipeStorage({ ...recipe, image, native: { hostId: sourceId, libraryId, image } });
        db().prepare('INSERT INTO server_launch_recipes(deployment_id,recipe,recipe_summary) VALUES(?,?,?)').run(id, stored.recipe, stored.recipeSummary);
        for (const file of recipe.files) recordScript(name, file.path, file.versionId, file.sha256, file.size, 'native-create');
    }
    await heartbeat(sourceId, 75); await heartbeat(targetId, 6.25);
    const backend = await import('../../src/lib/server/backends/native.js');
    vi.spyOn(backend, 'nativeRequest').mockImplementation(async (_instance, kind) => {
        if (kind !== 'instance.inspect') throw new Error(`Unexpected host request ${kind}`);
        return { state: 'running', healthy: true, exists: true } as never;
    });
});
afterEach(async () => { vi.useRealTimers(); vi.restoreAllMocks(); await fixture.cleanup(); vi.unstubAllEnvs(); });

async function start() {
    const service = await import('../../src/lib/server/rebalancing/index.js');
    const preview = service.previewRebalance(settings());
    expect(preview.moves.length, JSON.stringify(preview.exclusions)).toBeGreaterThanOrEqual(2);
    return service.startRebalance(preview.id);
}
async function currentMove(runId: string): Promise<RebalanceMove> {
    const { getRebalanceRun } = await import('../../src/lib/server/rebalancing/index.js');
    return getRebalanceRun(runId)!.moves.find(move => move.state === 'moving')!;
}
async function provision() {
    const { processProvisioning } = await import('../../src/lib/server/provisioning/worker.js');
    await processProvisioning();
}
async function finishMove(move: RebalanceMove) {
    const { hostMigration } = await import('../../src/lib/server/provisioning/host-drain.js');
    const hosts = await import('../../src/lib/server/hosts/index.js');
    const { db } = await import('../../src/lib/server/db/index.js');
    const removal = hostMigration(move.operationId!)!.remove_operation_id;
    expect(removal).toBeTruthy();
    clockStep(); hosts.reportHostOperation(sourceId, removal!, { state: 'completed' });
    await provision(); await provision();
    const created = db().prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?').get(move.operationId) as { operation_id: string | null };
    expect(created.operation_id).toBeTruthy();
    clockStep(); hosts.reportHostOperation(targetId, created.operation_id!, { state: 'completed', result: { state: 'running', healthy: true } });
    clockStep(); await provision();
    expect(db().prepare('SELECT state FROM server_deployments WHERE id=?').get(move.operationId)).toEqual({ state: 'done' });
}

describe('manual rebalance runs', () => {
    it('previews without reserving capacity and rejects stale plans', async () => {
        const { db } = await import('../../src/lib/server/db/index.js');
        const { previewRebalance, startRebalance } = await import('../../src/lib/server/rebalancing/index.js');
        const ports = db().prepare('SELECT * FROM ec_host_ports').all();
        const commands = db().prepare('SELECT id FROM ec_host_operations').all();
        const preview = previewRebalance(settings());
        expect(preview.moves.length).toBeGreaterThanOrEqual(2);
        expect(db().prepare('SELECT * FROM ec_host_ports').all()).toEqual(ports);
        expect(db().prepare('SELECT id FROM ec_host_operations').all()).toEqual(commands);
        expect(db().prepare('SELECT * FROM rebalancing_runs').all()).toEqual([]);
        db().prepare("UPDATE servers SET desired_state='stopped' WHERE name=?").run(preview.moves[0].server);
        expect(() => startRebalance(preview.id)).toThrow('fleet changed');
    });

    it('starts idempotently, rejects concurrent runs, and queues exactly one move', async () => {
        const service = await import('../../src/lib/server/rebalancing/index.js');
        const preview = service.previewRebalance(settings());
        const competing = service.previewRebalance(settings());
        const run = service.startRebalance(preview.id);
        expect(service.startRebalance(preview.id).id).toBe(run.id);
        expect(() => service.startRebalance(competing.id)).toThrow('current rebalance');
        expect(await service.runRebalancing()).toBe(true);
        expect(await service.runRebalancing()).toBe(false);
        const moves = service.getRebalanceRun(run.id)!.moves;
        expect(moves.filter(move => move.operationId)).toHaveLength(1);
        expect(moves.filter(move => move.state === 'planned').length).toBeGreaterThan(0);
    });

    it('rejects expired previews but returns an already-started run on a late retry', async () => {
        const service = await import('../../src/lib/server/rebalancing/index.js');
        const first = service.previewRebalance(settings());
        const expired = service.previewRebalance(settings());
        const run = service.startRebalance(first.id);
        vi.setSystemTime(Date.now() + 6 * 60_000);
        expect(service.startRebalance(first.id).id).toBe(run.id);
        expect(() => service.startRebalance(expired.id)).toThrow('Preview expired');
    });

    it('stops before removal without changing placement or releasing the key, and allows a later host drain', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const stopped = service.stopRebalance(run.id);
        expect(stopped.state).toBe('stopped');
        expect(stopped.moves.every(move => move.state === 'skipped')).toBe(true);
        await provision();
        const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
        const { getServerKey } = await import('../../src/lib/server/fleet/server-keys.js');
        expect(nativeInstance(move.server)?.host_id).toBe(sourceId);
        expect(getServerKey(move.server)).toBe(`saved-key-${move.server}`);
        const { db } = await import('../../src/lib/server/db/index.js');
        expect(db().prepare("SELECT id FROM ec_host_operations WHERE kind='instance.remove'").all()).toEqual([]);
        const { updateHost } = await import('../../src/lib/server/hosts/index.js');
        updateHost(sourceId, { mode: 'draining' });
        const { reconcileHostDrains } = await import('../../src/lib/server/provisioning/host-drain.js');
        reconcileHostDrains();
        expect(db().prepare("SELECT id FROM server_deployments WHERE server=? AND triggered_by='host-drain'").get(move.server)).toBeTruthy();
    });

    it('finishes a committed move after Stop while withdrawing the remaining moves', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing(); await provision();
        const move = await currentMove(run.id);
        expect(service.stopRebalance(run.id).state).toBe('stopping');
        const { cancelProvisioningJob } = await import('../../src/lib/server/provisioning/cancellation.js');
        expect(cancelProvisioningJob(move.operationId!)).toBe(false);
        await finishMove(move); await service.runRebalancing();
        const finished = service.getRebalanceRun(run.id)!;
        expect(finished.state).toBe('stopped');
        expect(finished.moves.filter(move => move.state === 'done')).toHaveLength(1);
        expect(finished.moves.filter(move => move.state === 'skipped')).toHaveLength(run.moves.length - 1);
        const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
        const { getHost } = await import('../../src/lib/server/hosts/index.js');
        expect(nativeInstance(move.server)).toMatchObject({ id: move.instanceId, generation: 2, host_id: targetId });
        expect(getHost(sourceId)?.mode).toBe('active');
    });

    it('treats cancellation from Provisioning as stopping the rest of the run', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const { cancelProvisioningJob } = await import('../../src/lib/server/provisioning/cancellation.js');
        expect(cancelProvisioningJob(move.operationId!)).toBe(true);
        expect(service.getRebalanceRun(run.id)?.state).toBe('stopping');
        await service.runRebalancing();
        expect(service.getRebalanceRun(run.id)?.state).toBe('stopped');
        expect(service.getRebalanceRun(run.id)?.moves.every(move => move.state === 'skipped')).toBe(true);
    });

    it('waits for both hosts to report fresh measurements before admitting the next move', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing(); await provision();
        await finishMove(await currentMove(run.id));
        expect(await service.runRebalancing()).toBe(false);
        expect(service.getRebalanceRun(run.id)?.message).toContain('fresh source and destination');
        expect(service.getRebalanceRun(run.id)?.moves.filter(move => move.operationId)).toHaveLength(1);
        clockStep(); await heartbeat(sourceId, 56.25);
        expect(await service.runRebalancing()).toBe(false);
        await heartbeat(targetId, 25);
        expect(await service.runRebalancing()).toBe(true);
        expect(service.getRebalanceRun(run.id)?.moves.filter(move => move.operationId)).toHaveLength(2);
    });

    it('skips an uncommitted move that becomes occupied and continues other approved work', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const { db } = await import('../../src/lib/server/db/index.js');
        db().prepare('UPDATE servers SET last_players=2 WHERE name=?').run(move.server);
        await service.runRebalancing();
        const fresh = service.getRebalanceRun(run.id)!;
        expect(fresh.state).toBe('running');
        expect(fresh.moves.find(item => item.id === move.id)?.state).toBe('skipped');
        expect(fresh.moves.filter(item => item.state === 'moving')).toHaveLength(1);
    });

    it('halts remaining moves after a provisioning failure', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const { db } = await import('../../src/lib/server/db/index.js');
        db().prepare("UPDATE server_deployments SET state='failed',error='Recovery needed' WHERE id=?").run(move.operationId);
        expect(await service.runRebalancing()).toBe(false);
        const failed = service.getRebalanceRun(run.id)!;
        expect(failed.state).toBe('failed');
        expect(failed.moves.find(item => item.id === move.id)).toMatchObject({ state: 'failed', message: 'Recovery needed' });
        expect(failed.moves.filter(item => item.state === 'skipped')).toHaveLength(run.moves.length - 1);
    });

    it('keeps the source running until a global startup slot is available', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const { db } = await import('../../src/lib/server/db/index.js');
        const { beginDeployment } = await import('../../src/lib/server/provisioning/history.js');
        const occupied = beginDeployment('another-server', 'create', 'fixture');
        await provision();
        const { hostMigration } = await import('../../src/lib/server/provisioning/host-drain.js');
        expect(hostMigration(move.operationId!)?.remove_operation_id).toBeNull();
        db().prepare("UPDATE server_deployments SET state='done' WHERE id=?").run(occupied);
        await provision();
        expect(hostMigration(move.operationId!)?.remove_operation_id).toEqual(expect.any(String));
    });

    it('retains its global startup reservation between source removal and destination creation', async () => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing(); await provision();
        const move = await currentMove(run.id);
        const { hostMigration } = await import('../../src/lib/server/provisioning/migrations.js');
        const { reportHostOperation } = await import('../../src/lib/server/hosts/index.js');
        reportHostOperation(sourceId, hostMigration(move.operationId!)!.remove_operation_id!, { state: 'completed' });
        await provision();
        const { db } = await import('../../src/lib/server/db/index.js');
        expect(db().prepare('SELECT state FROM server_deployments WHERE id=?').get(move.operationId)).toEqual({ state: 'waiting' });
        const { queueNativeCreation } = await import('../../src/lib/server/provisioning/native.js');
        const extra = queueNativeCreation({ game: 't5', gameMode: 'zombie', mapSlug: 'theater',
            hostId: targetId, libraryId, image, plutoKey: 'another-key' });
        expect(extra.ok).toBe(true);
        if (!extra.ok) throw new Error(extra.message);
        await provision();
        expect(db().prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?').get(extra.deploymentId)).toBeUndefined();
        expect(db().prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?').get(move.operationId))
            .toEqual({ operation_id: expect.any(String) });
    });

    it.each([{ state: 'running', healthy: false }, { state: 'stopped', healthy: false }])('holds removal when fresh inspection returns $state with healthy=$healthy', async health => {
        const run = await start();
        const service = await import('../../src/lib/server/rebalancing/index.js');
        await service.runRebalancing();
        const move = await currentMove(run.id);
        const { nativeRequest } = await import('../../src/lib/server/backends/native.js');
        vi.mocked(nativeRequest).mockResolvedValueOnce(health);
        await provision();
        const { hostMigration } = await import('../../src/lib/server/provisioning/host-drain.js');
        expect(hostMigration(move.operationId!)?.remove_operation_id).toBeNull();
        const { db } = await import('../../src/lib/server/db/index.js');
        expect(db().prepare('SELECT stage FROM server_deployments WHERE id=?').get(move.operationId))
            .toEqual({ stage: expect.stringContaining('no longer healthy') });
    });
});