import * as plugins from './plugins.js';
import type { IConnectionRecord, IDomainConfig, IPortProxySettings } from './classes.pp.interfaces.js';
import { ConnectionManager } from './classes.pp.connectionmanager.js';
import { SecurityManager } from './classes.pp.securitymanager.js';
import { DomainConfigManager } from './classes.pp.domainconfigmanager.js';
import { TlsManager } from './classes.pp.tlsmanager.js';
import { NetworkProxyBridge } from './classes.pp.networkproxybridge.js';
import { TimeoutManager } from './classes.pp.timeoutmanager.js';
import { PortRangeManager } from './classes.pp.portrangemanager.js';

/**
 * Handles new connection processing and setup logic
 */
export class ConnectionHandler {
  constructor(
    private settings: IPortProxySettings,
    private connectionManager: ConnectionManager,
    private securityManager: SecurityManager,
    private domainConfigManager: DomainConfigManager,
    private tlsManager: TlsManager,
    private networkProxyBridge: NetworkProxyBridge,
    private timeoutManager: TimeoutManager,
    private portRangeManager: PortRangeManager
  ) {}

  /**
   * Handle a new incoming connection
   */
  public handleConnection(socket: plugins.net.Socket): void {
    const remoteIP = socket.remoteAddress || '';
    const localPort = socket.localPort || 0;

    // Validate IP against rate limits and connection limits
    const ipValidation = this.securityManager.validateIP(remoteIP);
    if (!ipValidation.allowed) {
      console.log(`Connection rejected from ${remoteIP}: ${ipValidation.reason}`);
      socket.end();
      socket.destroy();
      return;
    }

    // Create a new connection record
    const record = this.connectionManager.createConnection(socket);
    const connectionId = record.id;

    // Apply socket optimizations
    socket.setNoDelay(this.settings.noDelay);

    // Apply keep-alive settings if enabled
    if (this.settings.keepAlive) {
      socket.setKeepAlive(true, this.settings.keepAliveInitialDelay);
      record.hasKeepAlive = true;

      // Apply enhanced TCP keep-alive options if enabled
      if (this.settings.enableKeepAliveProbes) {
        try {
          // These are platform-specific and may not be available
          if ('setKeepAliveProbes' in socket) {
            (socket as any).setKeepAliveProbes(10);
          }
          if ('setKeepAliveInterval' in socket) {
            (socket as any).setKeepAliveInterval(1000);
          }
        } catch (err) {
          // Ignore errors - these are optional enhancements
          if (this.settings.enableDetailedLogging) {
            console.log(`[${connectionId}] Enhanced TCP keep-alive settings not supported: ${err}`);
          }
        }
      }
    }

    if (this.settings.enableDetailedLogging) {
      console.log(
        `[${connectionId}] New connection from ${remoteIP} on port ${localPort}. ` +
        `Keep-Alive: ${record.hasKeepAlive ? 'Enabled' : 'Disabled'}. ` +
        `Active connections: ${this.connectionManager.getConnectionCount()}`
      );
    } else {
      console.log(
        `New connection from ${remoteIP} on port ${localPort}. Active connections: ${this.connectionManager.getConnectionCount()}`
      );
    }

    // Check if this connection should be forwarded directly to NetworkProxy
    if (this.portRangeManager.shouldUseNetworkProxy(localPort)) {
      this.handleNetworkProxyConnection(socket, record);
    } else {
      // For non-NetworkProxy ports, proceed with normal processing
      this.handleStandardConnection(socket, record);
    }
  }

  /**
   * Handle a connection that should be forwarded to NetworkProxy
   */
  private handleNetworkProxyConnection(socket: plugins.net.Socket, record: IConnectionRecord): void {
    const connectionId = record.id;
    let initialDataReceived = false;

    // Set an initial timeout for handshake data
    let initialTimeout: NodeJS.Timeout | null = setTimeout(() => {
      if (!initialDataReceived) {
        console.log(
          `[${connectionId}] Initial data warning (${this.settings.initialDataTimeout}ms) for connection from ${record.remoteIP}`
        );
        
        // Add a grace period instead of immediate termination
        setTimeout(() => {
          if (!initialDataReceived) {
            console.log(`[${connectionId}] Final initial data timeout after grace period`);
            if (record.incomingTerminationReason === null) {
              record.incomingTerminationReason = 'initial_timeout';
              this.connectionManager.incrementTerminationStat('incoming', 'initial_timeout');
            }
            socket.end();
            this.connectionManager.cleanupConnection(record, 'initial_timeout');
          }
        }, 30000); // 30 second grace period
      }
    }, this.settings.initialDataTimeout!);

    // Make sure timeout doesn't keep the process alive
    if (initialTimeout.unref) {
      initialTimeout.unref();
    }

    // Set up error handler
    socket.on('error', this.connectionManager.handleError('incoming', record));

    // First data handler to capture initial TLS handshake for NetworkProxy
    socket.once('data', (chunk: Buffer) => {
      // Clear the initial timeout since we've received data
      if (initialTimeout) {
        clearTimeout(initialTimeout);
        initialTimeout = null;
      }

      initialDataReceived = true;
      record.hasReceivedInitialData = true;

      // Block non-TLS connections on port 443
      const localPort = record.localPort;
      if (!this.tlsManager.isTlsHandshake(chunk) && localPort === 443) {
        console.log(
          `[${connectionId}] Non-TLS connection detected on port 443. ` +
          `Terminating connection - only TLS traffic is allowed on standard HTTPS port.`
        );
        if (record.incomingTerminationReason === null) {
          record.incomingTerminationReason = 'non_tls_blocked';
          this.connectionManager.incrementTerminationStat('incoming', 'non_tls_blocked');
        }
        socket.end();
        this.connectionManager.cleanupConnection(record, 'non_tls_blocked');
        return;
      }

      // Check if this looks like a TLS handshake
      if (this.tlsManager.isTlsHandshake(chunk)) {
        record.isTLS = true;

        // Check session tickets if they're disabled
        if (this.settings.allowSessionTicket === false && this.tlsManager.isClientHello(chunk)) {
          // Create connection info for SNI extraction
          const connInfo = {
            sourceIp: record.remoteIP,
            sourcePort: socket.remotePort || 0,
            destIp: socket.localAddress || '',
            destPort: socket.localPort || 0,
          };

          // Extract SNI for domain-specific NetworkProxy handling
          const serverName = this.tlsManager.extractSNI(chunk, connInfo);

          if (serverName) {
            // If we got an SNI, check for domain-specific NetworkProxy settings
            const domainConfig = this.domainConfigManager.findDomainConfig(serverName);

            // Save domain config and SNI in connection record
            record.domainConfig = domainConfig;
            record.lockedDomain = serverName;

            // Use domain-specific NetworkProxy port if configured
            if (domainConfig && this.domainConfigManager.shouldUseNetworkProxy(domainConfig)) {
              const networkProxyPort = this.domainConfigManager.getNetworkProxyPort(domainConfig);

              if (this.settings.enableDetailedLogging) {
                console.log(
                  `[${connectionId}] Using domain-specific NetworkProxy for ${serverName} on port ${networkProxyPort}`
                );
              }

              // Forward to NetworkProxy with domain-specific port
              this.networkProxyBridge.forwardToNetworkProxy(
                connectionId,
                socket,
                record,
                chunk,
                networkProxyPort,
                (reason) => this.connectionManager.initiateCleanupOnce(record, reason)
              );
              return;
            }
          }
        }

        // Forward directly to NetworkProxy without domain-specific settings
        this.networkProxyBridge.forwardToNetworkProxy(
          connectionId,
          socket,
          record,
          chunk,
          undefined,
          (reason) => this.connectionManager.initiateCleanupOnce(record, reason)
        );
      } else {
        // If not TLS, use normal direct connection
        console.log(`[${connectionId}] Non-TLS connection on NetworkProxy port ${record.localPort}`);
        this.setupDirectConnection(
          socket,
          record,
          undefined,
          undefined,
          chunk
        );
      }
    });
  }

  /**
   * Handle a standard (non-NetworkProxy) connection
   */
  private handleStandardConnection(socket: plugins.net.Socket, record: IConnectionRecord): void {
    const connectionId = record.id;
    const localPort = record.localPort;

    // Define helpers for rejecting connections
    const rejectIncomingConnection = (reason: string, logMessage: string) => {
      console.log(`[${connectionId}] ${logMessage}`);
      socket.end();
      if (record.incomingTerminationReason === null) {
        record.incomingTerminationReason = reason;
        this.connectionManager.incrementTerminationStat('incoming', reason);
      }
      this.connectionManager.cleanupConnection(record, reason);
    };

    let initialDataReceived = false;

    // Set an initial timeout for SNI data if needed
    let initialTimeout: NodeJS.Timeout | null = null;
    if (this.settings.sniEnabled) {
      initialTimeout = setTimeout(() => {
        if (!initialDataReceived) {
          console.log(
            `[${connectionId}] Initial data warning (${this.settings.initialDataTimeout}ms) for connection from ${record.remoteIP}`
          );
          
          // Add a grace period instead of immediate termination
          setTimeout(() => {
            if (!initialDataReceived) {
              console.log(`[${connectionId}] Final initial data timeout after grace period`);
              if (record.incomingTerminationReason === null) {
                record.incomingTerminationReason = 'initial_timeout';
                this.connectionManager.incrementTerminationStat('incoming', 'initial_timeout');
              }
              socket.end();
              this.connectionManager.cleanupConnection(record, 'initial_timeout');
            }
          }, 30000); // 30 second grace period
        }
      }, this.settings.initialDataTimeout!);

      // Make sure timeout doesn't keep the process alive
      if (initialTimeout.unref) {
        initialTimeout.unref();
      }
    } else {
      initialDataReceived = true;
      record.hasReceivedInitialData = true;
    }

    socket.on('error', this.connectionManager.handleError('incoming', record));

    // Track data for bytes counting
    socket.on('data', (chunk: Buffer) => {
      record.bytesReceived += chunk.length;
      this.timeoutManager.updateActivity(record);

      // Check for TLS handshake if this is the first chunk
      if (!record.isTLS && this.tlsManager.isTlsHandshake(chunk)) {
        record.isTLS = true;

        if (this.settings.enableTlsDebugLogging) {
          console.log(
            `[${connectionId}] TLS handshake detected from ${record.remoteIP}, ${chunk.length} bytes`
          );
        }
      }
    });

    /**
     * Sets up the connection to the target host.
     */
    const setupConnection = (
      serverName: string,
      initialChunk?: Buffer,
      forcedDomain?: IDomainConfig,
      overridePort?: number
    ) => {
      // Clear the initial timeout since we've received data
      if (initialTimeout) {
        clearTimeout(initialTimeout);
        initialTimeout = null;
      }

      // Mark that we've received initial data
      initialDataReceived = true;
      record.hasReceivedInitialData = true;

      // Check if this looks like a TLS handshake
      if (initialChunk && this.tlsManager.isTlsHandshake(initialChunk)) {
        record.isTLS = true;

        if (this.settings.enableTlsDebugLogging) {
          console.log(
            `[${connectionId}] TLS handshake detected in setup, ${initialChunk.length} bytes`
          );
        }
      }

      // If a forcedDomain is provided (port-based routing), use it; otherwise, use SNI-based lookup.
      const domainConfig = forcedDomain
        ? forcedDomain
        : serverName
        ? this.domainConfigManager.findDomainConfig(serverName)
        : undefined;

      // Save domain config in connection record
      record.domainConfig = domainConfig;

      // Check if this domain should use NetworkProxy (domain-specific setting)
      if (domainConfig && 
          this.domainConfigManager.shouldUseNetworkProxy(domainConfig) && 
          this.networkProxyBridge.getNetworkProxy()) {
        
        if (this.settings.enableDetailedLogging) {
          console.log(
            `[${connectionId}] Domain ${serverName} is configured to use NetworkProxy`
          );
        }

        const networkProxyPort = this.domainConfigManager.getNetworkProxyPort(domainConfig);

        if (initialChunk && record.isTLS) {
          // For TLS connections with initial chunk, forward to NetworkProxy
          this.networkProxyBridge.forwardToNetworkProxy(
            connectionId,
            socket,
            record,
            initialChunk,
            networkProxyPort,
            (reason) => this.connectionManager.initiateCleanupOnce(record, reason)
          );
          return; // Skip normal connection setup
        }
      }

      // IP validation
      if (domainConfig) {
        const ipRules = this.domainConfigManager.getEffectiveIPRules(domainConfig);
        
        // Skip IP validation if allowedIPs is empty
        if (
          domainConfig.allowedIPs.length > 0 &&
          !this.securityManager.isIPAuthorized(record.remoteIP, ipRules.allowedIPs, ipRules.blockedIPs)
        ) {
          return rejectIncomingConnection(
            'rejected',
            `Connection rejected: IP ${record.remoteIP} not allowed for domain ${domainConfig.domains.join(
              ', '
            )}`
          );
        }
      } else if (
        this.settings.defaultAllowedIPs &&
        this.settings.defaultAllowedIPs.length > 0
      ) {
        if (
          !this.securityManager.isIPAuthorized(
            record.remoteIP,
            this.settings.defaultAllowedIPs,
            this.settings.defaultBlockedIPs || []
          )
        ) {
          return rejectIncomingConnection(
            'rejected',
            `Connection rejected: IP ${record.remoteIP} not allowed by default allowed list`
          );
        }
      }

      // Save the initial SNI
      if (serverName) {
        record.lockedDomain = serverName;
      }

      // Set up the direct connection
      this.setupDirectConnection(
        socket,
        record,
        domainConfig,
        serverName,
        initialChunk,
        overridePort
      );
    };

    // --- PORT RANGE-BASED HANDLING ---
    // Only apply port-based rules if the incoming port is within one of the global port ranges.
    if (this.portRangeManager.isPortInGlobalRanges(localPort)) {
      if (this.portRangeManager.shouldUseGlobalForwarding(localPort)) {
        if (
          this.settings.defaultAllowedIPs &&
          this.settings.defaultAllowedIPs.length > 0 &&
          !this.securityManager.isIPAuthorized(record.remoteIP, this.settings.defaultAllowedIPs)
        ) {
          console.log(
            `[${connectionId}] Connection from ${record.remoteIP} rejected: IP ${record.remoteIP} not allowed in global default allowed list.`
          );
          socket.end();
          return;
        }
        if (this.settings.enableDetailedLogging) {
          console.log(
            `[${connectionId}] Port-based connection from ${record.remoteIP} on port ${localPort} forwarded to global target IP ${this.settings.targetIP}.`
          );
        }
        setupConnection(
          '',
          undefined,
          {
            domains: ['global'],
            allowedIPs: this.settings.defaultAllowedIPs || [],
            blockedIPs: this.settings.defaultBlockedIPs || [],
            targetIPs: [this.settings.targetIP!],
            portRanges: [],
          },
          localPort
        );
        return;
      } else {
        // Attempt to find a matching forced domain config based on the local port.
        const forcedDomain = this.domainConfigManager.findDomainConfigForPort(localPort);
        
        if (forcedDomain) {
          const ipRules = this.domainConfigManager.getEffectiveIPRules(forcedDomain);
          
          if (!this.securityManager.isIPAuthorized(record.remoteIP, ipRules.allowedIPs, ipRules.blockedIPs)) {
            console.log(
              `[${connectionId}] Connection from ${record.remoteIP} rejected: IP not allowed for domain ${forcedDomain.domains.join(
                ', '
              )} on port ${localPort}.`
            );
            socket.end();
            return;
          }
          
          if (this.settings.enableDetailedLogging) {
            console.log(
              `[${connectionId}] Port-based connection from ${record.remoteIP} on port ${localPort} matched domain ${forcedDomain.domains.join(
                ', '
              )}.`
            );
          }
          
          setupConnection('', undefined, forcedDomain, localPort);
          return;
        }
        // Fall through to SNI/default handling if no forced domain config is found.
      }
    }

    // --- FALLBACK: SNI-BASED HANDLING (or default when SNI is disabled) ---
    if (this.settings.sniEnabled) {
      initialDataReceived = false;

      socket.once('data', (chunk: Buffer) => {
        // Clear timeout immediately
        if (initialTimeout) {
          clearTimeout(initialTimeout);
          initialTimeout = null;
        }
        
        initialDataReceived = true;
        
        // Block non-TLS connections on port 443
        if (!this.tlsManager.isTlsHandshake(chunk) && localPort === 443) {
          console.log(
            `[${connectionId}] Non-TLS connection detected on port 443 in SNI handler. ` +
            `Terminating connection - only TLS traffic is allowed on standard HTTPS port.`
          );
          if (record.incomingTerminationReason === null) {
            record.incomingTerminationReason = 'non_tls_blocked';
            this.connectionManager.incrementTerminationStat('incoming', 'non_tls_blocked');
          }
          socket.end();
          this.connectionManager.cleanupConnection(record, 'non_tls_blocked');
          return;
        }

        // Try to extract SNI
        let serverName = '';

        if (this.tlsManager.isTlsHandshake(chunk)) {
          record.isTLS = true;

          if (this.settings.enableTlsDebugLogging) {
            console.log(
              `[${connectionId}] Extracting SNI from TLS handshake, ${chunk.length} bytes`
            );
          }

          // Create connection info object for SNI extraction
          const connInfo = {
            sourceIp: record.remoteIP,
            sourcePort: socket.remotePort || 0,
            destIp: socket.localAddress || '',
            destPort: socket.localPort || 0,
          };

          // Extract SNI
          serverName = this.tlsManager.extractSNI(chunk, connInfo) || '';
        }

        // Lock the connection to the negotiated SNI.
        record.lockedDomain = serverName;

        if (this.settings.enableDetailedLogging) {
          console.log(
            `[${connectionId}] Received connection from ${record.remoteIP} with SNI: ${
              serverName || '(empty)'
            }`
          );
        }

        setupConnection(serverName, chunk);
      });
    } else {
      initialDataReceived = true;
      record.hasReceivedInitialData = true;

      if (
        this.settings.defaultAllowedIPs &&
        this.settings.defaultAllowedIPs.length > 0 &&
        !this.securityManager.isIPAuthorized(record.remoteIP, this.settings.defaultAllowedIPs)
      ) {
        return rejectIncomingConnection(
          'rejected',
          `Connection rejected: IP ${record.remoteIP} not allowed for non-SNI connection`
        );
      }

      setupConnection('');
    }
  }

  /**
   * Sets up a direct connection to the target
   */
  private setupDirectConnection(
    socket: plugins.net.Socket,
    record: IConnectionRecord,
    domainConfig?: IDomainConfig,
    serverName?: string,
    initialChunk?: Buffer,
    overridePort?: number
  ): void {
    const connectionId = record.id;
    
    // Determine target host
    const targetHost = domainConfig 
      ? this.domainConfigManager.getTargetIP(domainConfig) 
      : this.settings.targetIP!;
    
    // Determine target port
    const targetPort = overridePort !== undefined 
      ? overridePort 
      : this.settings.toPort;
    
    // Setup connection options
    const connectionOptions: plugins.net.NetConnectOpts = {
      host: targetHost,
      port: targetPort,
    };
    
    // Preserve source IP if configured
    if (this.settings.preserveSourceIP) {
      connectionOptions.localAddress = record.remoteIP.replace('::ffff:', '');
    }

    // Create a safe queue for incoming data
    const dataQueue: Buffer[] = [];
    let queueSize = 0;
    let processingQueue = false;
    let drainPending = false;
    let pipingEstablished = false;

    // Pause the incoming socket to prevent buffer overflows
    socket.pause();

    // Function to safely process the data queue without losing events
    const processDataQueue = () => {
      if (processingQueue || dataQueue.length === 0 || pipingEstablished) return;

      processingQueue = true;

      try {
        // Process all queued chunks with the current active handler
        while (dataQueue.length > 0) {
          const chunk = dataQueue.shift()!;
          queueSize -= chunk.length;

          // Once piping is established, we shouldn't get here,
          // but just in case, pass to the outgoing socket directly
          if (pipingEstablished && record.outgoing) {
            record.outgoing.write(chunk);
            continue;
          }

          // Track bytes received
          record.bytesReceived += chunk.length;

          // Check for TLS handshake
          if (!record.isTLS && this.tlsManager.isTlsHandshake(chunk)) {
            record.isTLS = true;

            if (this.settings.enableTlsDebugLogging) {
              console.log(
                `[${connectionId}] TLS handshake detected in tempDataHandler, ${chunk.length} bytes`
              );
            }
          }

          // Check if adding this chunk would exceed the buffer limit
          const newSize = record.pendingDataSize + chunk.length;

          if (this.settings.maxPendingDataSize && newSize > this.settings.maxPendingDataSize) {
            console.log(
              `[${connectionId}] Buffer limit exceeded for connection from ${record.remoteIP}: ${newSize} bytes > ${this.settings.maxPendingDataSize} bytes`
            );
            socket.end(); // Gracefully close the socket
            this.connectionManager.initiateCleanupOnce(record, 'buffer_limit_exceeded');
            return;
          }

          // Buffer the chunk and update the size counter
          record.pendingData.push(Buffer.from(chunk));
          record.pendingDataSize = newSize;
          this.timeoutManager.updateActivity(record);
        }
      } finally {
        processingQueue = false;

        // If there's a pending drain and we've processed everything,
        // signal we're ready for more data if we haven't established piping yet
        if (drainPending && dataQueue.length === 0 && !pipingEstablished) {
          drainPending = false;
          socket.resume();
        }
      }
    };

    // Unified data handler that safely queues incoming data
    const safeDataHandler = (chunk: Buffer) => {
      // If piping is already established, just let the pipe handle it
      if (pipingEstablished) return;

      // Add to our queue for orderly processing
      dataQueue.push(Buffer.from(chunk)); // Make a copy to be safe
      queueSize += chunk.length;

      // If queue is getting large, pause socket until we catch up
      if (this.settings.maxPendingDataSize && queueSize > this.settings.maxPendingDataSize * 0.8) {
        socket.pause();
        drainPending = true;
      }

      // Process the queue
      processDataQueue();
    };

    // Add our safe data handler
    socket.on('data', safeDataHandler);

    // Add initial chunk to pending data if present
    if (initialChunk) {
      record.bytesReceived += initialChunk.length;
      record.pendingData.push(Buffer.from(initialChunk));
      record.pendingDataSize = initialChunk.length;
    }

    // Create the target socket but don't set up piping immediately
    const targetSocket = plugins.net.connect(connectionOptions);
    record.outgoing = targetSocket;
    record.outgoingStartTime = Date.now();

    // Apply socket optimizations
    targetSocket.setNoDelay(this.settings.noDelay);

    // Apply keep-alive settings to the outgoing connection as well
    if (this.settings.keepAlive) {
      targetSocket.setKeepAlive(true, this.settings.keepAliveInitialDelay);

      // Apply enhanced TCP keep-alive options if enabled
      if (this.settings.enableKeepAliveProbes) {
        try {
          if ('setKeepAliveProbes' in targetSocket) {
            (targetSocket as any).setKeepAliveProbes(10);
          }
          if ('setKeepAliveInterval' in targetSocket) {
            (targetSocket as any).setKeepAliveInterval(1000);
          }
        } catch (err) {
          // Ignore errors - these are optional enhancements
          if (this.settings.enableDetailedLogging) {
            console.log(
              `[${connectionId}] Enhanced TCP keep-alive not supported for outgoing socket: ${err}`
            );
          }
        }
      }
    }

    // Setup specific error handler for connection phase
    targetSocket.once('error', (err) => {
      // This handler runs only once during the initial connection phase
      const code = (err as any).code;
      console.log(
        `[${connectionId}] Connection setup error to ${targetHost}:${connectionOptions.port}: ${err.message} (${code})`
      );

      // Resume the incoming socket to prevent it from hanging
      socket.resume();

      if (code === 'ECONNREFUSED') {
        console.log(
          `[${connectionId}] Target ${targetHost}:${connectionOptions.port} refused connection`
        );
      } else if (code === 'ETIMEDOUT') {
        console.log(
          `[${connectionId}] Connection to ${targetHost}:${connectionOptions.port} timed out`
        );
      } else if (code === 'ECONNRESET') {
        console.log(
          `[${connectionId}] Connection to ${targetHost}:${connectionOptions.port} was reset`
        );
      } else if (code === 'EHOSTUNREACH') {
        console.log(`[${connectionId}] Host ${targetHost} is unreachable`);
      }

      // Clear any existing error handler after connection phase
      targetSocket.removeAllListeners('error');

      // Re-add the normal error handler for established connections
      targetSocket.on('error', this.connectionManager.handleError('outgoing', record));

      if (record.outgoingTerminationReason === null) {
        record.outgoingTerminationReason = 'connection_failed';
        this.connectionManager.incrementTerminationStat('outgoing', 'connection_failed');
      }

      // Clean up the connection
      this.connectionManager.initiateCleanupOnce(record, `connection_failed_${code}`);
    });

    // Setup close handler
    targetSocket.on('close', this.connectionManager.handleClose('outgoing', record));
    socket.on('close', this.connectionManager.handleClose('incoming', record));

    // Handle timeouts with keep-alive awareness
    socket.on('timeout', () => {
      // For keep-alive connections, just log a warning instead of closing
      if (record.hasKeepAlive) {
        console.log(
          `[${connectionId}] Timeout event on incoming keep-alive connection from ${
            record.remoteIP
          } after ${plugins.prettyMs(
            this.settings.socketTimeout || 3600000
          )}. Connection preserved.`
        );
        return;
      }

      // For non-keep-alive connections, proceed with normal cleanup
      console.log(
        `[${connectionId}] Timeout on incoming side from ${
          record.remoteIP
        } after ${plugins.prettyMs(this.settings.socketTimeout || 3600000)}`
      );
      if (record.incomingTerminationReason === null) {
        record.incomingTerminationReason = 'timeout';
        this.connectionManager.incrementTerminationStat('incoming', 'timeout');
      }
      this.connectionManager.initiateCleanupOnce(record, 'timeout_incoming');
    });

    targetSocket.on('timeout', () => {
      // For keep-alive connections, just log a warning instead of closing
      if (record.hasKeepAlive) {
        console.log(
          `[${connectionId}] Timeout event on outgoing keep-alive connection from ${
            record.remoteIP
          } after ${plugins.prettyMs(
            this.settings.socketTimeout || 3600000
          )}. Connection preserved.`
        );
        return;
      }

      // For non-keep-alive connections, proceed with normal cleanup
      console.log(
        `[${connectionId}] Timeout on outgoing side from ${
          record.remoteIP
        } after ${plugins.prettyMs(this.settings.socketTimeout || 3600000)}`
      );
      if (record.outgoingTerminationReason === null) {
        record.outgoingTerminationReason = 'timeout';
        this.connectionManager.incrementTerminationStat('outgoing', 'timeout');
      }
      this.connectionManager.initiateCleanupOnce(record, 'timeout_outgoing');
    });

    // Apply socket timeouts
    this.timeoutManager.applySocketTimeouts(record);

    // Track outgoing data for bytes counting
    targetSocket.on('data', (chunk: Buffer) => {
      record.bytesSent += chunk.length;
      this.timeoutManager.updateActivity(record);
    });

    // Wait for the outgoing connection to be ready before setting up piping
    targetSocket.once('connect', () => {
      // Clear the initial connection error handler
      targetSocket.removeAllListeners('error');

      // Add the normal error handler for established connections
      targetSocket.on('error', this.connectionManager.handleError('outgoing', record));

      // Process any remaining data in the queue before switching to piping
      processDataQueue();
      
      // Set up piping immediately
      pipingEstablished = true;
      
      // Flush all pending data to target
      if (record.pendingData.length > 0) {
        const combinedData = Buffer.concat(record.pendingData);
        
        if (this.settings.enableDetailedLogging) {
          console.log(`[${connectionId}] Forwarding ${combinedData.length} bytes of initial data to target`);
        }
        
        // Write pending data immediately
        targetSocket.write(combinedData, (err) => {
          if (err) {
            console.log(`[${connectionId}] Error writing pending data to target: ${err.message}`);
            return this.connectionManager.initiateCleanupOnce(record, 'write_error');
          }
        });
        
        // Clear the buffer now that we've processed it
        record.pendingData = [];
        record.pendingDataSize = 0;
      }
      
      // Setup piping in both directions without any delays
      socket.pipe(targetSocket);
      targetSocket.pipe(socket);
      
      // Resume the socket to ensure data flows
      socket.resume();
      
      // Process any data that might be queued in the interim
      if (dataQueue.length > 0) {
        // Write any remaining queued data directly to the target socket
        for (const chunk of dataQueue) {
          targetSocket.write(chunk);
        }
        // Clear the queue
        dataQueue.length = 0;
        queueSize = 0;
      }
      
      if (this.settings.enableDetailedLogging) {
        console.log(
          `[${connectionId}] Connection established: ${record.remoteIP} -> ${targetHost}:${connectionOptions.port}` +
            `${
              serverName
                ? ` (SNI: ${serverName})`
                : domainConfig
                ? ` (Port-based for domain: ${domainConfig.domains.join(', ')})`
                : ''
            }` +
            ` TLS: ${record.isTLS ? 'Yes' : 'No'}, Keep-Alive: ${
              record.hasKeepAlive ? 'Yes' : 'No'
            }`
        );
      } else {
        console.log(
          `Connection established: ${record.remoteIP} -> ${targetHost}:${connectionOptions.port}` +
            `${
              serverName
                ? ` (SNI: ${serverName})`
                : domainConfig
                ? ` (Port-based for domain: ${domainConfig.domains.join(', ')})`
                : ''
            }`
        );
      }

      // Add the renegotiation handler for SNI validation
      if (serverName) {
        // Create connection info object for the existing connection
        const connInfo = {
          sourceIp: record.remoteIP,
          sourcePort: record.incoming.remotePort || 0,
          destIp: record.incoming.localAddress || '',
          destPort: record.incoming.localPort || 0,
        };

        // Create a renegotiation handler function
        const renegotiationHandler = this.tlsManager.createRenegotiationHandler(
          connectionId,
          serverName,
          connInfo,
          (connectionId, reason) => this.connectionManager.initiateCleanupOnce(record, reason)
        );

        // Store the handler in the connection record so we can remove it during cleanup
        record.renegotiationHandler = renegotiationHandler;

        // Add the handler to the socket
        socket.on('data', renegotiationHandler);

        if (this.settings.enableDetailedLogging) {
          console.log(
            `[${connectionId}] TLS renegotiation handler installed for SNI domain: ${serverName}`
          );
          if (this.settings.allowSessionTicket === false) {
            console.log(
              `[${connectionId}] Session ticket usage is disabled. Connection will be reset on reconnection attempts.`
            );
          }
        }
      }

      // Set connection timeout
      record.cleanupTimer = this.timeoutManager.setupConnectionTimeout(
        record,
        (record, reason) => {
          console.log(
            `[${connectionId}] Connection from ${record.remoteIP} exceeded max lifetime, forcing cleanup.`
          );
          this.connectionManager.initiateCleanupOnce(record, reason);
        }
      );

      // Mark TLS handshake as complete for TLS connections
      if (record.isTLS) {
        record.tlsHandshakeComplete = true;

        if (this.settings.enableTlsDebugLogging) {
          console.log(
            `[${connectionId}] TLS handshake complete for connection from ${record.remoteIP}`
          );
        }
      }
    });
  }
}