#!/usr/bin/env node

/**
 * Compare two Trivy vulnerability JSON reports (supporting multi-arch) and produce a set-based Markdown diff.
 */

const fs = require('node:fs');
const path = require('node:path');

const SEVERITY_ORDER = ['CRITICAL', 'HIGH', 'MEDIUM', 'LOW', 'UNKNOWN'];

const SEVERITY_BADGES = {
  CRITICAL: '🔴 CRITICAL',
  HIGH: '🟠 HIGH',
  MEDIUM: '🟡 MEDIUM',
  LOW: '🔵 LOW',
  UNKNOWN: '⚪ UNKNOWN',
};

/**
 * Loads one or more JSON reports (path, comma-separated list, or directory) and merges findings with arch tracking.
 * @param {string} arg
 * @returns {Map<string, object>}
 */
function loadReportsForArg(arg) {
  const cves = new Map();
  const scannedArchs = new Set();
  const filePaths = [];

  for (const part of arg.split(',').map((s) => s.trim())) {
    if (!part || !fs.existsSync(part)) continue;
    const stat = fs.statSync(part);
    if (stat.isDirectory()) {
      const entries = fs.readdirSync(part).sort();
      for (const entry of entries) {
        if (entry.endsWith('.json')) {
          filePaths.push(path.join(part, entry));
        }
      }
    } else if (stat.isFile()) {
      filePaths.push(part);
    }
  }

  for (const filePath of filePaths) {
    let stat;
    try {
      stat = fs.statSync(filePath);
    } catch {
      continue;
    }
    if (stat.size === 0) continue;

    const baseName = path.basename(filePath, '.json').toLowerCase();
    let arch = 'all';
    const cveMatch = baseName.match(/^cve-([a-z0-9_.-]+)$/i);
    if (cveMatch) {
      arch = cveMatch[1];
    } else if (baseName.includes('amd64')) {
      arch = 'amd64';
    } else if (baseName.includes('arm64') || baseName.includes('aarch64')) {
      arch = 'arm64';
    }
    if (arch !== 'all') {
      scannedArchs.add(arch);
    }

    let data;
    try {
      const content = fs.readFileSync(filePath, 'utf8');
      data = JSON.parse(content);
    } catch {
      continue;
    }

    if (!data || typeof data !== 'object' || !Array.isArray(data.Results)) {
      continue;
    }

    for (const result of data.Results) {
      const target = result.Target || result.Type || 'App';
      const vulns = Array.isArray(result.Vulnerabilities) ? result.Vulnerabilities : [];
      for (const v of vulns) {
        const vid = v.VulnerabilityID || '';
        const pkg = v.PkgName || '';
        if (!vid || !pkg) continue;

        const key = `${vid}:::${pkg}`;
        if (!cves.has(key)) {
          cves.set(key, { ...v, _target: target, _archs: new Set([arch]) });
        } else {
          cves.get(key)._archs.add(arch);
        }
      }
    }
  }

  cves.scannedArchs = scannedArchs;
  return cves;
}

/**
 * Formats a vulnerability record into a Markdown table row.
 * @param {object} v
 * @param {Set<string>} [allScannedArchs]
 * @returns {string}
 */
function formatVulnRow(v, allScannedArchs = new Set()) {
  const sev = v.Severity || 'UNKNOWN';
  const badge = SEVERITY_BADGES[sev] || `⚪ ${sev}`;
  const vid = v.VulnerabilityID || '';
  const url = v.PrimaryURL || '';
  const link = url ? `[${vid}](${url})` : vid;
  const pkg = v.PkgName || '';

  const archs = v._archs || new Set(['all']);
  let archBadge;
  if (
    archs.has('all') ||
    (allScannedArchs && allScannedArchs.size > 1 && archs.size === allScannedArchs.size) ||
    (archs.has('amd64') && archs.has('arm64'))
  ) {
    archBadge = '`all`';
  } else {
    archBadge = `\`${Array.from(archs).sort().join(', ')}\``;
  }

  const installed = v.InstalledVersion || '-';
  const fixed = v.FixedVersion || '-';
  let title = (v.Title || v.Description || '').replace(/\n/g, ' ').trim();
  if (title.length > 80) {
    title = title.slice(0, 77) + '...';
  }

  return `| ${badge} | ${link} | \`${pkg}\` | ${archBadge} | ${installed} | ${fixed} | ${title} |`;
}

/**
 * Compares base and branch vulnerabilities and formats a Markdown diff.
 * @param {Map<string, object>} baseCves
 * @param {Map<string, object>} branchCves
 * @param {string} baseName
 * @returns {string}
 */
function renderMarkdownDiff(baseCves, branchCves, baseName = 'main') {
  if (baseCves.size === 0) {
    return [
      `### 🔄 Vulnerability Delta (vs \`${baseName}\`)`,
      '',
      '> [!NOTE]',
      `> No baseline scan from \`${baseName}\` was available to compare against.`,
      '',
    ].join('\n');
  }

  const baseKeys = new Set(baseCves.keys());
  const branchKeys = new Set(branchCves.keys());

  const newKeys = [];
  let unchangedCount = 0;
  for (const key of branchKeys) {
    if (!baseKeys.has(key)) {
      newKeys.push(key);
    } else {
      unchangedCount++;
    }
  }

  const resolvedKeys = [];
  for (const key of baseKeys) {
    if (!branchKeys.has(key)) {
      resolvedKeys.push(key);
    }
  }

  function sortVulns(a, b) {
    const aSev = a.Severity || 'UNKNOWN';
    const bSev = b.Severity || 'UNKNOWN';
    const aIdx = SEVERITY_ORDER.indexOf(aSev) !== -1 ? SEVERITY_ORDER.indexOf(aSev) : 99;
    const bIdx = SEVERITY_ORDER.indexOf(bSev) !== -1 ? SEVERITY_ORDER.indexOf(bSev) : 99;
    if (aIdx !== bIdx) return aIdx - bIdx;
    const aPkg = a.PkgName || '';
    const bPkg = b.PkgName || '';
    return aPkg.localeCompare(bPkg);
  }

  const allScannedArchs = new Set([
    ...(baseCves.scannedArchs || []),
    ...(branchCves.scannedArchs || []),
  ]);

  const newVulns = newKeys.map((k) => branchCves.get(k)).sort(sortVulns);
  const resolvedVulns = resolvedKeys.map((k) => baseCves.get(k)).sort(sortVulns);

  const lines = [
    `### 🔄 Vulnerability Delta (vs \`${baseName}\`)`,
    '',
    '| Metric | Count |',
    '| :--- | :--- |',
    `| 🚨 **New Vulnerabilities** | ${newVulns.length > 0 ? `🚨 **${newVulns.length}**` : '0'} |`,
    `| ✅ **Resolved Vulnerabilities** | ${resolvedVulns.length > 0 ? `✅ **${resolvedVulns.length}**` : '0'} |`,
    `| ⚪ **Unchanged Vulnerabilities** | ${unchangedCount} |`,
    '',
    `#### 🚨 New Vulnerabilities Introduced (${newVulns.length})`,
  ];

  if (newVulns.length > 0) {
    lines.push('| Severity | CVE | Package | Arch | Installed | Fixed In | Title |');
    lines.push('| :--- | :--- | :--- | :--- | :--- | :--- | :--- |');
    for (const v of newVulns) {
      lines.push(formatVulnRow(v, allScannedArchs));
    }
  } else {
    lines.push('✅ No new vulnerabilities introduced in this build.');
  }
  lines.push('');

  lines.push(`#### ✅ Resolved / Removed Vulnerabilities (${resolvedVulns.length})`);
  if (resolvedVulns.length > 0) {
    lines.push('| Severity | CVE | Package | Arch | Previously Installed | Fixed In | Title |');
    lines.push('| :--- | :--- | :--- | :--- | :--- | :--- | :--- |');
    for (const v of resolvedVulns) {
      lines.push(formatVulnRow(v, allScannedArchs));
    }
  } else {
    lines.push('⚪ No existing vulnerabilities were resolved.');
  }
  lines.push('');

  return lines.join('\n');
}

function main() {
  const args = process.argv.slice(2);
  if (args.length < 2) {
    console.error('Usage: cve-diff.js <base_report(s)> <branch_report(s)> [base_name]');
    process.exit(1);
  }

  const baseArg = args[0];
  const branchArg = args[1];
  const baseName = args[2] || 'main';

  const baseCves = loadReportsForArg(baseArg);
  const branchCves = loadReportsForArg(branchArg);

  const output = renderMarkdownDiff(baseCves, branchCves, baseName);
  console.log(output);
}

if (require.main === module) {
  main();
}

module.exports = {
  loadReportsForArg,
  formatVulnRow,
  renderMarkdownDiff,
};