import { expect, tap } from '@push.rocks/tapbundle';
import * as dgram from 'dgram';
import * as net from 'net';
import * as crypto from 'crypto';
import { RemoteIngressHub, RemoteIngressEdge } from '../ts/index.js';

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

async function findFreePorts(count: number): Promise<number[]> {
  const servers: net.Server[] = [];
  const ports: number[] = [];
  for (let i = 0; i < count; i++) {
    const server = net.createServer();
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    ports.push((server.address() as net.AddressInfo).port);
    servers.push(server);
  }
  await Promise.all(servers.map((s) => new Promise<void>((resolve) => s.close(() => resolve()))));
  return ports;
}

/**
 * Start a UDP echo server that:
 * 1. Receives the first datagram (PROXY v2 header — 28 bytes) and discards it
 * 2. Echoes all subsequent datagrams back to the sender
 */
function startUdpEchoServer(port: number, host: string): Promise<dgram.Socket> {
  return new Promise((resolve, reject) => {
    const server = dgram.createSocket('udp4');
    // Track which source endpoints have sent their PROXY v2 header.
    // The hub sends a 28-byte PROXY v2 header as the first datagram per session.
    const seenSources = new Set<string>();

    server.on('message', (msg, rinfo) => {
      const sourceKey = `${rinfo.address}:${rinfo.port}`;
      if (!seenSources.has(sourceKey)) {
        seenSources.add(sourceKey);
        // First datagram from this source is the PROXY v2 header — skip it
        return;
      }
      // Echo back
      server.send(msg, rinfo.port, rinfo.address);
    });

    server.on('error', reject);
    server.bind(port, host, () => resolve(server));
  });
}

/**
 * Send a UDP datagram through the tunnel and wait for the echo response.
 */
function udpSendAndReceive(
  port: number,
  data: Buffer,
  timeoutMs = 10000,
): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    const client = dgram.createSocket('udp4');
    let settled = false;

    const timer = setTimeout(() => {
      if (!settled) {
        settled = true;
        client.close();
        reject(new Error(`UDP timeout after ${timeoutMs}ms`));
      }
    }, timeoutMs);

    client.on('message', (msg) => {
      if (!settled) {
        settled = true;
        clearTimeout(timer);
        client.close();
        resolve(msg);
      }
    });

    client.on('error', (err) => {
      if (!settled) {
        settled = true;
        clearTimeout(timer);
        client.close();
        reject(err);
      }
    });

    client.send(data, port, '127.0.0.1');
  });
}

// ---------------------------------------------------------------------------
// Test state
// ---------------------------------------------------------------------------

let hub: RemoteIngressHub;
let edge: RemoteIngressEdge;
let echoServer: dgram.Socket;
let hubPort: number;
let edgeUdpPort: number;

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

tap.test('UDP/TLS setup: start UDP echo server and TCP+TLS tunnel with UDP ports', async () => {
  [hubPort, edgeUdpPort] = await findFreePorts(2);

  // Start UDP echo server on upstream (127.0.0.2)
  echoServer = await startUdpEchoServer(edgeUdpPort, '127.0.0.2');

  hub = new RemoteIngressHub();
  edge = new RemoteIngressEdge();

  await hub.start({ tunnelPort: hubPort, targetHost: '127.0.0.2' });
  await hub.updateAllowedEdges([
    { id: 'test-edge', secret: 'test-secret', listenPorts: [], listenPortsUdp: [edgeUdpPort] },
  ]);

  const connectedPromise = new Promise<void>((resolve, reject) => {
    const timeout = setTimeout(() => reject(new Error('Edge did not connect within 10s')), 10000);
    edge.once('tunnelConnected', () => {
      clearTimeout(timeout);
      resolve();
    });
  });

  await edge.start({
    hubHost: '127.0.0.1',
    hubPort,
    edgeId: 'test-edge',
    secret: 'test-secret',
    bindAddress: '127.0.0.1',
  });

  await connectedPromise;
  // Wait for UDP listener to bind
  await new Promise((resolve) => setTimeout(resolve, 500));

  const status = await edge.getStatus();
  expect(status.connected).toBeTrue();
});

tap.test('UDP/TLS: single UDP datagram echo — 64 bytes', async () => {
  const data = crypto.randomBytes(64);
  const received = await udpSendAndReceive(edgeUdpPort, data, 5000);
  expect(received.length).toEqual(64);
  expect(Buffer.compare(received, data)).toEqual(0);
});

tap.test('UDP/TLS: single UDP datagram echo — 1KB', async () => {
  const data = crypto.randomBytes(1024);
  const received = await udpSendAndReceive(edgeUdpPort, data, 5000);
  expect(received.length).toEqual(1024);
  expect(Buffer.compare(received, data)).toEqual(0);
});

tap.test('UDP/TLS: 10 sequential UDP datagrams', async () => {
  for (let i = 0; i < 10; i++) {
    const data = crypto.randomBytes(128);
    const received = await udpSendAndReceive(edgeUdpPort, data, 5000);
    expect(received.length).toEqual(128);
    expect(Buffer.compare(received, data)).toEqual(0);
  }
});

tap.test('UDP/TLS: 10 concurrent UDP datagrams from different source ports', async () => {
  const promises = Array.from({ length: 10 }, () => {
    const data = crypto.randomBytes(256);
    return udpSendAndReceive(edgeUdpPort, data, 5000).then((received) => ({
      sizeOk: received.length === 256,
      dataOk: Buffer.compare(received, data) === 0,
    }));
  });

  const results = await Promise.all(promises);
  const failures = results.filter((r) => !r.sizeOk || !r.dataOk);
  expect(failures.length).toEqual(0);
});

tap.test('UDP/TLS: tunnel still connected after UDP tests', async () => {
  const status = await edge.getStatus();
  expect(status.connected).toBeTrue();
});

tap.test('UDP/TLS teardown: stop tunnel and UDP echo server', async () => {
  await edge.stop();
  await hub.stop();
  await new Promise<void>((resolve) => echoServer.close(() => resolve()));
});

// ---------------------------------------------------------------------------
// QUIC transport UDP tests
// ---------------------------------------------------------------------------

let quicHub: RemoteIngressHub;
let quicEdge: RemoteIngressEdge;
let quicEchoServer: dgram.Socket;
let quicHubPort: number;
let quicEdgeUdpPort: number;

tap.test('UDP/QUIC setup: start UDP echo server and QUIC tunnel with UDP ports', async () => {
  [quicHubPort, quicEdgeUdpPort] = await findFreePorts(2);

  quicEchoServer = await startUdpEchoServer(quicEdgeUdpPort, '127.0.0.2');

  quicHub = new RemoteIngressHub();
  quicEdge = new RemoteIngressEdge();

  await quicHub.start({ tunnelPort: quicHubPort, targetHost: '127.0.0.2' });
  await quicHub.updateAllowedEdges([
    { id: 'test-edge', secret: 'test-secret', listenPorts: [], listenPortsUdp: [quicEdgeUdpPort] },
  ]);

  const connectedPromise = new Promise<void>((resolve, reject) => {
    const timeout = setTimeout(() => reject(new Error('QUIC edge did not connect within 10s')), 10000);
    quicEdge.once('tunnelConnected', () => {
      clearTimeout(timeout);
      resolve();
    });
  });

  await quicEdge.start({
    hubHost: '127.0.0.1',
    hubPort: quicHubPort,
    edgeId: 'test-edge',
    secret: 'test-secret',
    bindAddress: '127.0.0.1',
    transportMode: 'quic',
  });

  await connectedPromise;
  await new Promise((resolve) => setTimeout(resolve, 500));

  const status = await quicEdge.getStatus();
  expect(status.connected).toBeTrue();
});

tap.test('UDP/QUIC: single UDP datagram echo — 64 bytes', async () => {
  const data = crypto.randomBytes(64);
  const received = await udpSendAndReceive(quicEdgeUdpPort, data, 5000);
  expect(received.length).toEqual(64);
  expect(Buffer.compare(received, data)).toEqual(0);
});

tap.test('UDP/QUIC: single UDP datagram echo — 1KB', async () => {
  const data = crypto.randomBytes(1024);
  const received = await udpSendAndReceive(quicEdgeUdpPort, data, 5000);
  expect(received.length).toEqual(1024);
  expect(Buffer.compare(received, data)).toEqual(0);
});

tap.test('UDP/QUIC: 10 sequential UDP datagrams', async () => {
  for (let i = 0; i < 10; i++) {
    const data = crypto.randomBytes(128);
    const received = await udpSendAndReceive(quicEdgeUdpPort, data, 5000);
    expect(received.length).toEqual(128);
    expect(Buffer.compare(received, data)).toEqual(0);
  }
});

tap.test('UDP/QUIC: 10 concurrent UDP datagrams', async () => {
  const promises = Array.from({ length: 10 }, () => {
    const data = crypto.randomBytes(256);
    return udpSendAndReceive(quicEdgeUdpPort, data, 5000).then((received) => ({
      sizeOk: received.length === 256,
      dataOk: Buffer.compare(received, data) === 0,
    }));
  });

  const results = await Promise.all(promises);
  const failures = results.filter((r) => !r.sizeOk || !r.dataOk);
  expect(failures.length).toEqual(0);
});

tap.test('UDP/QUIC teardown: stop QUIC tunnel and UDP echo server', async () => {
  await quicEdge.stop();
  await quicHub.stop();
  await new Promise<void>((resolve) => quicEchoServer.close(() => resolve()));
});

export default tap.start();