import { createHash } from 'node:crypto';
import { createServer, type IncomingMessage, type ServerResponse } from 'node:http';
import { vi } from 'vitest';
import { LambdaMicrovmSandboxProvider, type LambdaMicrovmClientLike } from '../../src/sandbox/lambda-microvm.js';

type AuthTokenRequest = { microvmId: string; port: number; expirationInMinutes: number };

describe('LambdaMicrovmSandboxProvider', () => {
  it('creates handles that use the MicroVM endpoint as a bridge transport', async () => {
    const bridgeRequests: IncomingMessage[] = [];
    const server = createServer((request, response) => {
      bridgeRequests.push(request);
      handleBridgeRequest(request, response);
    });
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    const address = server.address();
    if (!address || typeof address === 'string') throw new Error('Expected local server address');

    let state = 'RUNNING';
    const runCalls: unknown[] = [];
    const authTokenPorts: number[] = [];
    const client: LambdaMicrovmClientLike = {
      async checkImage() {},
      async run(input) {
        runCalls.push(input);
        return {
          microvmId: 'mvm-1',
          endpoint: `http://127.0.0.1:${address.port}`,
          state,
          imageArn: 'arn:aws:lambda:us-east-2:123456789012:microvm-image:deputies',
          imageVersion: '1.0',
        };
      },
      async get(microvmId) {
        return { microvmId, endpoint: `http://127.0.0.1:${address.port}`, state };
      },
      async createAuthToken(input) {
        authTokenPorts.push(input.port);
        return `proxy-token-${input.port}`;
      },
      async suspend() {
        state = 'SUSPENDED';
      },
      async resume() {
        state = 'RUNNING';
      },
      async terminate() {
        state = 'TERMINATED';
      },
    };
    const provider = new LambdaMicrovmSandboxProvider({
      client,
      imageIdentifier: 'deputies-sandbox',
      imageVersion: '1.0',
      executionRoleArn: 'arn:aws:iam::123456789012:role/deputies-microvm-runtime',
      ingressNetworkConnectors: ['all-ingress'],
      egressNetworkConnectors: ['internet-egress'],
      idleTimeoutMs: 900_000,
      suspendedDurationMs: 3600_000,
      maximumDurationSeconds: 28_800,
      authTokenTtlMinutes: 30,
      workspacePath: '/workspace/custom',
      bridgePort: 3584,
      logGroup: '/aws/lambda/microvms/deputies',
    });

    try {
      const handle = await provider.create({ sessionId: 'session-1', metadata: { owner: 'test' } });

      expect(runCalls[0]).toMatchObject({
        imageIdentifier: 'deputies-sandbox',
        imageVersion: '1.0',
        executionRoleArn: 'arn:aws:iam::123456789012:role/deputies-microvm-runtime',
        ingressNetworkConnectors: ['all-ingress'],
        egressNetworkConnectors: ['internet-egress'],
        maximumDurationSeconds: 28_800,
        logGroup: '/aws/lambda/microvms/deputies',
      });
      expect(JSON.parse((runCalls[0] as { runHookPayload: string }).runHookPayload)).toMatchObject({
        sessionId: 'session-1',
        workspacePath: '/workspace/custom',
        bridgePort: 3584,
      });
      expect(handle).toMatchObject({
        provider: 'lambda-microvm',
        providerSandboxId: 'mvm-1',
        sessionId: 'session-1',
        workspacePath: '/workspace/custom',
        metadata: { owner: 'test', microvmId: 'mvm-1', imageVersion: '1.0' },
      });

      await expect(handle.exec({ command: 'echo ok', cwd: '/workspace/custom' })).resolves.toMatchObject({
        exitCode: 0,
        stdout: 'ran: echo ok',
        stderr: '',
      });
      await expect(handle.fs?.readFileBuffer('image.bin')).resolves.toEqual(new Uint8Array([0, 1, 2, 127, 128, 255]));
      await expect(handle.fs?.readFileBuffer('corrupt.bin')).rejects.toThrow(
        'Lambda MicroVM bridge read checksum mismatch for corrupt.bin',
      );
      const execRequest = bridgeRequests.find((request) => request.url === '/exec');
      expect(execRequest?.headers.authorization).toMatch(/^Bearer /);
      expect(execRequest?.headers['x-aws-proxy-auth']).toBe('proxy-token-3584');
      expect(execRequest?.headers['x-aws-proxy-port']).toBe('3584');

      const endpoint = await provider.getServiceEndpoint({
        providerSandboxId: 'mvm-1',
        sessionId: 'session-1',
        port: 3000,
      });
      expect(endpoint).toMatchObject({
        port: 3000,
        targetUrl: `http://127.0.0.1:${address.port}/preview/3000`,
        targetHeaders: {
          authorization: expect.stringMatching(/^Bearer /),
          'x-aws-proxy-auth': 'proxy-token-3584',
          'x-aws-proxy-port': '3584',
        },
        preserveTargetHost: true,
        forwardPreviewHost: true,
      });
      await expect(
        provider.getServiceEndpoint({ providerSandboxId: 'mvm-1', sessionId: 'session-1', port: 3584 }),
      ).resolves.toBeNull();
      await expect(
        provider.getServiceEndpoint({ providerSandboxId: 'mvm-1', sessionId: 'session-1', port: 9000 }),
      ).resolves.toBeNull();

      await provider.stop({ providerSandboxId: 'mvm-1', sessionId: 'session-1' });
      await expect(provider.health({ providerSandboxId: 'mvm-1', sessionId: 'session-1' })).resolves.toMatchObject({
        status: 'stopped',
        message: 'Lambda MicroVM is suspended with auto-resume',
      });
      await provider.start({ providerSandboxId: 'mvm-1', sessionId: 'session-1' });
      await expect(provider.connect({ providerSandboxId: 'mvm-1', sessionId: 'session-1' })).resolves.toMatchObject({
        providerSandboxId: 'mvm-1',
      });
      await expect(provider.health({ providerSandboxId: 'mvm-1', sessionId: 'session-1' })).resolves.toMatchObject({
        status: 'ready',
      });
      await provider.destroy({ providerSandboxId: 'mvm-1', sessionId: 'session-1' });
      expect(authTokenPorts).toEqual(expect.arrayContaining([3584]));
    } finally {
      await new Promise<void>((resolve, reject) => server.close((error) => (error ? reject(error) : resolve())));
    }
  });

  it('caches bridge auth tokens across requests for the same MicroVM port', async () => {
    const fetchCalls: Array<{ url: string; headers: Headers }> = [];
    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockImplementation(async (input, init) => {
      const url = fetchUrl(input);
      const headers = new Headers(init?.headers);
      fetchCalls.push({ url, headers });
      if (url.endsWith('/health')) return fetchJson({ status: 'ok' });
      if (url.endsWith('/exec')) return execJson();
      if (url.includes('/fs/exists?')) return fetchJson({ exists: true });
      return fetchJson({ error: 'not_found' }, 404);
    });
    const { client, authTokenRequests } = mockClient();
    const provider = new LambdaMicrovmSandboxProvider({ client, imageIdentifier: 'deputies-sandbox' });

    try {
      const handle = await provider.create({ sessionId: 'session-1' });
      if (!handle.fs) throw new Error('Expected Lambda MicroVM handle to expose filesystem');

      await expect(handle.exec({ command: 'echo ok' })).resolves.toMatchObject({ exitCode: 0 });
      await expect(handle.fs.exists('/tmp/ready')).resolves.toBe(true);
      await expect(
        provider.getServiceEndpoint({ providerSandboxId: 'mvm-1', sessionId: 'session-1', port: 3000 }),
      ).resolves.toMatchObject({ targetHeaders: { 'x-aws-proxy-auth': 'proxy-token-1' } });

      expect(authTokenRequests).toEqual([{ microvmId: 'mvm-1', port: 3584, expirationInMinutes: 30 }]);
      expect(fetchCalls.map((call) => call.headers.get('x-aws-proxy-auth'))).toEqual([
        'proxy-token-1',
        'proxy-token-1',
        'proxy-token-1',
      ]);
    } finally {
      fetchSpy.mockRestore();
    }
  });

  it('refreshes bridge auth tokens before cached expiry', async () => {
    vi.useFakeTimers({ toFake: ['Date'] });
    vi.setSystemTime(new Date(0));
    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockImplementation(async (input) => {
      const url = fetchUrl(input);
      if (url.endsWith('/health')) return fetchJson({ status: 'ok' });
      if (url.endsWith('/exec')) return execJson();
      return fetchJson({ error: 'not_found' }, 404);
    });
    const { client, authTokenRequests } = mockClient();
    const provider = new LambdaMicrovmSandboxProvider({
      client,
      imageIdentifier: 'deputies-sandbox',
      authTokenTtlMinutes: 1,
    });

    try {
      const handle = await provider.create({ sessionId: 'session-1' });
      await handle.exec({ command: 'echo first' });
      vi.setSystemTime(new Date(47_000));
      await handle.exec({ command: 'echo cached' });
      expect(authTokenRequests).toHaveLength(1);

      vi.setSystemTime(new Date(49_000));
      await handle.exec({ command: 'echo refreshed' });
      expect(authTokenRequests).toHaveLength(2);
      expect(authTokenRequests.map((request) => request.expirationInMinutes)).toEqual([1, 1]);
    } finally {
      fetchSpy.mockRestore();
      vi.useRealTimers();
    }
  });

  it('invalidates cached bridge auth tokens and retries once on auth rejection', async () => {
    let execRequests = 0;
    const execAuthHeaders: string[] = [];
    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockImplementation(async (input, init) => {
      const url = fetchUrl(input);
      const headers = new Headers(init?.headers);
      if (url.endsWith('/health')) return fetchJson({ status: 'ok' });
      if (url.endsWith('/exec')) {
        execRequests += 1;
        execAuthHeaders.push(headers.get('x-aws-proxy-auth') ?? '');
        if (execRequests === 1) return fetchJson({ error: 'expired' }, 401);
        return execJson();
      }
      return fetchJson({ error: 'not_found' }, 404);
    });
    const { client, authTokenRequests } = mockClient();
    const provider = new LambdaMicrovmSandboxProvider({ client, imageIdentifier: 'deputies-sandbox' });

    try {
      const handle = await provider.create({ sessionId: 'session-1' });

      await expect(handle.exec({ command: 'echo retry' })).resolves.toMatchObject({ exitCode: 0 });
      expect(execRequests).toBe(2);
      expect(execAuthHeaders).toEqual(['proxy-token-1', 'proxy-token-2']);
      expect(authTokenRequests).toHaveLength(2);
    } finally {
      fetchSpy.mockRestore();
    }
  });
});

function mockClient(): {
  client: LambdaMicrovmClientLike;
  authTokenRequests: AuthTokenRequest[];
} {
  const authTokenRequests: AuthTokenRequest[] = [];
  return {
    authTokenRequests,
    client: {
      async checkImage() {},
      async run() {
        return { microvmId: 'mvm-1', endpoint: 'http://lambda-microvm.localhost', state: 'RUNNING' };
      },
      async get(microvmId) {
        return { microvmId, endpoint: 'http://lambda-microvm.localhost', state: 'RUNNING' };
      },
      async createAuthToken(input) {
        authTokenRequests.push(input);
        return `proxy-token-${authTokenRequests.length}`;
      },
      async suspend() {},
      async resume() {},
      async terminate() {},
    },
  };
}

function fetchUrl(input: Parameters<typeof fetch>[0]): string {
  return input instanceof Request ? input.url : String(input);
}

function fetchJson(body: unknown, status = 200): Response {
  return new Response(JSON.stringify(body), { status, headers: { 'content-type': 'application/json; charset=utf-8' } });
}

function execJson(): Response {
  return fetchJson({
    exitCode: 0,
    stdout: 'ok',
    stderr: '',
    startedAt: new Date(0).toISOString(),
    completedAt: new Date(1).toISOString(),
  });
}

function handleBridgeRequest(request: IncomingMessage, response: ServerResponse): void {
  void (async () => {
    if (request.url === '/health') return json(response, 200, { status: 'ok' });
    if (request.url === '/exec') {
      const body = JSON.parse(await readBody(request)) as { command: string };
      return json(response, 200, {
        exitCode: 0,
        stdout: `ran: ${body.command}`,
        stderr: '',
        startedAt: new Date(0).toISOString(),
        completedAt: new Date(1).toISOString(),
      });
    }
    if (request.url?.startsWith('/fs/read?')) {
      const body = Buffer.from([0, 1, 2, 127, 128, 255]);
      const checksumBody = request.url.includes('corrupt.bin') ? Buffer.from('changed') : body;
      response.writeHead(200, {
        'content-length': String(body.byteLength),
        'content-type': 'application/octet-stream',
        'x-deputies-sha256': checksumSha256(checksumBody),
      });
      response.end(body);
      return;
    }
    json(response, 404, { error: 'not_found' });
  })().catch((error: unknown) => {
    response.writeHead(500, { 'content-type': 'text/plain; charset=utf-8' });
    response.end(error instanceof Error ? error.message : 'Unknown error');
  });
}

function readBody(request: IncomingMessage): Promise<string> {
  return new Promise((resolve, reject) => {
    const chunks: Buffer[] = [];
    request.on('data', (chunk: Buffer) => chunks.push(chunk));
    request.on('error', reject);
    request.on('end', () => resolve(Buffer.concat(chunks).toString('utf-8')));
  });
}

function json(response: ServerResponse, status: number, body: unknown): void {
  response.writeHead(status, { 'content-type': 'application/json; charset=utf-8' });
  response.end(JSON.stringify(body));
}

function checksumSha256(body: Uint8Array): string {
  return createHash('sha256').update(body).digest('hex');
}
