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

export interface IDomainConfig {
  domain: string;         // Glob pattern for domain
  allowedIPs: string[];   // Glob patterns for allowed IPs
  targetIP?: string;      // Optional target IP for this domain
}

export interface IPortProxySettings extends plugins.tls.TlsOptions {
  fromPort: number;
  toPort: number;
  toHost?: string; // Target host to proxy to, defaults to 'localhost'
  domains: IDomainConfig[];
  sniEnabled?: boolean;
  defaultAllowedIPs?: string[];
  preserveSourceIP?: boolean;
  maxConnectionLifetime?: number; // New option (in milliseconds) to force cleanup of long-lived connections
}

/**
 * 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 {
  incoming: plugins.net.Socket;
  outgoing: plugins.net.Socket | null;
  incomingStartTime: number;
  outgoingStartTime?: number;
  connectionClosed: boolean;
  cleanupTimer?: NodeJS.Timeout; // Timer to force cleanup after max lifetime/inactivity
}

export class PortProxy {
  netServer: plugins.net.Server;
  settings: IPortProxySettings;
  // Unified record tracking each connection pair.
  private connectionRecords: Set<IConnectionRecord> = new Set();
  private connectionLogger: NodeJS.Timeout | null = null;

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

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

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

  public async start() {
    // Helper to forcefully destroy sockets.
    const cleanUpSockets = (socketA: plugins.net.Socket, socketB?: plugins.net.Socket) => {
      if (!socketA.destroyed) socketA.destroy();
      if (socketB && !socketB.destroyed) socketB.destroy();
    };

    // Normalize an IP to include both IPv4 and IPv6 representations.
    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];
    };

    // Check if a given IP matches any of the glob patterns.
    const isAllowed = (ip: string, patterns: string[]): boolean => {
      const normalizedIPVariants = normalizeIP(ip);
      const expandedPatterns = patterns.flatMap(normalizeIP);
      return normalizedIPVariants.some(ipVariant =>
        expandedPatterns.some(pattern => plugins.minimatch(ipVariant, pattern))
      );
    };

    // Find a matching domain config based on the SNI.
    const findMatchingDomain = (serverName: string): IDomainConfig | undefined =>
      this.settings.domains.find(config => plugins.minimatch(serverName, config.domain));

    this.netServer = plugins.net.createServer((socket: plugins.net.Socket) => {
      const remoteIP = socket.remoteAddress || '';
      const connectionRecord: IConnectionRecord = {
        incoming: socket,
        outgoing: null,
        incomingStartTime: Date.now(),
        connectionClosed: false,
      };
      this.connectionRecords.add(connectionRecord);
      console.log(`New connection from ${remoteIP}. Active connections: ${this.connectionRecords.size}`);

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

      // Ensure cleanup happens only once for the entire connection record.
      const cleanupOnce = () => {
        if (!connectionRecord.connectionClosed) {
          connectionRecord.connectionClosed = true;
          if (connectionRecord.cleanupTimer) {
            clearTimeout(connectionRecord.cleanupTimer);
          }
          cleanUpSockets(connectionRecord.incoming, connectionRecord.outgoing || undefined);
          this.connectionRecords.delete(connectionRecord);
          console.log(`Connection from ${remoteIP} terminated. Active connections: ${this.connectionRecords.size}`);
        }
      };

      // 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();
      };

      socket.on('error', (err: Error) => {
        const errorMessage = initialDataReceived
          ? `(Immediate) Incoming socket error from ${remoteIP}: ${err.message}`
          : `(Premature) Incoming socket error from ${remoteIP} before data received: ${err.message}`;
        console.log(errorMessage);
      });

      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);
        }
        cleanupOnce();
      };

      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');
        }
        cleanupOnce();
      };

      const setupConnection = (serverName: string, initialChunk?: Buffer) => {
        const defaultAllowed = this.settings.defaultAllowedIPs && isAllowed(remoteIP, this.settings.defaultAllowedIPs);

        if (!defaultAllowed && serverName) {
          const domainConfig = findMatchingDomain(serverName);
          if (!domainConfig) {
            return rejectIncomingConnection('rejected', `Connection rejected: No matching domain config for ${serverName} from ${remoteIP}`);
          }
          if (!isAllowed(remoteIP, domainConfig.allowedIPs)) {
            return rejectIncomingConnection('rejected', `Connection rejected: IP ${remoteIP} not allowed for domain ${serverName}`);
          }
        } else if (!defaultAllowed && !serverName) {
          return rejectIncomingConnection('rejected', `Connection rejected: No SNI and IP ${remoteIP} not in default allowed list`);
        } else if (defaultAllowed && !serverName) {
          console.log(`Connection allowed: IP ${remoteIP} is in default allowed list`);
        }

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

        const targetSocket = plugins.net.connect(connectionOptions);
        connectionRecord.outgoing = targetSocket;
        connectionRecord.outgoingStartTime = Date.now();

        console.log(
          `Connection established: ${remoteIP} -> ${targetHost}:${this.settings.toPort}` +
          `${serverName ? ` (SNI: ${serverName})` : ''}`
        );

        if (initialChunk) {
          socket.unshift(initialChunk);
        }
        socket.setTimeout(120000);
        socket.pipe(targetSocket);
        targetSocket.pipe(socket);

        // Attach error and close handlers.
        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');
          }
          cleanupOnce();
        });
        targetSocket.on('timeout', () => {
          console.log(`Timeout on outgoing side from ${remoteIP}`);
          if (outgoingTerminationReason === null) {
            outgoingTerminationReason = 'timeout';
            this.incrementTerminationStat('outgoing', 'timeout');
          }
          cleanupOnce();
        });
        socket.on('end', handleClose('incoming'));
        targetSocket.on('end', handleClose('outgoing'));

        // If maxConnectionLifetime is set, initialize a cleanup timer that will be reset on data flow.
        if (this.settings.maxConnectionLifetime) {
          let incomingActive = false;
          let outgoingActive = false;
          const resetCleanupTimer = () => {
            if (this.settings.maxConnectionLifetime) {
              if (connectionRecord.cleanupTimer) {
                clearTimeout(connectionRecord.cleanupTimer);
              }
              connectionRecord.cleanupTimer = setTimeout(() => {
                console.log(`Connection from ${remoteIP} exceeded max lifetime with inactivity (${this.settings.maxConnectionLifetime}ms), forcing cleanup.`);
                cleanupOnce();
              }, this.settings.maxConnectionLifetime);
            }
          };

          // Start the cleanup timer.
          resetCleanupTimer();

          // Listen for data events on both sides and reset the timer when both are active.
          socket.on('data', () => {
            incomingActive = true;
            if (incomingActive && outgoingActive) {
              resetCleanupTimer();
            }
          });
          targetSocket.on('data', () => {
            outgoingActive = true;
            if (incomingActive && outgoingActive) {
              resetCleanupTimer();
            }
          });
        }
      };

      if (this.settings.sniEnabled) {
        socket.setTimeout(5000, () => {
          console.log(`Initial data timeout for ${remoteIP}`);
          socket.end();
          cleanupOnce();
        });

        socket.once('data', (chunk: Buffer) => {
          socket.setTimeout(0);
          initialDataReceived = true;
          const serverName = extractSNI(chunk) || '';
          console.log(`Received connection from ${remoteIP} with SNI: ${serverName}`);
          setupConnection(serverName, chunk);
        });
      } else {
        initialDataReceived = true;
        if (!this.settings.defaultAllowedIPs || !isAllowed(remoteIP, this.settings.defaultAllowedIPs)) {
          return rejectIncomingConnection('rejected', `Connection rejected: IP ${remoteIP} not allowed for non-SNI connection`);
        }
        setupConnection('');
      }
    })
      .on('error', (err: Error) => {
        console.log(`Server Error: ${err.message}`);
      })
      .listen(this.settings.fromPort, () => {
        console.log(
          `PortProxy -> OK: Now listening on port ${this.settings.fromPort}` +
          `${this.settings.sniEnabled ? ' (SNI passthrough enabled)' : ''}`
        );
      });

    // Every 10 seconds log active connection count and longest running durations.
    this.connectionLogger = setInterval(() => {
      const now = Date.now();
      let maxIncoming = 0;
      let maxOutgoing = 0;
      for (const record of this.connectionRecords) {
        maxIncoming = Math.max(maxIncoming, now - record.incomingStartTime);
        if (record.outgoingStartTime) {
          maxOutgoing = Math.max(maxOutgoing, now - record.outgoingStartTime);
        }
      }
      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() {
    const done = plugins.smartpromise.defer();
    this.netServer.close(() => {
      done.resolve();
    });
    if (this.connectionLogger) {
      clearInterval(this.connectionLogger);
      this.connectionLogger = null;
    }
    await done.promise;
  }
}