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

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

export interface IProxySettings extends plugins.tls.TlsOptions {
  // Port configuration
  fromPort: number;
  toPort: number;
  toHost?: string;        // Target host to proxy to, defaults to 'localhost'

  // Domain and security settings
  domains: IDomainConfig[];
  sniEnabled?: boolean;
  defaultAllowedIPs?: string[]; // Optional default IP patterns if no matching domain found
  preserveSourceIP?: boolean;   // Whether to preserve the client's source IP when proxying
}

/**
 * Extract SNI (Server Name Indication) from a TLS ClientHello packet.
 * Returns the server name if found, or undefined.
 */
function extractSNI(buffer: Buffer): string | undefined {
  let offset = 0;
  // We need at least 5 bytes for the record header.
  if (buffer.length < 5) {
    return undefined;
  }

  // TLS record header
  const recordType = buffer.readUInt8(0);
  if (recordType !== 22) { // 22 = handshake
    return undefined;
  }
  // Read record length
  const recordLength = buffer.readUInt16BE(3);
  if (buffer.length < 5 + recordLength) {
    // Not all data arrived yet; in production you might need to accumulate more data.
    return undefined;
  }

  offset = 5;
  // Handshake message type should be 1 for ClientHello.
  const handshakeType = buffer.readUInt8(offset);
  if (handshakeType !== 1) {
    return undefined;
  }
  // Skip handshake header (1 byte type + 3 bytes length)
  offset += 4;

  // Skip client version (2 bytes) and random (32 bytes)
  offset += 2 + 32;

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

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

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

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

  // Iterate over extensions
  while (offset + 4 <= extensionsEnd) {
    const extensionType = buffer.readUInt16BE(offset);
    const extensionLength = buffer.readUInt16BE(offset + 2);
    offset += 4;

    // Check for SNI extension (type 0)
    if (extensionType === 0x0000) {
      // SNI extension: first 2 bytes are the SNI list length.
      if (offset + 2 > buffer.length) {
        return undefined;
      }
      const sniListLength = buffer.readUInt16BE(offset);
      offset += 2;
      const sniListEnd = offset + sniListLength;
      // Loop through the list; typically there is one entry.
      while (offset + 3 < sniListEnd) {
        const nameType = buffer.readUInt8(offset);
        offset++;
        const nameLen = buffer.readUInt16BE(offset);
        offset += 2;
        if (nameType === 0) { // host_name
          if (offset + nameLen > buffer.length) {
            return undefined;
          }
          const serverName = buffer.toString('utf8', offset, offset + nameLen);
          return serverName;
        }
        offset += nameLen;
      }
      break;
    } else {
      offset += extensionLength;
    }
  }
  return undefined;
}

export class PortProxy {
  netServer: plugins.net.Server;
  settings: IProxySettings;
  // Track active incoming connections
  private activeConnections: Set<plugins.net.Socket> = new Set();
  // Record start times for incoming connections
  private incomingConnectionTimes: Map<plugins.net.Socket, number> = new Map();
  // Record start times for outgoing connections
  private outgoingConnectionTimes: Map<plugins.net.Socket, number> = new Map();
  private connectionLogger: NodeJS.Timeout | null = null;

  // Overall termination statistics
  private terminationStats: {
    incoming: Record<string, number>;
    outgoing: Record<string, number>;
  } = {
    incoming: {},
    outgoing: {},
  };

  constructor(settings: IProxySettings) {
    this.settings = {
      ...settings,
      toHost: settings.toHost || 'localhost'
    };
  }

  // Helper to update termination stats.
  private incrementTerminationStat(side: 'incoming' | 'outgoing', reason: string): void {
    if (!this.terminationStats[side][reason]) {
      this.terminationStats[side][reason] = 1;
    } else {
      this.terminationStats[side][reason]++;
    }
  }

  public async start() {
    // Adjusted cleanUpSockets to allow an optional outgoing socket.
    const cleanUpSockets = (from: plugins.net.Socket, to?: plugins.net.Socket) => {
      from.end();
      from.removeAllListeners();
      from.unpipe();
      from.destroy();
      if (to) {
        to.end();
        to.removeAllListeners();
        to.unpipe();
        to.destroy();
      }
    };

    const normalizeIP = (ip: string): string[] => {
      // Handle IPv4-mapped IPv6 addresses
      if (ip.startsWith('::ffff:')) {
        const ipv4 = ip.slice(7); // Remove '::ffff:' prefix
        return [ip, ipv4];
      }
      // Handle IPv4 addresses by adding IPv4-mapped IPv6 variant
      if (/^\d{1,3}(\.\d{1,3}){3}$/.test(ip)) {
        return [ip, `::ffff:${ip}`];
      }
      return [ip];
    };

    const isAllowed = (value: string, patterns: string[]): boolean => {
      // Expand patterns to include both IPv4 and IPv6 variants
      const expandedPatterns = patterns.flatMap(normalizeIP);
      // Check if any variant of the IP matches any expanded pattern
      return normalizeIP(value).some(ip =>
        expandedPatterns.some(pattern => plugins.minimatch(ip, pattern))
      );
    };

    const findMatchingDomain = (serverName: string): IDomainConfig | undefined => {
      return this.settings.domains.find(config => plugins.minimatch(serverName, config.domain));
    };

    // Create a plain net server for TLS passthrough.
    this.netServer = plugins.net.createServer((socket: plugins.net.Socket) => {
      const remoteIP = socket.remoteAddress || '';

      // Record start time for the incoming connection.
      this.activeConnections.add(socket);
      this.incomingConnectionTimes.set(socket, Date.now());
      console.log(`New connection from ${remoteIP}. Active connections: ${this.activeConnections.size}`);

      // Flag to detect if we've received the first data chunk.
      let initialDataReceived = false;

      // Local termination reason trackers for each side.
      let incomingTermReason: string | null = null;
      let outgoingTermReason: string | null = null;

      // Immediately attach an error handler to catch early errors.
      socket.on('error', (err: Error) => {
        if (!initialDataReceived) {
          console.log(`(Premature) Incoming socket error from ${remoteIP} before data received: ${err.message}`);
        } else {
          console.log(`(Immediate) Incoming socket error from ${remoteIP}: ${err.message}`);
        }
      });

      // Ensure cleanup happens only once.
      let connectionClosed = false;
      const cleanupOnce = () => {
        if (!connectionClosed) {
          connectionClosed = true;
          cleanUpSockets(socket, to || undefined);
          this.incomingConnectionTimes.delete(socket);
          if (to) {
            this.outgoingConnectionTimes.delete(to);
          }
          if (this.activeConnections.has(socket)) {
            this.activeConnections.delete(socket);
            console.log(`Connection from ${remoteIP} terminated. Active connections: ${this.activeConnections.size}`);
          }
        }
      };

      // Outgoing connection placeholder.
      let to: plugins.net.Socket | null = null;

      // Handle errors by recording termination reason and cleaning up.
      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' && incomingTermReason === null) {
          incomingTermReason = reason;
          this.incrementTerminationStat('incoming', reason);
        } else if (side === 'outgoing' && outgoingTermReason === null) {
          outgoingTermReason = reason;
          this.incrementTerminationStat('outgoing', reason);
        }
        cleanupOnce();
      };

      // Handle close events. If no termination reason was recorded, mark as "normal".
      const handleClose = (side: 'incoming' | 'outgoing') => () => {
        console.log(`Connection closed on ${side} side from ${remoteIP}`);
        if (side === 'incoming' && incomingTermReason === null) {
          incomingTermReason = 'normal';
          this.incrementTerminationStat('incoming', 'normal');
        } else if (side === 'outgoing' && outgoingTermReason === null) {
          outgoingTermReason = 'normal';
          this.incrementTerminationStat('outgoing', 'normal');
        }
        cleanupOnce();
      };

      // Setup connection, optionally accepting the initial data chunk.
      const setupConnection = (serverName: string, initialChunk?: Buffer) => {
        // Check if the IP is allowed by default.
        const isDefaultAllowed = this.settings.defaultAllowedIPs && isAllowed(remoteIP, this.settings.defaultAllowedIPs);
        if (!isDefaultAllowed && serverName) {
          const domainConfig = findMatchingDomain(serverName);
          if (!domainConfig) {
            console.log(`Connection rejected: No matching domain config for ${serverName} from ${remoteIP}`);
            socket.end();
            if (incomingTermReason === null) {
              incomingTermReason = 'rejected';
              this.incrementTerminationStat('incoming', 'rejected');
            }
            cleanupOnce();
            return;
          }
          if (!isAllowed(remoteIP, domainConfig.allowedIPs)) {
            console.log(`Connection rejected: IP ${remoteIP} not allowed for domain ${serverName}`);
            socket.end();
            if (incomingTermReason === null) {
              incomingTermReason = 'rejected';
              this.incrementTerminationStat('incoming', 'rejected');
            }
            cleanupOnce();
            return;
          }
        } else if (!isDefaultAllowed && !serverName) {
          console.log(`Connection rejected: No SNI and IP ${remoteIP} not in default allowed list`);
          socket.end();
          if (incomingTermReason === null) {
            incomingTermReason = 'rejected';
            this.incrementTerminationStat('incoming', 'rejected');
          }
          cleanupOnce();
          return;
        } else {
          console.log(`Connection allowed: IP ${remoteIP} is in default allowed list`);
        }

        // Determine target host.
        const domainConfig = serverName ? findMatchingDomain(serverName) : undefined;
        const targetHost = domainConfig?.targetIP || this.settings.toHost!;

        // Create connection options.
        const connectionOptions: plugins.net.NetConnectOpts = {
          host: targetHost,
          port: this.settings.toPort,
        };
        if (this.settings.preserveSourceIP) {
          connectionOptions.localAddress = remoteIP.replace('::ffff:', '');
        }

        // Establish outgoing connection.
        to = plugins.net.connect(connectionOptions);
        if (to) {
          this.outgoingConnectionTimes.set(to, Date.now());
        }
        console.log(`Connection established: ${remoteIP} -> ${targetHost}:${this.settings.toPort}${serverName ? ` (SNI: ${serverName})` : ''}`);
        
        // Push back the initial chunk if provided.
        if (initialChunk) {
          socket.unshift(initialChunk);
        }
        socket.setTimeout(120000);
        socket.pipe(to!);
        to!.pipe(socket);

        // Attach event handlers for both sockets.
        socket.on('error', handleError('incoming'));
        to!.on('error', handleError('outgoing'));
        socket.on('close', handleClose('incoming'));
        to!.on('close', handleClose('outgoing'));
        socket.on('timeout', () => {
          console.log(`Timeout on incoming side from ${remoteIP}`);
          if (incomingTermReason === null) {
            incomingTermReason = 'timeout';
            this.incrementTerminationStat('incoming', 'timeout');
          }
          cleanupOnce();
        });
        to!.on('timeout', () => {
          console.log(`Timeout on outgoing side from ${remoteIP}`);
          if (outgoingTermReason === null) {
            outgoingTermReason = 'timeout';
            this.incrementTerminationStat('outgoing', 'timeout');
          }
          cleanupOnce();
        });
        socket.on('end', handleClose('incoming'));
        to!.on('end', handleClose('outgoing'));
      };

      // For SNI-enabled connections, peek at the first chunk.
      if (this.settings.sniEnabled) {
        socket.once('data', (chunk: Buffer) => {
          initialDataReceived = true;
          const serverName = extractSNI(chunk) || '';
          console.log(`Received connection from ${remoteIP} with SNI: ${serverName}`);
          setupConnection(serverName, chunk);
        });
      } else {
        // For non-SNI connections, simply check defaultAllowedIPs.
        initialDataReceived = true;
        if (!this.settings.defaultAllowedIPs || !isAllowed(remoteIP, this.settings.defaultAllowedIPs)) {
          console.log(`Connection rejected: IP ${remoteIP} not allowed for non-SNI connection`);
          socket.end();
          if (incomingTermReason === null) {
            incomingTermReason = 'rejected';
            this.incrementTerminationStat('incoming', 'rejected');
          }
          cleanupOnce();
          return;
        }
        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)' : ''}`);
    });

    // Log active connection count, longest running connection durations,
    // and termination statistics every 10 seconds.
    this.connectionLogger = setInterval(() => {
      const now = Date.now();
      let maxIncoming = 0;
      for (const startTime of this.incomingConnectionTimes.values()) {
        const duration = now - startTime;
        if (duration > maxIncoming) {
          maxIncoming = duration;
        }
      }
      let maxOutgoing = 0;
      for (const startTime of this.outgoingConnectionTimes.values()) {
        const duration = now - startTime;
        if (duration > maxOutgoing) {
          maxOutgoing = duration;
        }
      }
      console.log(`(Interval Log) Active connections: ${this.activeConnections.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;
  }
}