import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { newFixture, type Fixture } from '../helpers/fixture.js';
import { prepareNativeHost } from '../helpers/native-host.js';

vi.mock('../../src/lib/server/pipeline/worker.js', () => ({
    kick: vi.fn(async () => {}),
    kickProvisioning: vi.fn(async () => {})
}));

let fixture: Fixture;
const server = 't5_zombie_theater_1';
beforeEach(async () => {
    vi.resetModules();
    fixture = await newFixture([]);
});
afterEach(async () => {
    await fixture.cleanup();
});

async function recovering(step: 'restart' | 'rebuild') {
    const { host } = await prepareNativeHost([server]);
    const hosts = await import('../../src/lib/server/hosts/index.js');
    const { db } = await import('../../src/lib/server/db/index.js');
    const { nativeInstance } = await import('../../src/lib/server/backends/native.js');
    const { beginDeployment, getDeployment } =
        await import('../../src/lib/server/provisioning/history.js');
    const { cancelCreation, canCancelCreation } =
        await import('../../src/lib/server/provisioning/cancel-creation.js');
    const { processNativeProvisioning, reconcileNativeProvisioningReceipt } =
        await import('../../src/lib/server/provisioning/native.js');
    hosts.heartbeat(host.id, {
        ...hosts.getHost(host.id)!.inventory,
        capabilities: ['instance.cancel-create', 'instance.unhealthy-recovery']
    });
    db().prepare('INSERT INTO servers(name) VALUES(?)').run(server);
    const instance = nativeInstance(server)!;
    hosts.reserveHostPort(host.id, instance.id, { port: instance.port, protocol: 'udp' });
    hosts.reserveHostPort(host.id, instance.id, { port: instance.port, protocol: 'tcp' });
    const deploymentId = beginDeployment(server, 'create');
    const original = hosts.queueHostOperation(
        host.id,
        'instance.create',
        {},
        { instanceId: instance.id, generation: 1 }
    );
    hosts.reportHostOperation(host.id, original.id, {
        state: 'completed',
        result: { state: 'unhealthy', healthy: false }
    });
    const recovery = hosts.queueHostOperation(
        host.id,
        `instance.${step}`,
        { unhealthyRecovery: true, ...(step === 'rebuild' ? { nextGeneration: 2 } : {}) },
        { instanceId: instance.id, generation: 1 }
    );
    db()
        .prepare(
            'UPDATE native_provisioning SET operation_id=?,recovery_step=? WHERE deployment_id=?'
        )
        .run(recovery.id, step, deploymentId);
    db()
        .prepare("UPDATE server_deployments SET state='awaiting-health' WHERE id=?")
        .run(deploymentId);
    const commandFor = (id: string) => {
        const row = db()
            .prepare('SELECT operation_id FROM native_provisioning WHERE deployment_id=?')
            .get(id) as { operation_id: string };
        return hosts.getHostOperation(row.operation_id)!;
    };
    return {
        host,
        instance,
        deploymentId,
        original,
        recovery,
        ...hosts,
        db,
        nativeInstance,
        cancelCreation,
        canCancelCreation,
        getDeployment,
        reconcileNativeProvisioningReceipt,
        commandFor,
        advance: () => processNativeProvisioning(() => null)
    };
}

describe('cancelling creation during unhealthy recovery', () => {
    it('keeps the original creation identity and suppresses recovery restart replay', async () => {
        const api = await recovering('restart');
        expect(api.pendingHostOperations(api.host.id).map((operation) => operation.id)).toContain(
            api.recovery.id
        );
        const cleanup = api.cancelCreation(api.deploymentId)!;
        const removal = api.commandFor(cleanup);
        expect(removal).toMatchObject({
            kind: 'instance.cancel-create',
            generation: 1,
            payload: { createOperationId: api.original.id }
        });
        expect(removal.payload).toEqual({ createOperationId: api.original.id });
        expect(api.pendingCreationCancellations(api.host.id)).toEqual([api.original.id]);
        expect(
            api
                .pendingHostOperations(api.host.id)
                .filter((operation) => operation.kind.startsWith('instance.'))
                .map((operation) => operation.id)
        ).toEqual([removal.id]);
        api.reportHostOperation(api.host.id, removal.id, {
            state: 'completed',
            result: { exists: false }
        });
        await api.advance();
        expect(api.getDeployment(cleanup)?.deployment.state).toBe('done');
        expect(api.getHostOperation(api.recovery.id)?.state).toBe('failed');
        expect(api.db().prepare('SELECT * FROM ec_host_ports').all()).toEqual([]);
    });

    it.each(['pending', 'running', 'unknown', 'completed'] as const)(
        'cleans up both rebuild generations when recovery is %s',
        async (state) => {
            const api = await recovering('rebuild');
            if (state !== 'pending') {
                api.pendingHostOperations(api.host.id);
                api.reportHostOperation(api.host.id, api.recovery.id, {
                    state,
                    result: { state: 'unhealthy', healthy: false }
                });
            }
            if (state === 'completed') api.reconcileNativeProvisioningReceipt(api.deploymentId);
            const registeredGeneration = state === 'completed' ? 2 : 1;
            expect(api.nativeInstance(server)?.generation).toBe(registeredGeneration);
            expect(api.canCancelCreation(api.deploymentId)).toBe(true);
            const cleanup = api.cancelCreation(api.deploymentId)!;
            const removal = api.commandFor(cleanup);
            expect(removal).toMatchObject({
                kind: 'instance.cancel-create',
                generation: 1,
                payload: {
                    createOperationId: api.original.id,
                    recoveryRebuildOperationId: api.recovery.id,
                    recoveryGeneration: 2
                }
            });
            expect(api.pendingCreationCancellations(api.host.id)).toEqual([
                api.original.id,
                api.recovery.id
            ]);
            expect(
                api
                    .db()
                    .prepare('SELECT generation FROM native_provisioning WHERE deployment_id=?')
                    .get(cleanup)
            ).toEqual({ generation: registeredGeneration });
            expect(
                api
                    .pendingHostOperations(api.host.id)
                    .filter((operation) => operation.kind.startsWith('instance.'))
                    .map((operation) => operation.id)
            ).toEqual([removal.id]);
            // A delivered rebuild may report success after cancellation, while the
            // cleanup command already owns deletion of both possible containers.
            api.reportHostOperation(api.host.id, api.recovery.id, {
                state: 'completed',
                result: { state: 'unhealthy', healthy: false }
            });
            api.reconcileNativeProvisioningReceipt(api.deploymentId);
            expect(api.nativeInstance(server)?.generation).toBe(registeredGeneration);
            expect(api.db().prepare('SELECT * FROM ec_host_ports').all()).toHaveLength(2);
            api.reportHostOperation(api.host.id, removal.id, {
                state: 'completed',
                result: { exists: false }
            });
            await api.advance();
            expect(api.getDeployment(cleanup)?.deployment.state).toBe('done');
            expect(api.getDeployment(api.deploymentId)?.deployment).toMatchObject({
                state: 'cancelled',
                server_deleted_at: expect.any(String)
            });
            expect(api.nativeInstance(server)).toBeNull();
            expect(api.db().prepare('SELECT * FROM ec_host_ports').all()).toEqual([]);
            expect(
                api.db().prepare('SELECT name FROM servers WHERE name=?').get(server)
            ).toBeUndefined();
            expect(api.pendingCreationCancellations(api.host.id)).toEqual([]);
            expect(api.cancelCreation(api.deploymentId)).toBe(cleanup);
        }
    );

    it('does not adopt a changed generation without a completed recovery receipt', async () => {
        const api = await recovering('rebuild');
        api.db()
            .prepare('UPDATE native_instances SET generation=2 WHERE id=?')
            .run(api.instance.id);
        expect(api.canCancelCreation(api.deploymentId)).toBe(false);
        expect(api.cancelCreation(api.deploymentId)).toBeNull();
        api.reportHostOperation(api.host.id, api.recovery.id, { state: 'completed' });
        api.db()
            .prepare('UPDATE native_instances SET generation=3 WHERE id=?')
            .run(api.instance.id);
        expect(api.cancelCreation(api.deploymentId)).toBeNull();
    });

    it('holds extended cleanup until the host supports recovery cancellation', async () => {
        const api = await recovering('rebuild');
        const cleanup = api.cancelCreation(api.deploymentId)!;
        api.heartbeat(api.host.id, {
            ...api.getHost(api.host.id)!.inventory,
            capabilities: ['instance.cancel-create']
        });
        expect(api.pendingHostOperations(api.host.id)).toEqual([]);
        await api.advance();
        expect(api.getDeployment(cleanup)?.deployment.stage).toContain('Update the host agent');
        expect(api.db().prepare('SELECT * FROM ec_host_ports').all()).toHaveLength(2);
        api.heartbeat(api.host.id, {
            ...api.getHost(api.host.id)!.inventory,
            capabilities: ['instance.cancel-create', 'instance.unhealthy-recovery']
        });
        expect(
            api
                .pendingHostOperations(api.host.id)
                .filter((operation) => operation.kind.startsWith('instance.'))
                .map((operation) => operation.id)
        ).toEqual([api.commandFor(cleanup).id]);
    });
});