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

/**
 * ARCHITECT mode agent — Threat Modeling.
 *
 * Performs execution engine threat modeling using STRIDE + DREAD.
 * Ported from execution engine/modes/architect.py.
 *
 * @module execution-engine/modes/architect
 */

import { ModeAgentConfig, ToolRegistry } from '../framework.js';
import { ThreatModelValidator } from '../validators/threat-model.js';

// ---------------------------------------------------------------------------
// Architecture analysis keywords for system decomposition
// ---------------------------------------------------------------------------

const _PROCESS_KEYWORDS = [
  'service', 'server', 'api', 'gateway', 'proxy', 'worker',
  'handler', 'controller', 'middleware', 'microservice',
  'application', 'app', 'frontend', 'backend', 'engine',
  'processor', 'scheduler', 'daemon', 'auth', 'authentication',
  'authorization', 'load balancer', 'reverse proxy',
];

const _DATA_STORE_KEYWORDS = [
  'database', 'db', 'postgresql', 'postgres', 'mysql', 'mariadb',
  'mongodb', 'redis', 'memcached', 'cache', 'elasticsearch',
  'sqlite', 'dynamodb', 'cassandra', 's3', 'blob storage',
  'file system', 'data warehouse', 'data lake', 'queue',
  'message broker', 'kafka', 'rabbitmq', 'sqs',
];

const _EXTERNAL_ENTITY_KEYWORDS = [
  'external api', 'third-party', 'third party', '3rd party',
  'stripe', 'twilio', 'sendgrid', 'aws', 'azure', 'gcp',
  'oauth', 'sso', 'ldap', 'saml', 'webhook', 'cdn',
  'payment', 'email service', 'sms', 'push notification',
  'analytics', 'monitoring service', 'logging service',
];

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

/**
 * Decompose a system description into architectural components.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _architectAnalyzeArchitecture(context, toolInput) {
  const description = toolInput.system_description || '';
  if (!description || !description.trim()) {
    return "ERROR: 'system_description' parameter is required and must not be empty.";
  }

  const descLower = description.toLowerCase();

  // Identify processes/services
  const components = [];
  for (const keyword of _PROCESS_KEYWORDS) {
    if (descLower.includes(keyword)) {
      const idx = descLower.indexOf(keyword);
      const start = Math.max(0, idx - 40);
      const end = Math.min(description.length, idx + keyword.length + 40);
      const snippet = description.slice(start, end).trim();
      components.push({
        name: keyword.split(' ').map(w => w.charAt(0).toUpperCase() + w.slice(1)).join(' '),
        type: 'process',
        source: snippet,
      });
    }
  }

  // Identify data stores
  const dataStores = [];
  for (const keyword of _DATA_STORE_KEYWORDS) {
    if (descLower.includes(keyword)) {
      dataStores.push({
        name: keyword.split(' ').map(w => w.charAt(0).toUpperCase() + w.slice(1)).join(' '),
        type: 'data_store',
      });
    }
  }

  // Identify external entities
  const externalEntities = [];
  for (const keyword of _EXTERNAL_ENTITY_KEYWORDS) {
    if (descLower.includes(keyword)) {
      externalEntities.push({
        name: keyword.split(' ').map(w => w.charAt(0).toUpperCase() + w.slice(1)).join(' '),
        type: 'external_entity',
      });
    }
  }

  // Infer data flows between consecutive identified components
  const allNames = [
    ...components.map(c => c.name),
    ...dataStores.map(d => d.name),
    ...externalEntities.map(e => e.name),
  ];
  const dataFlows = [];
  for (let i = 0; i < allNames.length - 1; i++) {
    dataFlows.push({
      from: allNames[i],
      to: allNames[i + 1],
      description: `Data flow from ${allNames[i]} to ${allNames[i + 1]}`,
    });
  }

  return JSON.stringify({
    components,
    data_stores: dataStores,
    external_entities: externalEntities,
    data_flows: dataFlows,
    total_identified: components.length + dataStores.length + externalEntities.length,
  }, null, 2);
}

/**
 * Validate and score threats using DREAD risk methodology.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _architectScoreThreats(context, toolInput) {
  const threats = toolInput.threats || [];
  if (!threats.length) {
    return "ERROR: 'threats' parameter is required and must not be empty.";
  }

  const dreadDimensions = [
    'Damage', 'Reproducibility', 'Exploitability',
    'Affected Users', 'Discoverability',
  ];

  const scoredThreats = [];
  let highestRisk = 'Low';
  const riskPriority = { Critical: 4, High: 3, Medium: 2, Low: 1 };

  for (let i = 0; i < threats.length; i++) {
    const threat = threats[i];
    if (typeof threat !== 'object' || threat === null) {
      return `ERROR: threats[${i}] must be a dict (got ${typeof threat}).`;
    }

    const dread = threat.dread || {};
    if (typeof dread !== 'object' || dread === null) {
      return `ERROR: threats[${i}].dread must be a dict (got ${typeof dread}).`;
    }

    // Validate each DREAD dimension
    let total = 0;
    for (const dim of dreadDimensions) {
      const val = dread[dim];
      if (val === undefined || val === null) {
        return `ERROR: threats[${i}].dread missing dimension: ${dim}`;
      }
      if (typeof val !== 'number') {
        return `ERROR: threats[${i}].dread.${dim} must be an integer (got ${typeof val})`;
      }
      const valInt = Math.floor(val);
      if (valInt < 1 || valInt > 10) {
        return `ERROR: threats[${i}].dread.${dim} must be 1-10 (got ${valInt})`;
      }
      total += valInt;
    }

    const average = total / dreadDimensions.length;

    // Assign risk level
    let riskLevel;
    if (average >= 8) riskLevel = 'Critical';
    else if (average >= 6) riskLevel = 'High';
    else if (average >= 4) riskLevel = 'Medium';
    else riskLevel = 'Low';

    if ((riskPriority[riskLevel] || 0) > (riskPriority[highestRisk] || 0)) {
      highestRisk = riskLevel;
    }

    scoredThreats.push({
      component: threat.component || 'unknown',
      threat_type: threat.threat_type || 'unknown',
      description: threat.description || '',
      dread,
      average_score: Math.round(average * 100) / 100,
      risk_level: riskLevel,
    });
  }

  return JSON.stringify({
    scored_threats: scoredThreats,
    total_threats: scoredThreats.length,
    highest_risk_level: highestRisk,
  }, null, 2);
}

/**
 * Store a structured JSON threat model report in context.
 * @param {*} context
 * @param {Object} toolInput
 * @returns {string}
 */
export function _architectWriteThreatModel(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 || 'threat_model.json';

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

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

const _ARCHITECT_ANALYZE_ARCHITECTURE_SCHEMA = {
  name: 'analyze_architecture',
  description:
    'Decompose a system description into architectural components: ' +
    'processes/services, data stores, external entities, and data flows.',
  input_schema: {
    type: 'object',
    properties: {
      system_description: {
        type: 'string',
        description: 'Natural language description of the system architecture',
      },
    },
    required: ['system_description'],
  },
};

const _ARCHITECT_SCORE_THREATS_SCHEMA = {
  name: 'score_threats',
  description:
    'Validate and score threats using DREAD risk methodology. Each threat ' +
    'needs DREAD dimension scores (1-10). Returns scored threats with risk levels.',
  input_schema: {
    type: 'object',
    properties: {
      threats: {
        type: 'array',
        description: 'List of threat objects with component, threat_type, description, dread scores',
        items: {
          type: 'object',
          properties: {
            component: { type: 'string', description: 'Component being threatened' },
            threat_type: {
              type: 'string',
              description: 'STRIDE category: Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, or Elevation of Privilege',
            },
            description: { type: 'string', description: 'Description of the threat' },
            dread: {
              type: 'object',
              description: 'DREAD scores: Damage, Reproducibility, Exploitability, Affected Users, Discoverability (each 1-10)',
            },
          },
          required: ['component', 'threat_type', 'description', 'dread'],
        },
      },
    },
    required: ['threats'],
  },
};

const _ARCHITECT_WRITE_THREAT_MODEL_SCHEMA = {
  name: 'write_threat_model',
  description:
    'Submit the completed threat model report as JSON with required sections: ' +
    'system_description, components, trust_boundaries, stride_analysis, ' +
    'dread_scores, mitigations, dfd_mermaid.',
  input_schema: {
    type: 'object',
    properties: {
      report: {
        type: 'string',
        description: 'Full JSON threat model report as a string.',
      },
      filename: {
        type: 'string',
        description: 'Filename for the report (e.g. threat_model.json)',
      },
    },
    required: ['report'],
  },
};

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

const _ARCHITECT_SYSTEM_PROMPT = `\
You are an expert security architect performing threat modeling using the \
STRIDE methodology with DREAD risk scoring. Analyze the system description \
and produce a comprehensive structured threat model.

## System Description
{system_description}

## Instructions
1. Use \`analyze_architecture\` to decompose the system into components.
2. Perform STRIDE analysis on each component and data flow.
3. Use \`score_threats\` to compute DREAD risk scores for each threat.
4. Generate a Mermaid DFD showing components, trust boundaries, and data flows.
5. Use \`write_threat_model\` to submit the final report.

## Rules
- Analyze EVERY component for ALL 6 STRIDE categories
- Score EVERY threat with ALL 5 DREAD dimensions (integers 1-10)
- Include trust boundaries between all architectural tiers
- Generate valid Mermaid DFD syntax (starts with 'graph' or 'flowchart')
- Provide specific, actionable mitigations
`;

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

/**
 * Extract JSON threat model from LLM text output.
 * @param {string} text
 * @returns {Object}
 */
export function _architectOutputParser(text) {
  if (!text || !text.trim()) {
    return { raw_text: text, parse_error: 'empty output' };
  }

  let matches = [...text.matchAll(/```json\s*\n(.*?)```/gs)].map(m => m[1]);
  if (matches.length > 0) {
    const jsonText = matches.join('\n');
    try { return JSON.parse(jsonText); } catch (e) {
      return { raw_text: text, parse_error: e.message };
    }
  }

  matches = [...text.matchAll(/```\s*\n(.*?)```/gs)].map(m => m[1]);
  if (matches.length > 0) {
    const jsonText = matches.join('\n');
    try { return JSON.parse(jsonText); } catch (e) {
      return { raw_text: text, parse_error: e.message };
    }
  }

  try { return JSON.parse(text); } catch (e) {
    return { raw_text: text, parse_error: e.message };
  }
}

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

/**
 * Build an ARCHITECT-mode ModeAgentConfig for threat modeling.
 * @returns {ModeAgentConfig}
 */
function _makeArchitectConfig() {
  const reg = new ToolRegistry();
  reg.register('analyze_architecture', _ARCHITECT_ANALYZE_ARCHITECTURE_SCHEMA, _architectAnalyzeArchitecture);
  reg.register('score_threats', _ARCHITECT_SCORE_THREATS_SCHEMA, _architectScoreThreats);
  reg.register('write_threat_model', _ARCHITECT_WRITE_THREAT_MODEL_SCHEMA, _architectWriteThreatModel);

  return new ModeAgentConfig({
    mode: 'ARCHITECT',
    toolRegistry: reg,
    systemPromptTemplate: _ARCHITECT_SYSTEM_PROMPT,
    validator: new ThreatModelValidator(),
    maxTurns: 15,
    requiresSandbox: false,
    completionCheck: null,
    outputParser: _architectOutputParser,
    outputFormat: 'json',
  });
}

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

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