import * as plugins from '../plugins.js';
import * as paths from '../paths.js';
import { Email } from './classes.email.js';
import type { MtaService } from './classes.mta.js';
import { logger } from '../logger.js';
import { 
  SecurityLogger, 
  SecurityLogLevel, 
  SecurityEventType,
  IPReputationChecker, 
  ReputationThreshold 
} from '../security/index.js';

export interface ISmtpServerOptions {
  port: number;
  key: string;
  cert: string;
  hostname?: string;
}

// SMTP Session States
enum SmtpState {
  GREETING,
  AFTER_EHLO,
  MAIL_FROM,
  RCPT_TO,
  DATA,
  DATA_RECEIVING,
  FINISHED
}

// Structure to store session information
interface SmtpSession {
  state: SmtpState;
  clientHostname: string;
  mailFrom: string;
  rcptTo: string[];
  emailData: string;
  useTLS: boolean;
  connectionEnded: boolean;
}

export class SMTPServer {
  public mtaRef: MtaService;
  private smtpServerOptions: ISmtpServerOptions;
  private server: plugins.net.Server;
  private sessions: Map<plugins.net.Socket | plugins.tls.TLSSocket, SmtpSession>;
  private hostname: string;

  constructor(mtaRefArg: MtaService, optionsArg: ISmtpServerOptions) {
    console.log('SMTPServer instance is being created...');

    this.mtaRef = mtaRefArg;
    this.smtpServerOptions = optionsArg;
    this.sessions = new Map();
    this.hostname = optionsArg.hostname || 'mta.lossless.one';

    this.server = plugins.net.createServer((socket) => {
      this.handleNewConnection(socket);
    });
  }

  private async handleNewConnection(socket: plugins.net.Socket): Promise<void> {
    const clientIp = socket.remoteAddress;
    const clientPort = socket.remotePort;
    console.log(`New connection from ${clientIp}:${clientPort}`);
    
    // Initialize a new session
    this.sessions.set(socket, {
      state: SmtpState.GREETING,
      clientHostname: '',
      mailFrom: '',
      rcptTo: [],
      emailData: '',
      useTLS: false,
      connectionEnded: false
    });
    
    // Check IP reputation
    try {
      if (this.mtaRef.config.security?.checkIPReputation !== false && clientIp) {
        const reputationChecker = IPReputationChecker.getInstance();
        const reputation = await reputationChecker.checkReputation(clientIp);
        
        // Log the reputation check
        SecurityLogger.getInstance().logEvent({
          level: reputation.score < ReputationThreshold.HIGH_RISK 
            ? SecurityLogLevel.WARN 
            : SecurityLogLevel.INFO,
          type: SecurityEventType.IP_REPUTATION,
          message: `IP reputation checked for new SMTP connection: score=${reputation.score}`,
          ipAddress: clientIp,
          details: {
            clientPort,
            score: reputation.score,
            isSpam: reputation.isSpam,
            isProxy: reputation.isProxy,
            isTor: reputation.isTor,
            isVPN: reputation.isVPN,
            country: reputation.country,
            blacklists: reputation.blacklists,
            socketId: socket.remotePort.toString() + socket.remoteFamily
          }
        });
        
        // Handle high-risk IPs - add delay or reject based on score
        if (reputation.score < ReputationThreshold.HIGH_RISK) {
          // For high-risk connections, add an artificial delay to slow down potential spam
          const delayMs = Math.min(5000, Math.max(1000, (ReputationThreshold.HIGH_RISK - reputation.score) * 100));
          await new Promise(resolve => setTimeout(resolve, delayMs));
          
          if (reputation.score < 5) {
            // Very high risk - can optionally reject the connection
            if (this.mtaRef.config.security?.rejectHighRiskIPs) {
              this.sendResponse(socket, `554 Transaction failed - IP is on spam blocklist`);
              socket.destroy();
              return;
            }
          }
        }
      }
    } catch (error) {
      logger.log('error', `Error checking IP reputation: ${error.message}`, {
        ip: clientIp,
        error: error.message
      });
    }
    
    // Log the connection as a security event
    SecurityLogger.getInstance().logEvent({
      level: SecurityLogLevel.INFO,
      type: SecurityEventType.CONNECTION,
      message: `New SMTP connection established`,
      ipAddress: clientIp,
      details: {
        clientPort,
        socketId: socket.remotePort.toString() + socket.remoteFamily
      }
    });

    // Send greeting
    this.sendResponse(socket, `220 ${this.hostname} ESMTP Service Ready`);

    socket.on('data', (data) => {
      this.processData(socket, data);
    });

    socket.on('end', () => {
      const clientIp = socket.remoteAddress;
      const clientPort = socket.remotePort;
      console.log(`Connection ended from ${clientIp}:${clientPort}`);
      
      const session = this.sessions.get(socket);
      if (session) {
        session.connectionEnded = true;
        
        // Log connection end as security event
        SecurityLogger.getInstance().logEvent({
          level: SecurityLogLevel.INFO,
          type: SecurityEventType.CONNECTION,
          message: `SMTP connection ended normally`,
          ipAddress: clientIp,
          details: {
            clientPort,
            state: SmtpState[session.state],
            from: session.mailFrom || 'not set'
          }
        });
      }
    });

    socket.on('error', (err) => {
      const clientIp = socket.remoteAddress;
      const clientPort = socket.remotePort;
      console.error(`Socket error: ${err.message}`);
      
      // Log connection error as security event
      SecurityLogger.getInstance().logEvent({
        level: SecurityLogLevel.WARN,
        type: SecurityEventType.CONNECTION,
        message: `SMTP connection error`,
        ipAddress: clientIp,
        details: {
          clientPort,
          error: err.message,
          errorCode: (err as any).code,
          from: this.sessions.get(socket)?.mailFrom || 'not set'
        }
      });
      
      this.sessions.delete(socket);
      socket.destroy();
    });

    socket.on('close', () => {
      const clientIp = socket.remoteAddress;
      const clientPort = socket.remotePort;
      console.log(`Connection closed from ${clientIp}:${clientPort}`);
      
      // Log connection closure as security event
      SecurityLogger.getInstance().logEvent({
        level: SecurityLogLevel.INFO,
        type: SecurityEventType.CONNECTION,
        message: `SMTP connection closed`,
        ipAddress: clientIp,
        details: {
          clientPort,
          sessionEnded: this.sessions.get(socket)?.connectionEnded || false
        }
      });
      
      this.sessions.delete(socket);
    });
  }

  private sendResponse(socket: plugins.net.Socket | plugins.tls.TLSSocket, response: string): void {
    try {
      socket.write(`${response}\r\n`);
      console.log(`→ ${response}`);
    } catch (error) {
      console.error(`Error sending response: ${error.message}`);
      socket.destroy();
    }
  }

  private processData(socket: plugins.net.Socket | plugins.tls.TLSSocket, data: Buffer): void {
    const session = this.sessions.get(socket);
    if (!session) {
      console.error('No session found for socket. Closing connection.');
      socket.destroy();
      return;
    }

    // If we're in DATA_RECEIVING state, handle differently
    if (session.state === SmtpState.DATA_RECEIVING) {
      // Call async method but don't return the promise
      this.processEmailData(socket, data.toString()).catch(err => {
        console.error(`Error processing email data: ${err.message}`);
      });
      return;
    }

    // Process normal SMTP commands
    const lines = data.toString().split('\r\n').filter(line => line.length > 0);
    for (const line of lines) {
      console.log(`← ${line}`);
      this.processCommand(socket, line);
    }
  }

  private processCommand(socket: plugins.net.Socket | plugins.tls.TLSSocket, commandLine: string): void {
    const session = this.sessions.get(socket);
    if (!session || session.connectionEnded) return;

    const [command, ...args] = commandLine.split(' ');
    const upperCommand = command.toUpperCase();

    switch (upperCommand) {
      case 'EHLO':
      case 'HELO':
        this.handleEhlo(socket, args.join(' '));
        break;
      case 'STARTTLS':
        this.handleStartTls(socket);
        break;
      case 'MAIL':
        this.handleMailFrom(socket, args.join(' '));
        break;
      case 'RCPT':
        this.handleRcptTo(socket, args.join(' '));
        break;
      case 'DATA':
        this.handleData(socket);
        break;
      case 'RSET':
        this.handleRset(socket);
        break;
      case 'QUIT':
        this.handleQuit(socket);
        break;
      case 'NOOP':
        this.sendResponse(socket, '250 OK');
        break;
      default:
        this.sendResponse(socket, '502 Command not implemented');
    }
  }

  private handleEhlo(socket: plugins.net.Socket | plugins.tls.TLSSocket, clientHostname: string): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    if (!clientHostname) {
      this.sendResponse(socket, '501 Syntax error in parameters or arguments');
      return;
    }

    session.clientHostname = clientHostname;
    session.state = SmtpState.AFTER_EHLO;

    // List available extensions
    this.sendResponse(socket, `250-${this.hostname} Hello ${clientHostname}`);
    this.sendResponse(socket, '250-SIZE 10485760'); // 10MB max
    this.sendResponse(socket, '250-8BITMIME');
    
    // Only offer STARTTLS if we haven't already established it
    if (!session.useTLS) {
      this.sendResponse(socket, '250-STARTTLS');
    }
    
    this.sendResponse(socket, '250 HELP');
  }

  private handleStartTls(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    if (session.state !== SmtpState.AFTER_EHLO) {
      this.sendResponse(socket, '503 Bad sequence of commands');
      return;
    }

    if (session.useTLS) {
      this.sendResponse(socket, '503 TLS already active');
      return;
    }

    this.sendResponse(socket, '220 Ready to start TLS');
    this.startTLS(socket);
  }

  private handleMailFrom(socket: plugins.net.Socket | plugins.tls.TLSSocket, args: string): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    if (session.state !== SmtpState.AFTER_EHLO) {
      this.sendResponse(socket, '503 Bad sequence of commands');
      return;
    }

    // Extract email from MAIL FROM:<user@example.com>
    const emailMatch = args.match(/FROM:<([^>]*)>/i);
    if (!emailMatch) {
      this.sendResponse(socket, '501 Syntax error in parameters or arguments');
      return;
    }

    const email = emailMatch[1];
    if (!this.isValidEmail(email)) {
      this.sendResponse(socket, '501 Invalid email address');
      return;
    }

    session.mailFrom = email;
    session.state = SmtpState.MAIL_FROM;
    this.sendResponse(socket, '250 OK');
  }

  private handleRcptTo(socket: plugins.net.Socket | plugins.tls.TLSSocket, args: string): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    if (session.state !== SmtpState.MAIL_FROM && session.state !== SmtpState.RCPT_TO) {
      this.sendResponse(socket, '503 Bad sequence of commands');
      return;
    }

    // Extract email from RCPT TO:<user@example.com>
    const emailMatch = args.match(/TO:<([^>]*)>/i);
    if (!emailMatch) {
      this.sendResponse(socket, '501 Syntax error in parameters or arguments');
      return;
    }

    const email = emailMatch[1];
    if (!this.isValidEmail(email)) {
      this.sendResponse(socket, '501 Invalid email address');
      return;
    }

    session.rcptTo.push(email);
    session.state = SmtpState.RCPT_TO;
    this.sendResponse(socket, '250 OK');
  }

  private handleData(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    if (session.state !== SmtpState.RCPT_TO) {
      this.sendResponse(socket, '503 Bad sequence of commands');
      return;
    }

    session.state = SmtpState.DATA_RECEIVING;
    session.emailData = '';
    this.sendResponse(socket, '354 End data with <CR><LF>.<CR><LF>');
  }

  private handleRset(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    // Reset the session data but keep connection information
    session.state = SmtpState.AFTER_EHLO;
    session.mailFrom = '';
    session.rcptTo = [];
    session.emailData = '';
    
    this.sendResponse(socket, '250 OK');
  }

  private handleQuit(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    this.sendResponse(socket, '221 Goodbye');
    
    // If we have collected email data, try to parse it before closing
    if (session.state === SmtpState.FINISHED && session.emailData.length > 0) {
      this.parseEmail(socket);
    }
    
    socket.end();
    this.sessions.delete(socket);
  }

  private async processEmailData(socket: plugins.net.Socket | plugins.tls.TLSSocket, data: string): Promise<void> {
    const session = this.sessions.get(socket);
    if (!session) return;

    // Check for end of data marker
    if (data.endsWith('\r\n.\r\n')) {
      // Remove the end of data marker
      const emailData = data.slice(0, -5);
      session.emailData += emailData;
      session.state = SmtpState.FINISHED;
      
      // Save and process the email
      this.saveEmail(socket);
      this.sendResponse(socket, '250 OK: Message accepted for delivery');
    } else {
      // Accumulate the data
      session.emailData += data;
    }
  }

  private saveEmail(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
    const session = this.sessions.get(socket);
    if (!session) return;

    try {
      // Ensure the directory exists
      plugins.smartfile.fs.ensureDirSync(paths.receivedEmailsDir);
      
      // Write the email to disk
      plugins.smartfile.memory.toFsSync(
        session.emailData,
        plugins.path.join(paths.receivedEmailsDir, `${Date.now()}.eml`)
      );
      
      // Parse the email
      this.parseEmail(socket);
    } catch (error) {
      console.error('Error saving email:', error);
    }
  }

  private async parseEmail(socket: plugins.net.Socket | plugins.tls.TLSSocket): Promise<void> {
    const session = this.sessions.get(socket);
    if (!session || !session.emailData) {
      console.error('No email data found for session.');
      return;
    }

    let mightBeSpam = false;
    // Prepare headers for DKIM verification results
    const customHeaders: Record<string, string> = {};

    // Authentication results
    let dkimResult = { domain: '', result: false };
    let spfResult = { domain: '', result: false };
    
    // Check security configuration
    const securityConfig = this.mtaRef.config.security || {};
    
    // 1. Verify DKIM signature if enabled
    if (securityConfig.verifyDkim !== false) {
      try {
        const verificationResult = await this.mtaRef.dkimVerifier.verify(session.emailData, {
          useCache: true,
          returnDetails: false
        });
      
        dkimResult.result = verificationResult.isValid;
        dkimResult.domain = verificationResult.domain || '';
        
        if (!verificationResult.isValid) {
          logger.log('warn', `DKIM verification failed for incoming email: ${verificationResult.errorMessage || 'Unknown error'}`);
          
          // Enhanced security logging
          SecurityLogger.getInstance().logEvent({
            level: SecurityLogLevel.WARN,
            type: SecurityEventType.DKIM,
            message: `DKIM verification failed for incoming email`,
            domain: verificationResult.domain || session.mailFrom.split('@')[1],
            details: {
              error: verificationResult.errorMessage || 'Unknown error',
              status: verificationResult.status,
              selector: verificationResult.selector,
              senderIP: socket.remoteAddress
            },
            ipAddress: socket.remoteAddress,
            success: false
          });
        } else {
          logger.log('info', `DKIM verification passed for incoming email from domain ${verificationResult.domain}`);
          
          // Enhanced security logging
          SecurityLogger.getInstance().logEvent({
            level: SecurityLogLevel.INFO,
            type: SecurityEventType.DKIM,
            message: `DKIM verification passed for incoming email`,
            domain: verificationResult.domain,
            details: {
              selector: verificationResult.selector,
              status: verificationResult.status,
              senderIP: socket.remoteAddress
            },
            ipAddress: socket.remoteAddress,
            success: true
          });
        }
        
        // Store verification results in headers
        if (verificationResult.domain) {
          customHeaders['X-DKIM-Domain'] = verificationResult.domain;
        }
        
        customHeaders['X-DKIM-Status'] = verificationResult.status || 'unknown';
        customHeaders['X-DKIM-Result'] = verificationResult.isValid ? 'pass' : 'fail';
      } catch (error) {
        logger.log('error', `Failed to verify DKIM signature: ${error.message}`);
        customHeaders['X-DKIM-Status'] = 'error';
        customHeaders['X-DKIM-Result'] = 'error';
      }
    }
    
    // 2. Verify SPF if enabled
    if (securityConfig.verifySpf !== false) {
      try {
        // Get the client IP and hostname
        const clientIp = socket.remoteAddress || '127.0.0.1';
        const clientHostname = session.clientHostname || 'localhost';
      
        // Parse the email to get envelope from
        const parsedEmail = await plugins.mailparser.simpleParser(session.emailData);
        
        // Create a temporary Email object for SPF verification
        const tempEmail = new Email({
          from: parsedEmail.from?.value[0].address || session.mailFrom,
          to: session.rcptTo[0],
          subject: "Temporary Email for SPF Verification",
          text: "This is a temporary email for SPF verification"
        });
        
        // Set envelope from for SPF verification
        tempEmail.setEnvelopeFrom(session.mailFrom);
        
        // Verify SPF
        const spfVerified = await this.mtaRef.spfVerifier.verifyAndApply(
          tempEmail,
          clientIp,
          clientHostname
        );
        
        // Update SPF result
        spfResult.result = spfVerified;
        spfResult.domain = session.mailFrom.split('@')[1] || '';
        
        // Copy SPF headers from the temp email
        if (tempEmail.headers['Received-SPF']) {
          customHeaders['Received-SPF'] = tempEmail.headers['Received-SPF'];
        }
        
        // Set spam flag if SPF fails badly
        if (tempEmail.mightBeSpam) {
          mightBeSpam = true;
        }
      } catch (error) {
        logger.log('error', `Failed to verify SPF: ${error.message}`);
        customHeaders['Received-SPF'] = `error (${error.message})`;
      }
    }
    
    // 3. Verify DMARC if enabled
    if (securityConfig.verifyDmarc !== false) {
      try {
        // Parse the email again
        const parsedEmail = await plugins.mailparser.simpleParser(session.emailData);
      
        // Create a temporary Email object for DMARC verification
        const tempEmail = new Email({
          from: parsedEmail.from?.value[0].address || session.mailFrom,
          to: session.rcptTo[0],
          subject: "Temporary Email for DMARC Verification",
          text: "This is a temporary email for DMARC verification"
        });
        
        // Verify DMARC
        const dmarcResult = await this.mtaRef.dmarcVerifier.verify(
          tempEmail,
          spfResult,
          dkimResult
        );
        
        // Apply DMARC policy
        const dmarcPassed = this.mtaRef.dmarcVerifier.applyPolicy(tempEmail, dmarcResult);
        
        // Add DMARC result to headers
        if (tempEmail.headers['X-DMARC-Result']) {
          customHeaders['X-DMARC-Result'] = tempEmail.headers['X-DMARC-Result'];
        }
        
        // Add Authentication-Results header combining all authentication results
        customHeaders['Authentication-Results'] = `${this.mtaRef.config.smtp.hostname}; ` +
          `spf=${spfResult.result ? 'pass' : 'fail'} smtp.mailfrom=${session.mailFrom}; ` +
          `dkim=${dkimResult.result ? 'pass' : 'fail'} header.d=${dkimResult.domain || 'unknown'}; ` +
          `dmarc=${dmarcPassed ? 'pass' : 'fail'} header.from=${tempEmail.getFromDomain()}`;
        
        // Set spam flag if DMARC fails
        if (tempEmail.mightBeSpam) {
          mightBeSpam = true;
        }
      } catch (error) {
        logger.log('error', `Failed to verify DMARC: ${error.message}`);
        customHeaders['X-DMARC-Result'] = `error (${error.message})`;
      }
    }

    try {
      const parsedEmail = await plugins.mailparser.simpleParser(session.emailData);
      
      const email = new Email({
        from: parsedEmail.from?.value[0].address || session.mailFrom,
        to: session.rcptTo[0], // Use the first recipient
        headers: customHeaders, // Add our custom headers with DKIM verification results
        subject: parsedEmail.subject || '',
        text: parsedEmail.html || parsedEmail.text || '',
        attachments: parsedEmail.attachments?.map((attachment) => ({
          filename: attachment.filename || '',
          content: attachment.content,
          contentType: attachment.contentType,
        })) || [],
        mightBeSpam: mightBeSpam,
      });

      console.log('Email received and parsed:', {
        from: email.from,
        to: email.to,
        subject: email.subject,
        attachments: email.attachments.length,
        mightBeSpam: email.mightBeSpam
      });
      
      // Enhanced security logging for received email
      SecurityLogger.getInstance().logEvent({
        level: mightBeSpam ? SecurityLogLevel.WARN : SecurityLogLevel.INFO,
        type: mightBeSpam ? SecurityEventType.SPAM : SecurityEventType.EMAIL_VALIDATION,
        message: `Email received and ${mightBeSpam ? 'flagged as potential spam' : 'validated successfully'}`,
        domain: email.from.split('@')[1],
        ipAddress: socket.remoteAddress,
        details: {
          from: email.from,
          subject: email.subject,
          recipientCount: email.getAllRecipients().length,
          attachmentCount: email.attachments.length,
          hasAttachments: email.hasAttachments(),
          dkimStatus: customHeaders['X-DKIM-Result'] || 'unknown'
        },
        success: !mightBeSpam
      });

      // Process or forward the email via MTA service
      try {
        await this.mtaRef.processIncomingEmail(email);
      } catch (err) {
        console.error('Error in MTA processing of incoming email:', err);
        
        // Log processing errors
        SecurityLogger.getInstance().logEvent({
          level: SecurityLogLevel.ERROR,
          type: SecurityEventType.EMAIL_VALIDATION,
          message: `Error processing incoming email`,
          domain: email.from.split('@')[1],
          ipAddress: socket.remoteAddress,
          details: {
            error: err.message,
            from: email.from,
            stack: err.stack
          },
          success: false
        });
      }
    } catch (error) {
      console.error('Error parsing email:', error);
      
      // Log parsing errors
      SecurityLogger.getInstance().logEvent({
        level: SecurityLogLevel.ERROR,
        type: SecurityEventType.EMAIL_VALIDATION,
        message: `Error parsing incoming email`,
        ipAddress: socket.remoteAddress,
        details: {
          error: error.message,
          sender: session.mailFrom,
          stack: error.stack
        },
        success: false
      });
    }
  }

  private startTLS(socket: plugins.net.Socket): void {
    try {
      const secureContext = plugins.tls.createSecureContext({
        key: this.smtpServerOptions.key,
        cert: this.smtpServerOptions.cert,
      });

      const tlsSocket = new plugins.tls.TLSSocket(socket, {
        secureContext: secureContext,
        isServer: true,
        server: this.server
      });

      const originalSession = this.sessions.get(socket);
      if (!originalSession) {
        console.error('No session found when upgrading to TLS');
        return;
      }

      // Transfer the session data to the new TLS socket
      this.sessions.set(tlsSocket, {
        ...originalSession,
        useTLS: true,
        state: SmtpState.GREETING // Reset state to require a new EHLO
      });
      
      this.sessions.delete(socket);

      tlsSocket.on('secure', () => {
        console.log('TLS negotiation successful');
      });

      tlsSocket.on('data', (data: Buffer) => {
        this.processData(tlsSocket, data);
      });

      tlsSocket.on('end', () => {
        console.log('TLS socket ended');
        const session = this.sessions.get(tlsSocket);
        if (session) {
          session.connectionEnded = true;
        }
      });

      tlsSocket.on('error', (err) => {
        console.error('TLS socket error:', err);
        this.sessions.delete(tlsSocket);
        tlsSocket.destroy();
      });

      tlsSocket.on('close', () => {
        console.log('TLS socket closed');
        this.sessions.delete(tlsSocket);
      });
    } catch (error) {
      console.error('Error upgrading connection to TLS:', error);
      socket.destroy();
    }
  }

  private isValidEmail(email: string): boolean {
    // Basic email validation - more comprehensive validation could be implemented
    const emailRegex = /^[^\s@]+@[^\s@]+\.[^\s@]+$/;
    return emailRegex.test(email);
  }

  public start(): void {
    this.server.listen(this.smtpServerOptions.port, () => {
      console.log(`SMTP Server is now running on port ${this.smtpServerOptions.port}`);
    });
  }

  public stop(): void {
    this.server.getConnections((err, count) => {
      if (err) throw err;
      console.log('Number of active connections: ', count);
    });
    
    this.server.close(() => {
      console.log('SMTP Server is now stopped');
    });
  }
}