import { expect, tap } from '@push.rocks/tapbundle';
import * as net from 'net';
import { PortProxy } from '../ts/classes.portproxy.js';

let testServer: net.Server;
let portProxy: PortProxy;
const TEST_SERVER_PORT = 4000;
const PROXY_PORT = 4001;
const TEST_DATA = 'Hello through port proxy!';

// Helper: Creates a test TCP server that listens on a given port and host.
function createTestServer(port: number, host: string = 'localhost'): Promise<net.Server> {
  return new Promise((resolve) => {
    const server = net.createServer((socket) => {
      socket.on('data', (data) => {
        // Echo the received data back with a prefix.
        socket.write(`Echo: ${data.toString()}`);
      });
      socket.on('error', (error) => {
        console.error(`[Test Server] Socket error on ${host}:${port}:`, error);
      });
    });
    server.listen(port, host, () => {
      console.log(`[Test Server] Listening on ${host}:${port}`);
      resolve(server);
    });
  });
}

// Helper: Creates a test client connection.
function createTestClient(port: number, data: string): Promise<string> {
  return new Promise((resolve, reject) => {
    const client = new net.Socket();
    let response = '';
    client.connect(port, 'localhost', () => {
      console.log('[Test Client] Connected to server');
      client.write(data);
    });
    client.on('data', (chunk) => {
      response += chunk.toString();
      client.end();
    });
    client.on('end', () => resolve(response));
    client.on('error', (error) => reject(error));
  });
}

// SETUP: Create a test server and a PortProxy instance.
tap.test('setup port proxy test environment', async () => {
  testServer = await createTestServer(TEST_SERVER_PORT);
  portProxy = new PortProxy({
    fromPort: PROXY_PORT,
    toPort: TEST_SERVER_PORT,
    targetIP: 'localhost',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1'],
    globalPortRanges: []
  });
});

// Test that the proxy starts and its servers are listening.
tap.test('should start port proxy', async () => {
  await portProxy.start();
  expect((portProxy as any).netServers.every((server: net.Server) => server.listening)).toBeTrue();
});

// Test basic TCP forwarding.
tap.test('should forward TCP connections and data to localhost', async () => {
  const response = await createTestClient(PROXY_PORT, TEST_DATA);
  expect(response).toEqual(`Echo: ${TEST_DATA}`);
});

// Test proxy with a custom target host.
tap.test('should forward TCP connections to custom host', async () => {
  const customHostProxy = new PortProxy({
    fromPort: PROXY_PORT + 1,
    toPort: TEST_SERVER_PORT,
    targetIP: '127.0.0.1',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1'],
    globalPortRanges: []
  });
  
  await customHostProxy.start();
  const response = await createTestClient(PROXY_PORT + 1, TEST_DATA);
  expect(response).toEqual(`Echo: ${TEST_DATA}`);
  await customHostProxy.stop();
});

// Test forced domain routing via port-range configuration.
// In this test, we want to forward to a different IP (using '127.0.0.2')
// while keeping the same port. We create a test server on '127.0.0.2'.
tap.test('should forward connections based on domain-specific target IP (forced domain via port-range)', async () => {
  const forcedProxyPort = PROXY_PORT + 2;
  // Create a test server listening on '127.0.0.2' at forcedProxyPort.
  const testServer2 = await createTestServer(forcedProxyPort, '127.0.0.2');

  const domainProxy = new PortProxy({
    fromPort: forcedProxyPort,
    toPort: TEST_SERVER_PORT, // default target port (unused for forced domain)
    targetIP: 'localhost',
    domainConfigs: [{
      domains: ['forced.test'],
      allowedIPs: ['127.0.0.1'],
      targetIPs: ['127.0.0.2'], // Use a different IP than the default.
      portRanges: [{ from: forcedProxyPort, to: forcedProxyPort }]
    }],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1'],
    globalPortRanges: [{ from: forcedProxyPort, to: forcedProxyPort }]
  });

  await domainProxy.start();

  // When connecting to forcedProxyPort, forced domain handling triggers,
  // so the proxy will connect to '127.0.0.2' on the same port.
  const response = await createTestClient(forcedProxyPort, TEST_DATA);
  expect(response).toEqual(`Echo: ${TEST_DATA}`);

  await domainProxy.stop();
  await new Promise<void>((resolve) => testServer2.close(() => resolve()));
});

// Test handling of multiple concurrent connections.
tap.test('should handle multiple concurrent connections', async () => {
  const concurrentRequests = 5;
  const requests = Array(concurrentRequests).fill(null).map((_, i) =>
    createTestClient(PROXY_PORT, `${TEST_DATA} ${i + 1}`)
  );
  const responses = await Promise.all(requests);
  responses.forEach((response, i) => {
    expect(response).toEqual(`Echo: ${TEST_DATA} ${i + 1}`);
  });
});

// Test connection timeout handling.
tap.test('should handle connection timeouts', async () => {
  const client = new net.Socket();
  await new Promise<void>((resolve) => {
    client.connect(PROXY_PORT, 'localhost', () => {
      // Do not send any data to trigger a timeout.
      client.on('close', () => resolve());
    });
  });
});

// Test stopping the port proxy.
tap.test('should stop port proxy', async () => {
  await portProxy.stop();
  expect((portProxy as any).netServers.every((server: net.Server) => !server.listening)).toBeTrue();
});

// Test chained proxies with and without source IP preservation.
tap.test('should support optional source IP preservation in chained proxies', async () => {
  // Chained proxies without IP preservation.
  const firstProxyDefault = new PortProxy({
    fromPort: PROXY_PORT + 4,
    toPort: PROXY_PORT + 5,
    targetIP: 'localhost',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1', '::ffff:127.0.0.1'],
    globalPortRanges: []
  });
  const secondProxyDefault = new PortProxy({
    fromPort: PROXY_PORT + 5,
    toPort: TEST_SERVER_PORT,
    targetIP: 'localhost',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1', '::ffff:127.0.0.1'],
    globalPortRanges: []
  });
  await secondProxyDefault.start();
  await firstProxyDefault.start();
  const response1 = await createTestClient(PROXY_PORT + 4, TEST_DATA);
  expect(response1).toEqual(`Echo: ${TEST_DATA}`);
  await firstProxyDefault.stop();
  await secondProxyDefault.stop();

  // Chained proxies with IP preservation.
  const firstProxyPreserved = new PortProxy({
    fromPort: PROXY_PORT + 6,
    toPort: PROXY_PORT + 7,
    targetIP: 'localhost',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1'],
    preserveSourceIP: true,
    globalPortRanges: []
  });
  const secondProxyPreserved = new PortProxy({
    fromPort: PROXY_PORT + 7,
    toPort: TEST_SERVER_PORT,
    targetIP: 'localhost',
    domainConfigs: [],
    sniEnabled: false,
    defaultAllowedIPs: ['127.0.0.1'],
    preserveSourceIP: true,
    globalPortRanges: []
  });
  await secondProxyPreserved.start();
  await firstProxyPreserved.start();
  const response2 = await createTestClient(PROXY_PORT + 6, TEST_DATA);
  expect(response2).toEqual(`Echo: ${TEST_DATA}`);
  await firstProxyPreserved.stop();
  await secondProxyPreserved.stop();
});

// Test round-robin behavior for multiple target IPs in a domain config.
tap.test('should use round robin for multiple target IPs in domain config', async () => {
  const domainConfig = {
    domains: ['rr.test'],
    allowedIPs: ['127.0.0.1'],
    targetIPs: ['hostA', 'hostB']
  } as any;
  
  const proxyInstance = new PortProxy({
    fromPort: 0,
    toPort: 0,
    targetIP: 'localhost',
    domainConfigs: [domainConfig],
    sniEnabled: false,
    defaultAllowedIPs: [],
    globalPortRanges: []
  });
  
  const firstTarget = (proxyInstance as any).getTargetIP(domainConfig);
  const secondTarget = (proxyInstance as any).getTargetIP(domainConfig);
  expect(firstTarget).toEqual('hostA');
  expect(secondTarget).toEqual('hostB');
});

// CLEANUP: Tear down the test server.
tap.test('cleanup port proxy test environment', async () => {
  await new Promise<void>((resolve) => testServer.close(() => resolve()));
});

process.on('exit', () => {
  if (testServer) {
    testServer.close();
  }
  if (portProxy && (portProxy as any).netServers) {
    portProxy.stop();
  }
});

export default tap.start();