import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as smartproxy from '../ts/index.js';
import * as net from 'net';
import * as crypto from 'crypto';

tap.test('Connection Stability - Fragment Cleanup', async () => {
  // Create a simple TCP server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, '127.0.0.1', () => resolve());
  });
  const serverPort = (server.address() as net.AddressInfo).port;

  // Configure a route
  const routes: smartproxy.IRouteConfig[] = [{
    match: {
      ports: 9000,
      domains: '*'
    },
    action: {
      type: 'forward',
      target: {
        host: '127.0.0.1',
        port: serverPort
      }
    }
  }];
  
  // Create SmartProxy instance with routes
  const proxy = new smartproxy.SmartProxy({ 
    keepAliveTimeoutMs: 5000,
    routes 
  });
  
  await proxy.start();

  // Test 1: Send fragmented TLS hello
  const tlsHello = Buffer.concat([
    Buffer.from([0x16, 0x03, 0x03]), // TLS handshake, version 1.2
    Buffer.from([0x00, 0x50]), // Length: 80 bytes
    Buffer.from([0x01]), // ClientHello
    Buffer.from([0x00, 0x00, 0x4c]), // Handshake length
    Buffer.from([0x03, 0x03]), // TLS 1.2
    crypto.randomBytes(32), // Random
    Buffer.from([0x00]), // Session ID length
    Buffer.from([0x00, 0x04]), // Cipher suites length
    Buffer.from([0xc0, 0x2f, 0xc0, 0x30]), // Cipher suites
    Buffer.from([0x01, 0x00]), // Compression methods
    Buffer.from([0x00, 0x1f]), // Extensions length
    // SNI extension
    Buffer.from([0x00, 0x00]), // Server name extension
    Buffer.from([0x00, 0x1b]), // Extension length
    Buffer.from([0x00, 0x19]), // Server name list length
    Buffer.from([0x00]), // Host name type
    Buffer.from([0x00, 0x16]), // Name length
    Buffer.from('test.example.com') // Server name
  ]);

  // Function to check fragment manager size
  const getFragmentCount = () => {
    // Access the fragment manager through the singleton
    const detector = (smartproxy.detection.ProtocolDetector as any).getInstance();
    const tlsFragments = detector.fragmentManager.getHandler('tls');
    const httpFragments = detector.fragmentManager.getHandler('http');
    return tlsFragments.size + httpFragments.size;
  };

  // Test fragmented connections
  const connections: net.Socket[] = [];
  
  // Create multiple fragmented connections
  for (let i = 0; i < 5; i++) {
    const client = new net.Socket();
    connections.push(client);
    
    await new Promise<void>((resolve, reject) => {
      client.connect(9000, '127.0.0.1', () => {
        // Send first fragment
        client.write(tlsHello.subarray(0, 20));
        resolve();
      });
      client.on('error', reject);
    });
  }

  // Give time for fragments to accumulate
  await new Promise(resolve => setTimeout(resolve, 100));

  // Check that fragments are being tracked
  const fragmentCount = getFragmentCount();
  expect(fragmentCount).toBeGreaterThan(0);

  // Send remaining fragments and close connections
  for (const client of connections) {
    client.write(tlsHello.subarray(20));
    client.end();
  }

  // Wait for connections to close
  await new Promise(resolve => setTimeout(resolve, 500));

  // Check that fragments are cleaned up
  const finalFragmentCount = getFragmentCount();
  expect(finalFragmentCount).toEqual(0);

  // Cleanup
  await proxy.stop();
  server.close();
});

tap.test('Connection Stability - Memory Leak Prevention', async () => {
  // Create a simple echo server
  const server = net.createServer((socket) => {
    socket.pipe(socket);
  });
  
  await new Promise<void>((resolve) => {
    server.listen(0, '127.0.0.1', () => resolve());
  });
  const serverPort = (server.address() as net.AddressInfo).port;

  // Configure a route
  const routes: smartproxy.IRouteConfig[] = [{
    match: {
      ports: 9001,
      domains: '*'
    },
    action: {
      type: 'forward',
      target: {
        host: '127.0.0.1',
        port: serverPort
      }
    }
  }];
  
  // Create SmartProxy instance with routes
  const proxy = new smartproxy.SmartProxy({ 
    keepAliveTimeoutMs: 5000,
    routes 
  });
  
  await proxy.start();

  // Function to get active connection count
  const getConnectionCount = () => {
    const connectionManager = (proxy as any).connectionManager;
    return connectionManager.getActiveConnectionCount();
  };

  // Create many short-lived connections
  const connectionPromises: Promise<void>[] = [];
  
  for (let i = 0; i < 20; i++) {
    const promise = new Promise<void>((resolve, reject) => {
      const client = new net.Socket();
      
      client.connect(9001, '127.0.0.1', () => {
        // Send some data
        client.write('Hello World');
        
        // Close after a short delay
        setTimeout(() => {
          client.end();
        }, 50);
      });
      
      client.on('close', () => resolve());
      client.on('error', reject);
    });
    
    connectionPromises.push(promise);
    
    // Stagger connection creation
    await new Promise(resolve => setTimeout(resolve, 10));
  }

  // Wait for all connections to complete
  await Promise.all(connectionPromises);

  // Give time for cleanup
  await new Promise(resolve => setTimeout(resolve, 1000));

  // Check that all connections are cleaned up
  const finalConnectionCount = getConnectionCount();
  expect(finalConnectionCount).toEqual(0);

  // Check fragment cleanup
  const fragmentCount = (() => {
    const detector = (smartproxy.detection.ProtocolDetector as any).getInstance();
    const tlsFragments = detector.fragmentManager.getHandler('tls');
    const httpFragments = detector.fragmentManager.getHandler('http');
    return tlsFragments.size + httpFragments.size;
  })();
  
  expect(fragmentCount).toEqual(0);

  // Cleanup
  await proxy.stop();
  server.close();
});

tap.test('Connection Stability - Rapid Connect/Disconnect', async () => {
  // Create a server that immediately closes connections
  const server = net.createServer((socket) => {
    socket.end();
  });
  
  await new Promise<void>((resolve) => {
    server.listen(0, '127.0.0.1', () => resolve());
  });
  const serverPort = (server.address() as net.AddressInfo).port;

  // Configure a route
  const routes: smartproxy.IRouteConfig[] = [{
    match: {
      ports: 9002,
      domains: '*'
    },
    action: {
      type: 'forward',
      target: {
        host: '127.0.0.1',
        port: serverPort
      }
    }
  }];
  
  // Create SmartProxy instance with routes
  const proxy = new smartproxy.SmartProxy({ 
    keepAliveTimeoutMs: 5000,
    routes 
  });
  
  await proxy.start();

  let errors = 0;
  const connections: Promise<void>[] = [];

  // Create many rapid connections
  for (let i = 0; i < 50; i++) {
    const promise = new Promise<void>((resolve) => {
      const client = new net.Socket();
      
      client.on('error', () => {
        errors++;
        resolve();
      });
      
      client.on('close', () => {
        resolve();
      });
      
      client.connect(9002, '127.0.0.1');
    });
    
    connections.push(promise);
  }

  // Wait for all to complete
  await Promise.all(connections);

  // Give time for cleanup
  await new Promise(resolve => setTimeout(resolve, 500));

  // Check that connections are cleaned up despite rapid connect/disconnect
  const connectionManager = (proxy as any).connectionManager;
  const finalConnectionCount = connectionManager.getActiveConnectionCount();
  expect(finalConnectionCount).toEqual(0);

  // Check fragment cleanup
  const fragmentCount = (() => {
    const detector = (smartproxy.detection.ProtocolDetector as any).getInstance();
    const tlsFragments = detector.fragmentManager.getHandler('tls');
    const httpFragments = detector.fragmentManager.getHandler('http');
    return tlsFragments.size + httpFragments.size;
  })();
  
  expect(fragmentCount).toEqual(0);

  // Cleanup
  await proxy.stop();
  server.close();
});

tap.start();