src/lib/fuzzer.ts

import { RingLwe } from './lattice';
import { ParamValidator, ValidationResult, Severity } from './validator';

export enum NoiseType {
  Uniform = "Uniform",
  Gaussian = "Gaussian",
  Cbd = "Cbd",
}

export interface FuzzConfig {
  iterations: number;
  n: number;
  q: number;
  noiseParam: number;
  noiseType: NoiseType;
  parallel: boolean;
  fuzzEdgeCases: boolean;
  fuzzCiphertextMalleability: boolean;
}

export interface FuzzStats {
  iterations: number;
  failures: number;
  failureRate: number;
  edgeCaseFailures: number;
  malleabilityPanics: number;

  avgKeygenUs: number;
  avgEncryptUs: number;
  avgDecryptUs: number;
  totalElapsedMs: number;

  dualAttackComplexityBits: number;
  primalAttackBits: number;
  keyRecoveryBits: number;
  sideChannelIndex: number;

  observedNoiseMax: number;
  theoreticalFailureProb: number;
  expectedFailures: number;
}

export enum Verdict {
  Pass = "Pass",
  CorrectnessFail = "CorrectnessFail",
  SecurityFail = "SecurityFail",
  CriticalFail = "CriticalFail",
}

export interface FuzzResult {
  config: FuzzConfig;
  stats: FuzzStats;
  validation: ValidationResult;
  verdict: Verdict;
  recommendations: string[];
}

export class FuzzEngine {
  private config: FuzzConfig;

  constructor(config: FuzzConfig) {
    this.config = config;
  }

  public async run(): Promise<FuzzResult> {
    const start = performance.now();

    // 1. Parameter Validation
    let validation: ValidationResult;
    switch (this.config.noiseType) {
      case NoiseType.Cbd:
        validation = ParamValidator.validateKyber(this.config.n, this.config.q, this.config.noiseParam);
        break;
      case NoiseType.Gaussian:
      default:
        validation = ParamValidator.validateGaussian(this.config.n, this.config.q, this.config.noiseParam);
        break;
    }

    // 2. Correctness Trials
    const trials = await this.runCorrectnessTrials();
    
    // 3. Edge Cases
    const edgeCaseFailures = this.config.fuzzEdgeCases ? await this.runEdgeCases() : 0;
    
    // 4. Malleability
    const malleabilityPanics = this.config.fuzzCiphertextMalleability ? await this.runMalleability() : 0;

    const failures = trials.filter(t => !t.success).length;
    const failureRate = failures / (trials.length || 1);

    const avgKeygenUs = trials.reduce((acc, t) => acc + t.keygenUs, 0) / (trials.length || 1);
    const avgEncryptUs = trials.reduce((acc, t) => acc + t.encryptUs, 0) / (trials.length || 1);
    const avgDecryptUs = trials.reduce((acc, t) => acc + t.decryptUs, 0) / (trials.length || 1);

    const totalElapsedMs = performance.now() - start;

    // 5. Security Estimation
    const security = this.estimateSecurity();

    const stats: FuzzStats = {
      iterations: trials.length,
      failures,
      failureRate,
      edgeCaseFailures,
      malleabilityPanics,
      avgKeygenUs,
      avgEncryptUs,
      avgDecryptUs,
      totalElapsedMs,
      dualAttackComplexityBits: security.dual,
      primalAttackBits: security.primal,
      keyRecoveryBits: security.keyRec,
      sideChannelIndex: security.sc,
      observedNoiseMax: 0, 
      theoreticalFailureProb: validation.failureProb,
      expectedFailures: validation.expectedFailuresIn500 * (this.config.iterations / 500.0),
    };

    const verdict = this.determineVerdict(stats, validation);
    const recommendations = this.buildRecommendations(stats, validation);

    return {
      config: this.config,
      stats,
      validation,
      verdict,
      recommendations,
    };
  }

  private async runCorrectnessTrials() {
    const trials = [];
    const lwe = new RingLwe(
        this.config.n, 
        this.config.q, 
        this.config.noiseParam, 
        this.config.noiseType === NoiseType.Cbd
    );

    for (let i = 0; i < this.config.iterations; i++) {
        const t0 = performance.now();
        const { pk, sk } = lwe.keygen();
        const keygenUs = (performance.now() - t0) * 1000;

        const ptxt = Math.random() > 0.5 ? 1 : 0;
        const t1 = performance.now();
        const ct = lwe.encrypt(pk, ptxt as 0 | 1);
        const encryptUs = (performance.now() - t1) * 1000;

        const t2 = performance.now();
        const dtxt = lwe.decrypt(sk, ct);
        const decryptUs = (performance.now() - t2) * 1000;

        trials.push({
            success: ptxt === dtxt,
            keygenUs,
            encryptUs,
            decryptUs
        });

        // Yield occasionally to prevent UI freezing
        if (i > 0 && i % 100 === 0) {
            await new Promise(r => setTimeout(r, 0));
        }
    }
    return trials;
  }

  private async runEdgeCases(): Promise<number> {
    const lwe = new RingLwe(
        this.config.n, 
        this.config.q, 
        this.config.noiseParam, 
        this.config.noiseType === NoiseType.Cbd
    );
    let failures = 0;

    // Zero messages
    for (let i = 0; i < 10; i++) {
        const { pk, sk } = lwe.keygen();
        const ct = lwe.encrypt(pk, 0);
        if (lwe.decrypt(sk, ct) !== 0) failures++;
    }

    // One messages
    for (let i = 0; i < 10; i++) {
        const { pk, sk } = lwe.keygen();
        const ct = lwe.encrypt(pk, 1);
        if (lwe.decrypt(sk, ct) !== 1) failures++;
    }

    return failures;
  }

  private async runMalleability(): Promise<number> {
    const lwe = new RingLwe(
        this.config.n, 
        this.config.q, 
        this.config.noiseParam, 
        this.config.noiseType === NoiseType.Cbd
    );
    const { pk, sk } = lwe.keygen();
    const ct = lwe.encrypt(pk, 1);
    
    // Corrupt v component
    ct.v[0] = (ct.v[0] + 1) % this.config.q;
    
    try {
        lwe.decrypt(sk, ct);
        return 0; // Decryption should not throw even on corrupted data
    } catch {
        return 1;
    }
  }

  private estimateSecurity() {
    const n = this.config.n;
    const q = this.config.q;
    const sigma = this.config.noiseParam;

    const ratio = Math.log2(q / Math.max(sigma, 0.1));
    const dual = 0.265 * n * ratio;
    const primal = dual * 1.05;
    const keyRec = (n * Math.log2(q)) / (Math.max(sigma, 0.1) * 8.0);
    const sc = Math.max(0.1, Math.min(15.0, sigma * 2.0 + n / 200.0));

    return { dual, primal, keyRec, sc };
  }

  private determineVerdict(stats: FuzzStats, val: ValidationResult): Verdict {
    const correctnessFail = stats.failures > 0 || stats.edgeCaseFailures > 0;
    const securityFail = stats.dualAttackComplexityBits < 80.0 || !val.passes;

    if (correctnessFail && securityFail) return Verdict.CriticalFail;
    if (correctnessFail) return Verdict.CorrectnessFail;
    if (securityFail) return Verdict.SecurityFail;
    return Verdict.Pass;
  }

  private buildRecommendations(stats: FuzzStats, val: ValidationResult): string[] {
    const recs = [];
    if (stats.failures > 0) {
        recs.push(`⚠️ ${stats.failures} / ${stats.iterations} failures. Noise ${val.noiseBudgetRms.toFixed(1)} ≥ q/4. Reduce β or increase q.`);
    }
    if (stats.dualAttackComplexityBits < 80.0) {
        recs.push(`🔴 Security critically low (${stats.dualAttackComplexityBits.toFixed(0)} bits).`);
    }
    if (stats.failures === 0 && val.passes) {
        recs.push(`✅ All trials passed. Parameters look sound.`);
    }
    return recs;
  }
}