tests/roundtrip.test.ts

import { describe, it, expect } from 'vitest';
import { embedWatermark } from '../core/embedder.js';
import { detectWatermark, detectWatermarkMultiFrame } from '../core/detector.js';
import { getPreset } from '../core/presets.js';
import type { PresetName } from '../core/types.js';

/**
 * Crop a Y plane by removing pixels from the edges.
 */
function cropYPlane(
  yPlane: Uint8Array,
  width: number,
  height: number,
  left: number,
  top: number,
  right: number,
  bottom: number,
): { cropped: Uint8Array; croppedW: number; croppedH: number } {
  const croppedW = width - left - right;
  const croppedH = height - top - bottom;
  const cropped = new Uint8Array(croppedW * croppedH);
  for (let y = 0; y < croppedH; y++) {
    for (let x = 0; x < croppedW; x++) {
      cropped[y * croppedW + x] = yPlane[(y + top) * width + (x + left)];
    }
  }
  return { cropped, croppedW, croppedH };
}

/**
 * Generate a synthetic Y plane (gradient + noise)
 */
function generateTestFrame(width: number, height: number): Uint8Array {
  const frame = new Uint8Array(width * height);
  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      const gradient = ((x + y) / (width + height)) * 200 + 20;
      const noise = (Math.random() - 0.5) * 20;
      frame[y * width + x] = Math.max(0, Math.min(255, Math.round(gradient + noise)));
    }
  }
  return frame;
}

describe('Embed → Detect Round-trip', () => {
  const width = 512;
  const height = 512;
  const payload = new Uint8Array([0xDE, 0xAD, 0xBE, 0xEF]);
  const key = 'test-secret-key-2024';

  // Test the presets that should work with a 512x512 frame
  for (const presetName of ['light', 'moderate', 'strong'] as PresetName[]) {
    it(`should round-trip with ${presetName} preset`, () => {
      const config = getPreset(presetName);
      const frame = generateTestFrame(width, height);

      // Embed
      const embedResult = embedWatermark(frame, width, height, payload, key, config);
      expect(embedResult.psnr).toBeGreaterThan(15); // Higher presets sacrifice invisibility for robustness

      // Detect
      const detectResult = detectWatermark(embedResult.yPlane, width, height, key, config);
      expect(detectResult.detected).toBe(true);
      expect(detectResult.payload).not.toBeNull();
      expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
    });
  }

  it('should not detect watermark with wrong key', () => {
    const config = getPreset('moderate');
    const frame = generateTestFrame(width, height);

    const embedResult = embedWatermark(frame, width, height, payload, key, config);

    const detectResult = detectWatermark(embedResult.yPlane, width, height, 'wrong-key', config);
    expect(detectResult.detected).toBe(false);
  });

  it('should not detect watermark on unwatermarked frame', () => {
    const config = getPreset('moderate');
    const frame = generateTestFrame(width, height);

    const detectResult = detectWatermark(frame, width, height, key, config);
    expect(detectResult.detected).toBe(false);
  });
});

describe('Crop-resilient detection', () => {
  const width = 512;
  const height = 512;
  const payload = new Uint8Array([0xDE, 0xAD, 0xBE, 0xEF]);
  const key = 'test-secret-key-2024';

  it('should detect after arbitrary crop with multiple frames', () => {
    const config = getPreset('strong');
    // Generate multiple distinct frames (simulating video — at least 32 frames)
    const frames: Uint8Array[] = [];
    for (let i = 0; i < 32; i++) {
      const frame = generateTestFrame(width, height);
      const embedResult = embedWatermark(frame, width, height, payload, key, config);
      // Crop by arbitrary offset (breaks DWT pixel pairing)
      const { cropped } = cropYPlane(
        embedResult.yPlane, width, height, 7, 3, 5, 9
      );
      frames.push(cropped);
    }
    const croppedW = width - 7 - 5;
    const croppedH = height - 3 - 9;

    const detectResult = detectWatermarkMultiFrame(
      frames, croppedW, croppedH, key, config, { cropResilient: true }
    );
    expect(detectResult.detected).toBe(true);
    expect(detectResult.payload).not.toBeNull();
    expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
  });

  it('should detect after large crop with multiple frames', () => {
    const config = getPreset('strong');
    const frames: Uint8Array[] = [];
    for (let i = 0; i < 32; i++) {
      const frame = generateTestFrame(width, height);
      const embedResult = embedWatermark(frame, width, height, payload, key, config);
      const { cropped } = cropYPlane(
        embedResult.yPlane, width, height, 50, 50, 50, 50
      );
      frames.push(cropped);
    }
    const croppedW = width - 100;
    const croppedH = height - 100;

    const detectResult = detectWatermarkMultiFrame(
      frames, croppedW, croppedH, key, config, { cropResilient: true }
    );
    expect(detectResult.detected).toBe(true);
    expect(detectResult.payload).not.toBeNull();
    expect(Array.from(detectResult.payload!)).toEqual(Array.from(payload));
  });
});