// Copyright (c) 2026 defconxt. All rights reserved.
// Licensed under AGPL-3.0 — see LICENSE file for details.
// NERF is a trademark of defconxt.

/**
 * INCIDENT mode agent — Forensic Triage.
 *
 * Analyzes PCAP data and produces structured forensic reports.
 * Ported from execution engine/modes/incident.py.
 *
 * PCAP parsing uses Node.js Buffer API (readUInt32LE/BE, readUInt16BE, readUInt8).
 * Critical: PCAP header magic 0xa1b2c3d4 (little-endian) vs 0xd4c3b2a1
 * (big-endian) determines byte order for all subsequent reads.
 *
 * @module execution-engine/modes/incident
 */

import { readFileSync } from 'node:fs';
import { ModeAgentConfig, ToolRegistry } from '../framework.js';
import { ForensicValidator } from '../validators/forensic.js';

// ---------------------------------------------------------------------------
// PCAP constants
// ---------------------------------------------------------------------------

const PCAP_MAGIC_LE = 0xA1B2C3D4;
const PCAP_MAGIC_BE = 0xD4C3B2A1;
const PCAP_GLOBAL_HEADER_SIZE = 24;
const PCAP_PACKET_HEADER_SIZE = 16;
const ETHERNET_HEADER_SIZE = 14;
const IPV4_MIN_HEADER_SIZE = 20;
const UDP_HEADER_SIZE = 8;
const DNS_HEADER_SIZE = 12;

const ETHERTYPE_IPV4 = 0x0800;
const PROTO_ICMP = 1;
const PROTO_TCP = 6;
const PROTO_UDP = 17;

// ---------------------------------------------------------------------------
// PCAP parsing helpers (Buffer API)
// ---------------------------------------------------------------------------

/**
 * Format 6-byte MAC address as colon-separated hex string.
 * @param {Buffer} buf
 * @param {number} offset
 * @returns {string}
 */
export function _formatMac(buf, offset) {
  const bytes = [];
  for (let i = 0; i < 6; i++) {
    bytes.push(buf.readUInt8(offset + i).toString(16).padStart(2, '0'));
  }
  return bytes.join(':');
}

/**
 * Format a 32-bit integer as a dotted-quad IPv4 address.
 * @param {number} ipInt - 32-bit integer in network (big-endian) byte order
 * @returns {string}
 */
export function _formatIpv4(ipInt) {
  return [
    (ipInt >>> 24) & 0xff,
    (ipInt >>> 16) & 0xff,
    (ipInt >>> 8) & 0xff,
    ipInt & 0xff,
  ].join('.');
}

/**
 * Parse 24-byte PCAP global header.
 *
 * @param {Buffer} data
 * @returns {Object} header with magic, endian ('<' or '>'), version info, snaplen, link_type
 * @throws {Error} If magic number is invalid
 */
export function _parsePcapGlobalHeader(data) {
  if (data.length < PCAP_GLOBAL_HEADER_SIZE) {
    throw new Error(
      `PCAP data too short for global header: ${data.length} bytes (need ${PCAP_GLOBAL_HEADER_SIZE})`
    );
  }

  // Read magic to determine endianness
  const magic = data.readUInt32LE(0);
  let endian;
  if (magic === PCAP_MAGIC_LE) {
    endian = '<'; // little-endian
  } else if (magic === PCAP_MAGIC_BE) {
    endian = '>'; // big-endian
  } else {
    throw new Error(
      `Invalid PCAP magic number: 0x${magic.toString(16).padStart(8, '0')} ` +
      `(expected 0x${PCAP_MAGIC_LE.toString(16)} or 0x${PCAP_MAGIC_BE.toString(16)})`
    );
  }

  const readU32 = endian === '<' ? (b, o) => b.readUInt32LE(o) : (b, o) => b.readUInt32BE(o);
  const readU16 = endian === '<' ? (b, o) => b.readUInt16LE(o) : (b, o) => b.readUInt16BE(o);
  const readI32 = endian === '<' ? (b, o) => b.readInt32LE(o) : (b, o) => b.readInt32BE(o);

  return {
    magic: readU32(data, 0),
    endian,
    version_major: readU16(data, 4),
    version_minor: readU16(data, 6),
    thiszone: readI32(data, 8),
    sigfigs: readU32(data, 12),
    snaplen: readU32(data, 16),
    link_type: readU32(data, 20),
  };
}

/**
 * Parse a DNS QNAME from data starting at offset.
 * @param {Buffer} data
 * @param {number} offset
 * @returns {string|null}
 */
export function _parseDnsQueryName(data, offset) {
  const labels = [];
  let pos = offset;
  const maxPos = Math.min(offset + 255, data.length);

  while (pos < maxPos) {
    let labelLen;
    try {
      labelLen = data.readUInt8(pos);
    } catch {
      return null;
    }

    if (labelLen === 0) break;

    // Pointer (compression) — not expected in queries but handle it
    if ((labelLen & 0xC0) === 0xC0) return null;

    pos += 1;
    if (pos + labelLen > data.length) return null;

    try {
      labels.push(data.subarray(pos, pos + labelLen).toString('ascii'));
    } catch {
      return null;
    }

    pos += labelLen;
  }

  return labels.length > 0 ? labels.join('.') : null;
}

/**
 * Parse packet records from PCAP data starting at offset.
 *
 * Parses Ethernet → IPv4 → TCP/UDP → DNS (if port 53).
 * Truncated/malformed packets are skipped gracefully.
 *
 * @param {Buffer} data
 * @param {number} offset
 * @param {string} endian - '<' for little-endian, '>' for big-endian
 * @returns {Array<Object>} list of packet dicts
 */
export function _parsePackets(data, offset, endian) {
  const packets = [];
  let pos = offset;

  const readU32 = endian === '<' ? (b, o) => b.readUInt32LE(o) : (b, o) => b.readUInt32BE(o);

  while (pos + PCAP_PACKET_HEADER_SIZE <= data.length) {
    let tsSec, tsUsec, inclLen;
    try {
      tsSec = readU32(data, pos);
      tsUsec = readU32(data, pos + 4);
      inclLen = readU32(data, pos + 8);
      // origLen = readU32(data, pos + 12); // not used
    } catch {
      break;
    }

    const pktStart = pos + PCAP_PACKET_HEADER_SIZE;
    const pktEnd = pktStart + inclLen;

    if (pktEnd > data.length) break; // Truncated file

    const pktData = data.subarray(pktStart, pktEnd);
    pos = pktEnd;

    const packet = {
      timestamp: tsSec + tsUsec / 1_000_000,
      protocols: [],
      src_ip: null,
      dst_ip: null,
      src_port: null,
      dst_port: null,
      dns_query: null,
    };

    // --- Ethernet frame ---
    if (pktData.length < ETHERNET_HEADER_SIZE) {
      packets.push(packet);
      continue;
    }

    let ethertype;
    try {
      ethertype = pktData.readUInt16BE(12);
    } catch {
      packets.push(packet);
      continue;
    }

    packet.protocols.push('Ethernet');

    if (ethertype !== ETHERTYPE_IPV4) {
      packets.push(packet);
      continue;
    }

    // --- IPv4 header ---
    const ipStart = ETHERNET_HEADER_SIZE;
    if (pktData.length < ipStart + IPV4_MIN_HEADER_SIZE) {
      packets.push(packet);
      continue;
    }

    let ihl, protocol, srcIpRaw, dstIpRaw;
    try {
      const versionIhl = pktData.readUInt8(ipStart);
      ihl = (versionIhl & 0x0F) * 4;
      if (ihl < IPV4_MIN_HEADER_SIZE) ihl = IPV4_MIN_HEADER_SIZE;

      protocol = pktData.readUInt8(ipStart + 9);
      srcIpRaw = pktData.readUInt32BE(ipStart + 12);
      dstIpRaw = pktData.readUInt32BE(ipStart + 16);
    } catch {
      packets.push(packet);
      continue;
    }

    packet.protocols.push('IPv4');
    packet.src_ip = _formatIpv4(srcIpRaw);
    packet.dst_ip = _formatIpv4(dstIpRaw);

    const transportStart = ipStart + ihl;

    // --- TCP ---
    if (protocol === PROTO_TCP) {
      if (pktData.length >= transportStart + 4) {
        try {
          packet.src_port = pktData.readUInt16BE(transportStart);
          packet.dst_port = pktData.readUInt16BE(transportStart + 2);
          packet.protocols.push('TCP');
        } catch {
          // skip
        }
      }
    }

    // --- UDP ---
    else if (protocol === PROTO_UDP) {
      if (pktData.length >= transportStart + UDP_HEADER_SIZE) {
        try {
          packet.src_port = pktData.readUInt16BE(transportStart);
          packet.dst_port = pktData.readUInt16BE(transportStart + 2);
          packet.protocols.push('UDP');

          // DNS query parsing (port 53)
          if (packet.dst_port === 53 || packet.src_port === 53) {
            const dnsStart = transportStart + UDP_HEADER_SIZE;
            if (pktData.length >= dnsStart + DNS_HEADER_SIZE) {
              const queryName = _parseDnsQueryName(pktData, dnsStart + DNS_HEADER_SIZE);
              if (queryName) {
                packet.dns_query = queryName;
                packet.protocols.push('DNS');
              }
            }
          }
        } catch {
          // skip
        }
      }
    }

    // --- ICMP ---
    else if (protocol === PROTO_ICMP) {
      packet.protocols.push('ICMP');
    }

    packets.push(packet);
  }

  return packets;
}

/**
 * Load PCAP data from file path or context.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {Buffer}
 */
function _loadPcapData(context, toolInput) {
  const pcapPath = toolInput.pcap_path || '';

  if (pcapPath) {
    try {
      return readFileSync(pcapPath);
    } catch {
      // Fall through to context
    }
  }

  if (typeof context === 'object' && context !== null && context.pcap_data) {
    return Buffer.isBuffer(context.pcap_data) ? context.pcap_data : Buffer.from(context.pcap_data);
  }

  throw new Error(`PCAP file not found: '${pcapPath}' and no pcap_data in context`);
}

// ---------------------------------------------------------------------------
// Tool handlers
// ---------------------------------------------------------------------------

/**
 * Analyze PCAP file: packet count, protocol breakdown, IP addresses.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _incidentAnalyzePcap(context, toolInput) {
  let data;
  try {
    data = _loadPcapData(context, toolInput);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  let header;
  try {
    header = _parsePcapGlobalHeader(data);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  const packets = _parsePackets(data, PCAP_GLOBAL_HEADER_SIZE, header.endian);

  if (packets.length === 0) {
    return (
      `PCAP Analysis: ${toolInput.pcap_path || 'in-memory'}\n` +
      `Packets: 0\n` +
      `No packets found in capture.`
    );
  }

  // Duration
  const timestamps = packets.map(p => p.timestamp);
  const duration = timestamps.length > 1
    ? Math.max(...timestamps) - Math.min(...timestamps)
    : 0.0;

  // Protocol counts
  const protoCounts = { TCP: 0, UDP: 0, ICMP: 0, Other: 0 };
  for (const pkt of packets) {
    const protos = pkt.protocols;
    if (protos.includes('TCP')) protoCounts.TCP++;
    else if (protos.includes('UDP')) protoCounts.UDP++;
    else if (protos.includes('ICMP')) protoCounts.ICMP++;
    else protoCounts.Other++;
  }

  // Unique IPs
  const ips = new Set();
  for (const pkt of packets) {
    if (pkt.src_ip) ips.add(pkt.src_ip);
    if (pkt.dst_ip) ips.add(pkt.dst_ip);
  }

  const protoStr = Object.entries(protoCounts)
    .filter(([, v]) => v > 0)
    .map(([k, v]) => `${k}: ${v}`)
    .join(', ');

  return (
    `PCAP Analysis: ${toolInput.pcap_path || 'in-memory'}\n` +
    `Packets: ${packets.length}\n` +
    `Duration: ${duration.toFixed(3)}s\n` +
    `Protocols: ${protoStr}\n` +
    `Unique IPs: ${[...ips].sort().join(', ')}`
  );
}

/**
 * Extract unique connection tuples from PCAP.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _incidentExtractConnections(context, toolInput) {
  let data;
  try {
    data = _loadPcapData(context, toolInput);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  let header;
  try {
    header = _parsePcapGlobalHeader(data);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  const packets = _parsePackets(data, PCAP_GLOBAL_HEADER_SIZE, header.endian);

  const connections = new Set();
  for (const pkt of packets) {
    if (pkt.src_ip && pkt.src_port !== null) {
      const proto = pkt.protocols.includes('TCP') ? 'TCP' : 'UDP';
      connections.add(
        `${pkt.src_ip}:${pkt.src_port} → ${pkt.dst_ip}:${pkt.dst_port} (${proto})`
      );
    }
  }

  if (connections.size === 0) {
    return 'No connections found in capture.';
  }

  const lines = [...connections].sort();
  return `Connections (${lines.length}):\n` + lines.map(c => `  ${c}`).join('\n');
}

/**
 * Extract DNS query names from UDP port 53 packets.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _incidentExtractDnsQueries(context, toolInput) {
  let data;
  try {
    data = _loadPcapData(context, toolInput);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  let header;
  try {
    header = _parsePcapGlobalHeader(data);
  } catch (e) {
    return `ERROR: ${e.message}`;
  }

  const packets = _parsePackets(data, PCAP_GLOBAL_HEADER_SIZE, header.endian);

  const domains = new Set();
  for (const pkt of packets) {
    if (pkt.dns_query) domains.add(pkt.dns_query);
  }

  if (domains.size === 0) {
    return 'No DNS queries found in capture.';
  }

  const sorted = [...domains].sort();
  return `DNS Queries (${sorted.length}):\n` + sorted.map(d => `  ${d}`).join('\n');
}

/**
 * Store forensic report JSON in context.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _incidentWriteForensicReport(context, toolInput) {
  const report = toolInput.report || '';

  if (typeof context !== 'object' || context === null) {
    return 'ERROR: Context must be a dict.';
  }

  let reportData;
  if (typeof report === 'string') {
    try {
      reportData = JSON.parse(report);
    } catch {
      reportData = report;
    }
  } else {
    reportData = report;
  }

  context.report = reportData;
  const filename = toolInput.filename || 'forensic_report.json';

  return (
    `Forensic report stored as ${filename}. ` +
    `Report is available in context['report'].`
  );
}

// ---------------------------------------------------------------------------
// Tool schemas (Anthropic format)
// ---------------------------------------------------------------------------

const _ANALYZE_PCAP_SCHEMA = {
  name: 'analyze_pcap',
  description:
    'Analyze a PCAP capture file. Parses packet headers, extracts ' +
    'protocol statistics (TCP/UDP/ICMP), packet counts, capture ' +
    'duration, and unique IP addresses. Returns a formatted overview.',
  input_schema: {
    type: 'object',
    properties: {
      pcap_path: {
        type: 'string',
        description: 'Path to the PCAP file to analyze',
      },
    },
    required: ['pcap_path'],
  },
};

const _EXTRACT_CONNECTIONS_SCHEMA = {
  name: 'extract_connections',
  description:
    'Extract unique network connections from a PCAP file. Returns ' +
    'connection tuples: source IP:port → destination IP:port with protocol.',
  input_schema: {
    type: 'object',
    properties: {
      pcap_path: {
        type: 'string',
        description: 'Path to the PCAP file to analyze',
      },
    },
    required: ['pcap_path'],
  },
};

const _EXTRACT_DNS_QUERIES_SCHEMA = {
  name: 'extract_dns_queries',
  description:
    'Extract DNS query names from a PCAP file. Parses UDP port 53 ' +
    'packets and returns queried domain names.',
  input_schema: {
    type: 'object',
    properties: {
      pcap_path: {
        type: 'string',
        description: 'Path to the PCAP file to analyze',
      },
    },
    required: ['pcap_path'],
  },
};

const _WRITE_FORENSIC_REPORT_SCHEMA = {
  name: 'write_forensic_report',
  description:
    'Write the completed forensic report as JSON. Stored in context ' +
    'and validated for structural correctness.',
  input_schema: {
    type: 'object',
    properties: {
      report: {
        type: 'string',
        description:
          'Full JSON forensic report with required sections: summary, ' +
          'services, connections, findings, recommendations',
      },
      filename: {
        type: 'string',
        description: 'Filename for the report (e.g. forensic_report.json)',
      },
    },
    required: ['report'],
  },
};

// ---------------------------------------------------------------------------
// System prompt template
// ---------------------------------------------------------------------------

const _INCIDENT_SYSTEM_PROMPT = `\
You are an expert digital forensics analyst performing incident triage on a \
network packet capture. Your task is to analyze the PCAP file, identify \
services, protocols, and anomalies, and produce a structured forensic report.

## Target
PCAP File: {pcap_path}
Description: {pcap_description}

## Workflow
1. Use \`analyze_pcap\` to get an overview of the capture
2. Use \`extract_connections\` to map the network topology
3. Use \`extract_dns_queries\` to identify domain name activity
4. Synthesize findings: identify services, flag suspicious activity, reference CVEs
5. Use \`write_forensic_report\` to submit the complete JSON forensic report

## Output Format
Your forensic report MUST be valid JSON with these required sections:
- **summary**: Brief text overview of the capture and key findings
- **services**: List of identified services (port, protocol, description)
- **connections**: List of notable connections (src, dst, port, notes)
- **findings**: List of findings (severity, description, evidence, CVE references)
- **recommendations**: List of recommended actions based on findings

## CVE Format
Use standard format: CVE-YYYY-NNNNN (e.g., CVE-2024-12345).

## Important Notes
- Focus on actionable findings — not every packet needs to be reported
- Prioritize findings by severity (critical, high, medium, low)
- Include evidence for each finding
- Flag any indicators of compromise (IOCs) prominently
`;

// ---------------------------------------------------------------------------
// Output parser (fallback for text-based output)
// ---------------------------------------------------------------------------

/**
 * Extract JSON forensic report from LLM text output.
 * @param {string} text
 * @returns {Object}
 */
export function _incidentOutputParser(text) {
  if (!text || !text.trim()) {
    return { report: {}, raw_text: text };
  }

  const matches = [...text.matchAll(/```(?:json)?\s*\n(.*?)```/gs)].map(m => m[1]);
  const jsonText = matches.length > 0 ? matches.join('\n') : text;

  try {
    const parsed = JSON.parse(jsonText);
    return typeof parsed === 'object' && parsed !== null ? parsed : { report: parsed };
  } catch (e) {
    return { raw_text: text, parse_error: e.message };
  }
}

// ---------------------------------------------------------------------------
// Factory function
// ---------------------------------------------------------------------------

/**
 * Build an INCIDENT-mode ModeAgentConfig for forensic triage.
 * @returns {ModeAgentConfig}
 */
function _makeIncidentConfig() {
  const reg = new ToolRegistry();
  reg.register('analyze_pcap', _ANALYZE_PCAP_SCHEMA, _incidentAnalyzePcap);
  reg.register('extract_connections', _EXTRACT_CONNECTIONS_SCHEMA, _incidentExtractConnections);
  reg.register('extract_dns_queries', _EXTRACT_DNS_QUERIES_SCHEMA, _incidentExtractDnsQueries);
  reg.register('write_forensic_report', _WRITE_FORENSIC_REPORT_SCHEMA, _incidentWriteForensicReport);

  return new ModeAgentConfig({
    mode: 'INCIDENT',
    toolRegistry: reg,
    systemPromptTemplate: _INCIDENT_SYSTEM_PROMPT,
    validator: new ForensicValidator(),
    maxTurns: 15,
    requiresSandbox: false,
    completionCheck: null,
    outputParser: _incidentOutputParser,
    outputFormat: 'json',
  });
}

// ---------------------------------------------------------------------------
// Registration function — called by runner.initModes()
// ---------------------------------------------------------------------------

/**
 * Register INCIDENT mode with the given registerMode function.
 * @param {Function} registerMode
 */
export function register(registerMode) {
  registerMode('INCIDENT', _makeIncidentConfig);
}
