import * as plugins from './plugins.js';

/** Domain configuration with per-domain allowed port ranges */
export interface IDomainConfig {
  domains: string[];         // Glob patterns for domain(s)
  allowedIPs: string[];      // Glob patterns for allowed IPs
  blockedIPs?: string[];     // Glob patterns for blocked IPs
  targetIPs?: string[];      // If multiple targetIPs are given, use round robin.
  portRanges?: Array<{ from: number; to: number }>; // Optional port ranges
}

/** Port proxy settings including global allowed port ranges */
export interface IPortProxySettings extends plugins.tls.TlsOptions {
  fromPort: number;
  toPort: number;
  targetIP?: string; // Global target host to proxy to, defaults to 'localhost'
  domainConfigs: IDomainConfig[];
  sniEnabled?: boolean;
  defaultAllowedIPs?: string[];
  defaultBlockedIPs?: string[];
  preserveSourceIP?: boolean;
  maxConnectionLifetime?: number; // (ms) force cleanup of long-lived connections
  globalPortRanges: Array<{ from: number; to: number }>; // Global allowed port ranges
  forwardAllGlobalRanges?: boolean; // When true, forwards all connections on global port ranges to the global targetIP
  gracefulShutdownTimeout?: number; // (ms) maximum time to wait for connections to close during shutdown
}

/**
 * Extracts the SNI (Server Name Indication) from a TLS ClientHello packet.
 * @param buffer - Buffer containing the TLS ClientHello.
 * @returns The server name if found, otherwise undefined.
 */
function extractSNI(buffer: Buffer): string | undefined {
  let offset = 0;
  if (buffer.length < 5) return undefined;

  const recordType = buffer.readUInt8(0);
  if (recordType !== 22) return undefined; // 22 = handshake

  const recordLength = buffer.readUInt16BE(3);
  if (buffer.length < 5 + recordLength) return undefined;

  offset = 5;
  const handshakeType = buffer.readUInt8(offset);
  if (handshakeType !== 1) return undefined; // 1 = ClientHello

  offset += 4; // Skip handshake header (type + length)
  offset += 2 + 32; // Skip client version and random

  const sessionIDLength = buffer.readUInt8(offset);
  offset += 1 + sessionIDLength; // Skip session ID

  const cipherSuitesLength = buffer.readUInt16BE(offset);
  offset += 2 + cipherSuitesLength; // Skip cipher suites

  const compressionMethodsLength = buffer.readUInt8(offset);
  offset += 1 + compressionMethodsLength; // Skip compression methods

  if (offset + 2 > buffer.length) return undefined;
  const extensionsLength = buffer.readUInt16BE(offset);
  offset += 2;
  const extensionsEnd = offset + extensionsLength;

  while (offset + 4 <= extensionsEnd) {
    const extensionType = buffer.readUInt16BE(offset);
    const extensionLength = buffer.readUInt16BE(offset + 2);
    offset += 4;
    if (extensionType === 0x0000) { // SNI extension
      if (offset + 2 > buffer.length) return undefined;
      const sniListLength = buffer.readUInt16BE(offset);
      offset += 2;
      const sniListEnd = offset + sniListLength;
      while (offset + 3 < sniListEnd) {
        const nameType = buffer.readUInt8(offset++);
        const nameLen = buffer.readUInt16BE(offset);
        offset += 2;
        if (nameType === 0) { // host_name
          if (offset + nameLen > buffer.length) return undefined;
          return buffer.toString('utf8', offset, offset + nameLen);
        }
        offset += nameLen;
      }
      break;
    } else {
      offset += extensionLength;
    }
  }
  return undefined;
}

interface IConnectionRecord {
  id: string;                       // Unique connection identifier
  incoming: plugins.net.Socket;
  outgoing: plugins.net.Socket | null;
  incomingStartTime: number;
  outgoingStartTime?: number;
  outgoingClosedTime?: number;
  lockedDomain?: string;            // Used to lock this connection to the initial SNI
  connectionClosed: boolean;        // Flag to prevent multiple cleanup attempts
  cleanupTimer?: NodeJS.Timeout;    // Timer for max lifetime/inactivity
  lastActivity: number;             // Last activity timestamp for inactivity detection
}

// Helper: Check if a port falls within any of the given port ranges
const isPortInRanges = (port: number, ranges: Array<{ from: number; to: number }>): boolean => {
  return ranges.some(range => port >= range.from && port <= range.to);
};

// Helper: Check if a given IP matches any of the glob patterns
const isAllowed = (ip: string, patterns: string[]): boolean => {
  const normalizeIP = (ip: string): string[] => {
    if (ip.startsWith('::ffff:')) {
      const ipv4 = ip.slice(7);
      return [ip, ipv4];
    }
    if (/^\d{1,3}(\.\d{1,3}){3}$/.test(ip)) {
      return [ip, `::ffff:${ip}`];
    }
    return [ip];
  };
  const normalizedIPVariants = normalizeIP(ip);
  const expandedPatterns = patterns.flatMap(normalizeIP);
  return normalizedIPVariants.some(ipVariant =>
    expandedPatterns.some(pattern => plugins.minimatch(ipVariant, pattern))
  );
};

// Helper: Check if an IP is allowed considering allowed and blocked glob patterns
const isGlobIPAllowed = (ip: string, allowed: string[], blocked: string[] = []): boolean => {
  if (blocked.length > 0 && isAllowed(ip, blocked)) return false;
  return isAllowed(ip, allowed);
};

// Helper: Generate a unique connection ID
const generateConnectionId = (): string => {
  return Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
};

export class PortProxy {
  private netServers: plugins.net.Server[] = [];
  settings: IPortProxySettings;
  private connectionRecords: Map<string, IConnectionRecord> = new Map();
  private connectionLogger: NodeJS.Timeout | null = null;
  private isShuttingDown: boolean = false;

  // Map to track round robin indices for each domain config
  private domainTargetIndices: Map<IDomainConfig, number> = new Map();

  private terminationStats: {
    incoming: Record<string, number>;
    outgoing: Record<string, number>;
  } = {
    incoming: {},
    outgoing: {},
  };

  constructor(settingsArg: IPortProxySettings) {
    this.settings = {
      ...settingsArg,
      targetIP: settingsArg.targetIP || 'localhost',
      maxConnectionLifetime: settingsArg.maxConnectionLifetime || 600000,
      gracefulShutdownTimeout: settingsArg.gracefulShutdownTimeout || 30000,
    };
  }

  private incrementTerminationStat(side: 'incoming' | 'outgoing', reason: string): void {
    this.terminationStats[side][reason] = (this.terminationStats[side][reason] || 0) + 1;
  }

  /**
   * Cleans up a connection record.
   * Destroys both incoming and outgoing sockets, clears timers, and removes the record.
   * @param record - The connection record to clean up
   * @param reason - Optional reason for cleanup (for logging)
   */
  private cleanupConnection(record: IConnectionRecord, reason: string = 'normal'): void {
    if (!record.connectionClosed) {
      record.connectionClosed = true;
      
      if (record.cleanupTimer) {
        clearTimeout(record.cleanupTimer);
        record.cleanupTimer = undefined;
      }
      
      try {
        if (!record.incoming.destroyed) {
          // Try graceful shutdown first, then force destroy after a short timeout
          record.incoming.end();
          setTimeout(() => {
            if (record && !record.incoming.destroyed) {
              record.incoming.destroy();
            }
          }, 1000);
        }
      } catch (err) {
        console.log(`Error closing incoming socket: ${err}`);
        if (!record.incoming.destroyed) {
          record.incoming.destroy();
        }
      }
      
      try {
        if (record.outgoing && !record.outgoing.destroyed) {
          // Try graceful shutdown first, then force destroy after a short timeout
          record.outgoing.end();
          setTimeout(() => {
            if (record && record.outgoing && !record.outgoing.destroyed) {
              record.outgoing.destroy();
            }
          }, 1000);
        }
      } catch (err) {
        console.log(`Error closing outgoing socket: ${err}`);
        if (record.outgoing && !record.outgoing.destroyed) {
          record.outgoing.destroy();
        }
      }
      
      // Remove the record from the tracking map
      this.connectionRecords.delete(record.id);
      
      const remoteIP = record.incoming.remoteAddress || 'unknown';
      console.log(`Connection from ${remoteIP} terminated (${reason}). Active connections: ${this.connectionRecords.size}`);
    }
  }

  private updateActivity(record: IConnectionRecord): void {
    record.lastActivity = Date.now();
  }

  private getTargetIP(domainConfig: IDomainConfig): string {
    if (domainConfig.targetIPs && domainConfig.targetIPs.length > 0) {
      const currentIndex = this.domainTargetIndices.get(domainConfig) || 0;
      const ip = domainConfig.targetIPs[currentIndex % domainConfig.targetIPs.length];
      this.domainTargetIndices.set(domainConfig, currentIndex + 1);
      return ip;
    }
    return this.settings.targetIP!;
  }

  public async start() {
    // Define a unified connection handler for all listening ports.
    const connectionHandler = (socket: plugins.net.Socket) => {
      if (this.isShuttingDown) {
        socket.end();
        socket.destroy();
        return;
      }

      const remoteIP = socket.remoteAddress || '';
      const localPort = socket.localPort; // The port on which this connection was accepted.
      
      const connectionId = generateConnectionId();
      const connectionRecord: IConnectionRecord = {
        id: connectionId,
        incoming: socket,
        outgoing: null,
        incomingStartTime: Date.now(),
        lastActivity: Date.now(),
        connectionClosed: false
      };
      this.connectionRecords.set(connectionId, connectionRecord);
      
      console.log(`New connection from ${remoteIP} on port ${localPort}. Active connections: ${this.connectionRecords.size}`);

      let initialDataReceived = false;
      let incomingTerminationReason: string | null = null;
      let outgoingTerminationReason: string | null = null;

      // Local function for cleanupOnce
      const cleanupOnce = () => {
        this.cleanupConnection(connectionRecord);
      };
      
      // Define initiateCleanupOnce for compatibility with potential future improvements
      const initiateCleanupOnce = (reason: string = 'normal') => {
        console.log(`Connection cleanup initiated for ${remoteIP} (${reason})`);
        cleanupOnce();
      };

      // Helper to reject an incoming connection
      const rejectIncomingConnection = (reason: string, logMessage: string) => {
        console.log(logMessage);
        socket.end();
        if (incomingTerminationReason === null) {
          incomingTerminationReason = reason;
          this.incrementTerminationStat('incoming', reason);
        }
        cleanupOnce();
      };

      // 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(`Initial data timeout for ${remoteIP}`);
            socket.end();
            cleanupOnce();
          }
        }, 5000);
      } else {
        initialDataReceived = true;
      }

      socket.on('error', (err: Error) => {
        console.log(`Incoming socket error from ${remoteIP}: ${err.message}`);
      });

      const handleError = (side: 'incoming' | 'outgoing') => (err: Error) => {
        const code = (err as any).code;
        let reason = 'error';
        if (code === 'ECONNRESET') {
          reason = 'econnreset';
          console.log(`ECONNRESET on ${side} side from ${remoteIP}: ${err.message}`);
        } else {
          console.log(`Error on ${side} side from ${remoteIP}: ${err.message}`);
        }
        if (side === 'incoming' && incomingTerminationReason === null) {
          incomingTerminationReason = reason;
          this.incrementTerminationStat('incoming', reason);
        } else if (side === 'outgoing' && outgoingTerminationReason === null) {
          outgoingTerminationReason = reason;
          this.incrementTerminationStat('outgoing', reason);
        }
        initiateCleanupOnce(reason);
      };

      const handleClose = (side: 'incoming' | 'outgoing') => () => {
        console.log(`Connection closed on ${side} side from ${remoteIP}`);
        if (side === 'incoming' && incomingTerminationReason === null) {
          incomingTerminationReason = 'normal';
          this.incrementTerminationStat('incoming', 'normal');
        } else if (side === 'outgoing' && outgoingTerminationReason === null) {
          outgoingTerminationReason = 'normal';
          this.incrementTerminationStat('outgoing', 'normal');
          // Record the time when outgoing socket closed.
          connectionRecord.outgoingClosedTime = Date.now();
        }
        initiateCleanupOnce('closed_' + side);
      };

      /**
       * Sets up the connection to the target host.
       * @param serverName - The SNI hostname (unused when forcedDomain is provided).
       * @param initialChunk - Optional initial data chunk.
       * @param forcedDomain - If provided, overrides SNI/domain lookup (used for port-based routing).
       * @param overridePort - If provided, use this port for the outgoing connection.
       */
      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;
        }
        
        // If a forcedDomain is provided (port-based routing), use it; otherwise, use SNI-based lookup.
        const domainConfig = forcedDomain
          ? forcedDomain
          : (serverName ? this.settings.domainConfigs.find(config =>
              config.domains.some(d => plugins.minimatch(serverName, d))
            ) : undefined);

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

        const targetHost = domainConfig ? this.getTargetIP(domainConfig) : this.settings.targetIP!;
        const connectionOptions: plugins.net.NetConnectOpts = {
          host: targetHost,
          port: overridePort !== undefined ? overridePort : this.settings.toPort,
        };
        if (this.settings.preserveSourceIP) {
          connectionOptions.localAddress = remoteIP.replace('::ffff:', '');
        }

        // Create the target socket and immediately set up data piping
        const targetSocket = plugins.net.connect(connectionOptions);
        connectionRecord.outgoing = targetSocket;
        connectionRecord.outgoingStartTime = Date.now();
        
        // Set up the pipe immediately to ensure data flows without delay
        if (initialChunk) {
          socket.unshift(initialChunk);
        }
        
        socket.pipe(targetSocket);
        targetSocket.pipe(socket);
        
        console.log(
          `Connection established: ${remoteIP} -> ${targetHost}:${connectionOptions.port}` +
          `${serverName ? ` (SNI: ${serverName})` : forcedDomain ? ` (Port-based for domain: ${forcedDomain.domains.join(', ')})` : ''}`
        );

        // Add appropriate handlers for connection management
        socket.on('error', handleError('incoming'));
        targetSocket.on('error', handleError('outgoing'));
        socket.on('close', handleClose('incoming'));
        targetSocket.on('close', handleClose('outgoing'));
        socket.on('timeout', () => {
          console.log(`Timeout on incoming side from ${remoteIP}`);
          if (incomingTerminationReason === null) {
            incomingTerminationReason = 'timeout';
            this.incrementTerminationStat('incoming', 'timeout');
          }
          initiateCleanupOnce('timeout_incoming');
        });
        targetSocket.on('timeout', () => {
          console.log(`Timeout on outgoing side from ${remoteIP}`);
          if (outgoingTerminationReason === null) {
            outgoingTerminationReason = 'timeout';
            this.incrementTerminationStat('outgoing', 'timeout');
          }
          initiateCleanupOnce('timeout_outgoing');
        });

        // Set appropriate timeouts
        socket.setTimeout(120000);
        targetSocket.setTimeout(120000);
        
        // Update activity for both sockets
        socket.on('data', () => {
          connectionRecord.lastActivity = Date.now();
        });
        
        targetSocket.on('data', () => {
          connectionRecord.lastActivity = Date.now();
        });

        // Initialize a cleanup timer for max connection lifetime
        if (this.settings.maxConnectionLifetime) {
          connectionRecord.cleanupTimer = setTimeout(() => {
            console.log(`Connection from ${remoteIP} exceeded max lifetime (${this.settings.maxConnectionLifetime}ms), forcing cleanup.`);
            initiateCleanupOnce('max_lifetime');
          }, this.settings.maxConnectionLifetime);
        }
      };

      // --- PORT RANGE-BASED HANDLING ---
      // Only apply port-based rules if the incoming port is within one of the global port ranges.
      if (this.settings.globalPortRanges && isPortInRanges(localPort, this.settings.globalPortRanges)) {
        if (this.settings.forwardAllGlobalRanges) {
          if (this.settings.defaultAllowedIPs && !isAllowed(remoteIP, this.settings.defaultAllowedIPs)) {
            console.log(`Connection from ${remoteIP} rejected: IP ${remoteIP} not allowed in global default allowed list.`);
            socket.end();
            return;
          }
          console.log(`Port-based connection from ${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.settings.domainConfigs.find(
            domain => domain.portRanges && domain.portRanges.length > 0 && isPortInRanges(localPort, domain.portRanges)
          );
          if (forcedDomain) {
            const effectiveAllowedIPs: string[] = [
              ...forcedDomain.allowedIPs,
              ...(this.settings.defaultAllowedIPs || [])
            ];
            const effectiveBlockedIPs: string[] = [
              ...(forcedDomain.blockedIPs || []),
              ...(this.settings.defaultBlockedIPs || [])
            ];
            if (!isGlobIPAllowed(remoteIP, effectiveAllowedIPs, effectiveBlockedIPs)) {
              console.log(`Connection from ${remoteIP} rejected: IP not allowed for domain ${forcedDomain.domains.join(', ')} on port ${localPort}.`);
              socket.end();
              return;
            }
            console.log(`Port-based connection from ${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) => {
          if (initialTimeout) {
            clearTimeout(initialTimeout);
            initialTimeout = null;
          }
          
          initialDataReceived = true;
          const serverName = extractSNI(chunk) || '';
          // Lock the connection to the negotiated SNI.
          connectionRecord.lockedDomain = serverName;
          console.log(`Received connection from ${remoteIP} with SNI: ${serverName}`);
          
          // Delay adding the renegotiation listener until the next tick,
          // so the initial ClientHello is not reprocessed.
          setImmediate(() => {
            socket.on('data', (renegChunk: Buffer) => {
              if (renegChunk.length > 0 && renegChunk.readUInt8(0) === 22) {
                try {
                  // Try to extract SNI from potential renegotiation
                  const newSNI = extractSNI(renegChunk);
                  if (newSNI && newSNI !== connectionRecord.lockedDomain) {
                    console.log(`Rehandshake detected with different SNI: ${newSNI} vs locked ${connectionRecord.lockedDomain}. Terminating connection.`);
                    initiateCleanupOnce('sni_mismatch');
                  } else if (newSNI) {
                    console.log(`Rehandshake detected with same SNI: ${newSNI}. Allowing.`);
                  }
                } catch (err) {
                  console.log(`Error processing potential renegotiation: ${err}. Allowing connection to continue.`);
                }
              }
            });
          });
          
          setupConnection(serverName, chunk);
        });
      } else {
        initialDataReceived = true;
        if (!this.settings.defaultAllowedIPs || this.settings.defaultAllowedIPs.length === 0 || !isAllowed(remoteIP, this.settings.defaultAllowedIPs)) {
          return rejectIncomingConnection('rejected', `Connection rejected: IP ${remoteIP} not allowed for non-SNI connection`);
        }
        setupConnection('');
      }
    };

    // --- SETUP LISTENERS ---
    // Determine which ports to listen on.
    const listeningPorts = new Set<number>();
    if (this.settings.globalPortRanges && this.settings.globalPortRanges.length > 0) {
      // Listen on every port defined by the global ranges.
      for (const range of this.settings.globalPortRanges) {
        for (let port = range.from; port <= range.to; port++) {
          listeningPorts.add(port);
        }
      }
      // Also ensure the default fromPort is listened to if it isn't already in the ranges.
      listeningPorts.add(this.settings.fromPort);
    } else {
      listeningPorts.add(this.settings.fromPort);
    }

    // Create a server for each port.
    for (const port of listeningPorts) {
      const server = plugins.net
        .createServer(connectionHandler)
        .on('error', (err: Error) => {
          console.log(`Server Error on port ${port}: ${err.message}`);
        });
      server.listen(port, () => {
        console.log(`PortProxy -> OK: Now listening on port ${port}${this.settings.sniEnabled ? ' (SNI passthrough enabled)' : ''}`);
      });
      this.netServers.push(server);
    }

    // Log active connection count, longest running durations, and run parity checks every 10 seconds.
    this.connectionLogger = setInterval(() => {
      if (this.isShuttingDown) return;
      
      const now = Date.now();
      let maxIncoming = 0;
      let maxOutgoing = 0;
      
      // Create a copy of the keys to avoid modification during iteration
      const connectionIds = [...this.connectionRecords.keys()];
      
      for (const id of connectionIds) {
        const record = this.connectionRecords.get(id);
        if (!record) continue;
        
        maxIncoming = Math.max(maxIncoming, now - record.incomingStartTime);
        if (record.outgoingStartTime) {
          maxOutgoing = Math.max(maxOutgoing, now - record.outgoingStartTime);
        }
        
        // Parity check: if outgoing socket closed and incoming remains active
        if (record.outgoingClosedTime && 
            !record.incoming.destroyed && 
            !record.connectionClosed && 
            (now - record.outgoingClosedTime > 30000)) {
          const remoteIP = record.incoming.remoteAddress || 'unknown';
          console.log(`Parity check: Incoming socket for ${remoteIP} still active ${plugins.prettyMs(now - record.outgoingClosedTime)} after outgoing closed.`);
          this.cleanupConnection(record, 'parity_check');
        }
        
        // Inactivity check
        const inactivityTime = now - record.lastActivity;
        if (inactivityTime > 180000 && // 3 minutes
            !record.connectionClosed) {
          const remoteIP = record.incoming.remoteAddress || 'unknown';
          console.log(`Inactivity check: No activity on connection from ${remoteIP} for ${plugins.prettyMs(inactivityTime)}.`);
          this.cleanupConnection(record, 'inactivity');
        }
      }
      
      console.log(
        `(Interval Log) Active connections: ${this.connectionRecords.size}. ` +
        `Longest running incoming: ${plugins.prettyMs(maxIncoming)}, outgoing: ${plugins.prettyMs(maxOutgoing)}. ` +
        `Termination stats (incoming): ${JSON.stringify(this.terminationStats.incoming)}, ` +
        `(outgoing): ${JSON.stringify(this.terminationStats.outgoing)}`
      );
    }, 10000);
  }

  public async stop() {
    console.log("PortProxy shutting down...");
    this.isShuttingDown = true;
    
    // Stop accepting new connections
    const closeServerPromises: Promise<void>[] = this.netServers.map(
      server =>
        new Promise<void>((resolve) => {
          server.close(() => resolve());
        })
    );
    
    // Stop the connection logger
    if (this.connectionLogger) {
      clearInterval(this.connectionLogger);
      this.connectionLogger = null;
    }

    // Wait for servers to close
    await Promise.all(closeServerPromises);
    console.log("All servers closed. Cleaning up active connections...");
    
    // Clean up active connections
    const connectionIds = [...this.connectionRecords.keys()];
    console.log(`Cleaning up ${connectionIds.length} active connections...`);
    
    for (const id of connectionIds) {
      const record = this.connectionRecords.get(id);
      if (record && !record.connectionClosed) {
        this.cleanupConnection(record, 'shutdown');
      }
    }
    
    // Wait for graceful shutdown or timeout
    const shutdownTimeout = this.settings.gracefulShutdownTimeout || 30000;
    await new Promise<void>((resolve) => {
      const checkInterval = setInterval(() => {
        if (this.connectionRecords.size === 0) {
          clearInterval(checkInterval);
          resolve();
        }
      }, 1000);
      
      // Force resolve after timeout
      setTimeout(() => {
        clearInterval(checkInterval);
        if (this.connectionRecords.size > 0) {
          console.log(`Forcing shutdown with ${this.connectionRecords.size} connections still active`);
          
          // Force destroy any remaining connections
          for (const record of this.connectionRecords.values()) {
            if (!record.incoming.destroyed) {
              record.incoming.destroy();
            }
            if (record.outgoing && !record.outgoing.destroyed) {
              record.outgoing.destroy();
            }
          }
          this.connectionRecords.clear();
        }
        resolve();
      }, shutdownTimeout);
    });
    
    console.log("PortProxy shutdown complete.");
  }
}