import { Buffer } from 'buffer';
import { TlsExtensionType, TlsUtils } from '../utils/tls-utils.js';
import {
  ClientHelloParser,
  type LoggerFunction
} from './client-hello-parser.js';

/**
 * Connection tracking information
 */
export interface ConnectionInfo {
  sourceIp: string;
  sourcePort: number;
  destIp: string;
  destPort: number;
  timestamp?: number;
}

/**
 * Utilities for extracting SNI information from TLS handshakes
 */
export class SniExtraction {
  /**
   * Extracts the SNI (Server Name Indication) from a TLS ClientHello message.
   * 
   * @param buffer The buffer containing the TLS ClientHello message
   * @param logger Optional logging function
   * @returns The extracted server name or undefined if not found
   */
  public static extractSNI(buffer: Buffer, logger?: LoggerFunction): string | undefined {
    const log = logger || (() => {});
    
    try {
      // Parse the ClientHello
      const parseResult = ClientHelloParser.parseClientHello(buffer, logger);
      if (!parseResult.isValid) {
        log(`Failed to parse ClientHello: ${parseResult.error}`);
        return undefined;
      }
      
      // Check if ServerName extension was found
      if (parseResult.serverNameList && parseResult.serverNameList.length > 0) {
        // Use the first hostname (most common case)
        const serverName = parseResult.serverNameList[0];
        log(`Found SNI: ${serverName}`);
        return serverName;
      }
      
      log('No SNI extension found in ClientHello');
      return undefined;
    } catch (error) {
      log(`Error extracting SNI: ${error instanceof Error ? error.message : String(error)}`);
      return undefined;
    }
  }
  
  /**
   * Attempts to extract SNI from the PSK extension in a TLS 1.3 ClientHello.
   * 
   * In TLS 1.3, when a client attempts to resume a session, it may include
   * the server name in the PSK identity hint rather than in the SNI extension.
   * 
   * @param buffer The buffer containing the TLS ClientHello message
   * @param logger Optional logging function
   * @returns The extracted server name or undefined if not found
   */
  public static extractSNIFromPSKExtension(
    buffer: Buffer,
    logger?: LoggerFunction
  ): string | undefined {
    const log = logger || (() => {});
    
    try {
      // Ensure this is a ClientHello
      if (!TlsUtils.isClientHello(buffer)) {
        log('Not a ClientHello message');
        return undefined;
      }
      
      // Parse the ClientHello to find PSK extension
      const parseResult = ClientHelloParser.parseClientHello(buffer, logger);
      if (!parseResult.isValid || !parseResult.extensions) {
        return undefined;
      }
      
      // Find the PSK extension
      const pskExtension = parseResult.extensions.find(ext => 
        ext.type === TlsExtensionType.PRE_SHARED_KEY);
      
      if (!pskExtension) {
        log('No PSK extension found');
        return undefined;
      }
      
      // Parse the PSK extension data
      const data = pskExtension.data;
      
      // PSK extension structure:
      // 2 bytes: identities list length
      if (data.length < 2) return undefined;
      
      const identitiesLength = (data[0] << 8) + data[1];
      let pos = 2;
      
      // End of identities list
      const identitiesEnd = pos + identitiesLength;
      if (identitiesEnd > data.length) return undefined;
      
      // Process each PSK identity
      while (pos + 2 <= identitiesEnd) {
        // Identity length (2 bytes)
        if (pos + 2 > identitiesEnd) break;
        
        const identityLength = (data[pos] << 8) + data[pos + 1];
        pos += 2;
        
        if (pos + identityLength > identitiesEnd) break;
        
        // Try to extract hostname from identity
        // Chrome often embeds the hostname in the PSK identity
        // This is a heuristic as there's no standard format
        if (identityLength > 0) {
          const identity = data.slice(pos, pos + identityLength);
          
          // Skip identity bytes
          pos += identityLength;
          
          // Skip obfuscated ticket age (4 bytes)
          if (pos + 4 <= identitiesEnd) {
            pos += 4;
          } else {
            break;
          }
          
          // Try to parse the identity as UTF-8
          try {
            const identityStr = identity.toString('utf8');
            log(`PSK identity: ${identityStr}`);
            
            // Check if the identity contains hostname hints
            // Chrome often embeds the hostname in a known format
            // Try to extract using common patterns
            
            // Pattern 1: Look for domain name pattern
            const domainPattern =
              /([a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?/i;
            const domainMatch = identityStr.match(domainPattern);
            if (domainMatch && domainMatch[0]) {
              log(`Found domain in PSK identity: ${domainMatch[0]}`);
              return domainMatch[0];
            }
            
            // Pattern 2: Chrome sometimes uses a specific format with delimiters
            // This is a heuristic approach since the format isn't standardized
            const parts = identityStr.split('|');
            if (parts.length > 1) {
              for (const part of parts) {
                if (part.includes('.') && !part.includes('/')) {
                  const possibleDomain = part.trim();
                  if (/^[a-z0-9.-]+$/i.test(possibleDomain)) {
                    log(`Found possible domain in PSK delimiter format: ${possibleDomain}`);
                    return possibleDomain;
                  }
                }
              }
            }
          } catch (e) {
            log('Failed to parse PSK identity as UTF-8');
          }
        }
      }
      
      log('No hostname found in PSK extension');
      return undefined;
    } catch (error) {
      log(`Error parsing PSK: ${error instanceof Error ? error.message : String(error)}`);
      return undefined;
    }
  }
  
  /**
   * Main entry point for SNI extraction with support for fragmented messages
   * and session resumption edge cases.
   * 
   * @param buffer The buffer containing TLS data
   * @param connectionInfo Connection tracking information
   * @param logger Optional logging function
   * @param cachedSni Optional previously cached SNI value
   * @returns The extracted server name or undefined
   */
  public static extractSNIWithResumptionSupport(
    buffer: Buffer,
    connectionInfo?: ConnectionInfo,
    logger?: LoggerFunction,
    cachedSni?: string
  ): string | undefined {
    const log = logger || (() => {});
    
    // Log buffer details for debugging
    if (logger) {
      log(`Buffer size: ${buffer.length} bytes`);
      log(`Buffer starts with: ${buffer.slice(0, Math.min(10, buffer.length)).toString('hex')}`);
      
      if (buffer.length >= 5) {
        const recordType = buffer[0];
        const majorVersion = buffer[1];
        const minorVersion = buffer[2];
        const recordLength = (buffer[3] << 8) + buffer[4];
        
        log(
          `TLS Record: type=${recordType}, version=${majorVersion}.${minorVersion}, length=${recordLength}`
        );
      }
    }
    
    // Check if we need to handle fragmented packets
    let processBuffer = buffer;
    if (connectionInfo) {
      const connectionId = TlsUtils.createConnectionId(connectionInfo);
      const reassembledBuffer = ClientHelloParser.handleFragmentedClientHello(
        buffer,
        connectionId,
        logger
      );
      
      if (!reassembledBuffer) {
        log(`Waiting for more fragments on connection ${connectionId}`);
        return undefined; // Need more fragments to complete ClientHello
      }
      
      processBuffer = reassembledBuffer;
      log(`Using reassembled buffer of length ${processBuffer.length}`);
    }
    
    // First try the standard SNI extraction
    const standardSni = this.extractSNI(processBuffer, logger);
    if (standardSni) {
      log(`Found standard SNI: ${standardSni}`);
      return standardSni;
    }
    
    // Check for session resumption when standard SNI extraction fails
    if (TlsUtils.isClientHello(processBuffer)) {
      const resumptionInfo = ClientHelloParser.hasSessionResumption(processBuffer, logger);
      
      if (resumptionInfo.isResumption) {
        log(`Detected session resumption in ClientHello without standard SNI`);
        
        // Try to extract SNI from PSK extension
        const pskSni = this.extractSNIFromPSKExtension(processBuffer, logger);
        if (pskSni) {
          log(`Extracted SNI from PSK extension: ${pskSni}`);
          return pskSni;
        }
      }
    }
    
    // If cached SNI was provided, use it for application data packets
    if (cachedSni && TlsUtils.isTlsApplicationData(buffer)) {
      log(`Using provided cached SNI for application data: ${cachedSni}`);
      return cachedSni;
    }
    
    return undefined;
  }
  
  /**
   * Unified method for processing a TLS packet and extracting SNI.
   * Main entry point for SNI extraction that handles all edge cases.
   * 
   * @param buffer The buffer containing TLS data
   * @param connectionInfo Connection tracking information
   * @param logger Optional logging function
   * @param cachedSni Optional previously cached SNI value
   * @returns The extracted server name or undefined
   */
  public static processTlsPacket(
    buffer: Buffer,
    connectionInfo: ConnectionInfo,
    logger?: LoggerFunction,
    cachedSni?: string
  ): string | undefined {
    const log = logger || (() => {});
    
    // Add timestamp if not provided
    if (!connectionInfo.timestamp) {
      connectionInfo.timestamp = Date.now();
    }
    
    // Check if this is a TLS handshake or application data
    if (!TlsUtils.isTlsHandshake(buffer) && !TlsUtils.isTlsApplicationData(buffer)) {
      log('Not a TLS handshake or application data packet');
      return undefined;
    }
    
    // Create connection ID for tracking
    const connectionId = TlsUtils.createConnectionId(connectionInfo);
    log(`Processing TLS packet for connection ${connectionId}, buffer length: ${buffer.length}`);
    
    // Handle application data with cached SNI (for connection racing)
    if (TlsUtils.isTlsApplicationData(buffer)) {
      // If explicit cachedSni was provided, use it
      if (cachedSni) {
        log(`Using provided cached SNI for application data: ${cachedSni}`);
        return cachedSni;
      }
      
      log('Application data packet without cached SNI, cannot determine hostname');
      return undefined;
    }
    
    // Enhanced session resumption detection
    if (TlsUtils.isClientHello(buffer)) {
      const resumptionInfo = ClientHelloParser.hasSessionResumption(buffer, logger);
      
      if (resumptionInfo.isResumption) {
        log(`Session resumption detected in TLS packet`);
        
        // Always try standard SNI extraction first
        const standardSni = this.extractSNI(buffer, logger);
        if (standardSni) {
          log(`Found standard SNI in session resumption: ${standardSni}`);
          return standardSni;
        }
        
        // Enhanced session resumption SNI extraction
        // Try extracting from PSK identity
        const pskSni = this.extractSNIFromPSKExtension(buffer, logger);
        if (pskSni) {
          log(`Extracted SNI from PSK extension: ${pskSni}`);
          return pskSni;
        }
        
        log(`Session resumption without extractable SNI`);
      }
    }
    
    // For handshake messages, try the full extraction process
    const sni = this.extractSNIWithResumptionSupport(buffer, connectionInfo, logger);
    
    if (sni) {
      log(`Successfully extracted SNI: ${sni}`);
      return sni;
    }
    
    // If we couldn't extract an SNI, check if this is a valid ClientHello
    if (TlsUtils.isClientHello(buffer)) {
      log('Valid ClientHello detected, but no SNI extracted - might need more data');
    }
    
    return undefined;
  }
}