/**
 * ATR Rule Scaffolder - Generates ATR rule YAML scaffolds from structured input
 * @module agent-threat-rules/rule-scaffolder
 */

import yaml from 'js-yaml';
import type {
  ATRCategory,
  ATRSeverity,
  ATRSourceType,
  ATRAction,
  ATRArrayCondition,
} from './types.js';

export interface ScaffoldInput {
  title: string;
  category: ATRCategory;
  severity?: ATRSeverity;
  attackDescription: string;
  examplePayloads: string[];
  agentSourceType?: ATRSourceType;
  owaspRefs?: string[];
  mitreRefs?: string[];
}

export interface ScaffoldResult {
  yaml: string;
  id: string;
  warnings: string[];
}

export interface ScaffoldOptions {
  author?: string;
  schemaVersion?: string;
}

const CATEGORY_TO_SOURCE_TYPE: Readonly<Record<ATRCategory, ATRSourceType>> = {
  'prompt-injection': 'llm_io',
  'tool-poisoning': 'tool_call',
  'context-exfiltration': 'context_window',
  'agent-manipulation': 'multi_agent_comm',
  'privilege-escalation': 'agent_behavior',
  'excessive-autonomy': 'agent_behavior',
  'data-poisoning': 'llm_io',
  'model-abuse': 'llm_io',
  'skill-compromise': 'skill_lifecycle',
};

const CATEGORY_TO_FIELD: Readonly<Record<ATRCategory, string>> = {
  'prompt-injection': 'user_input',
  'tool-poisoning': 'tool_response',
  'context-exfiltration': 'agent_output',
  'agent-manipulation': 'agent_message',
  'privilege-escalation': 'agent_action',
  'excessive-autonomy': 'agent_action',
  'data-poisoning': 'training_input',
  'model-abuse': 'user_input',
  'skill-compromise': 'skill_manifest',
};

const SEVERITY_TO_ACTIONS: Readonly<Record<ATRSeverity, readonly ATRAction[]>> = {
  critical: ['block_input', 'alert', 'escalate'],
  high: ['block_input', 'alert'],
  medium: ['alert', 'snapshot'],
  low: ['alert'],
  informational: ['alert'],
};

const REGEX_SPECIAL_CHARS = /[.*+?^${}()|[\]\\]/g;

function escapeRegex(str: string): string {
  return str.replace(REGEX_SPECIAL_CHARS, '\\$&');
}

/**
 * Attack pattern templates by category — reusable regex building blocks
 * that detect BEHAVIOR, not package names.
 */
export const ATTACK_PATTERN_INDICATORS: ReadonlyArray<{
  /** Regex to test if the payload contains this attack indicator */
  readonly test: RegExp;
  /** The detection regex to use in the rule */
  readonly pattern: string;
  /** Human-readable description */
  readonly description: string;
  /** Which categories this indicator applies to */
  readonly categories: readonly ATRCategory[];
}> = [
  // Shell execution patterns
  {
    test: /exec(Sync)?|spawn|child_process|shell|subprocess|popen|os\.system/i,
    pattern: '(execSync?|spawn|child_process|shell|subprocess|popen|os\\.system)\\s*\\(',
    description: 'Shell/command execution',
    categories: ['tool-poisoning', 'skill-compromise', 'privilege-escalation'],
  },
  // Dynamic shell with interpolation (RCE)
  {
    test: /exec.*\$\{|spawn.*\$\{|`.*\$\{.*`/i,
    pattern: '(exec|spawn|shell)\\s*\\(.*\\$\\{',
    description: 'Dynamic shell execution with variable interpolation',
    categories: ['tool-poisoning', 'skill-compromise'],
  },
  // Network exfiltration
  {
    test: /fetch|http|request|axios|got|node-fetch|urllib|curl|wget/i,
    pattern: '(fetch|https?://|request|axios|got|node-fetch|urllib|curl|wget)\\s*\\(?',
    description: 'Outbound network request',
    categories: ['context-exfiltration', 'tool-poisoning', 'data-poisoning'],
  },
  // Credential/secret access
  {
    test: /password|secret|token|credential|api[_\s]?key|auth|cookie/i,
    pattern: '(password|secret|token|credential|api[_ ]?key|auth_token|cookie)',
    description: 'Credential/secret access',
    categories: ['context-exfiltration', 'privilege-escalation', 'tool-poisoning'],
  },
  // Environment variable exfiltration
  {
    test: /process\.env|os\.environ|getenv|ENV\[/i,
    pattern: '(process\\.env|os\\.environ|getenv|ENV\\[)',
    description: 'Environment variable access',
    categories: ['context-exfiltration', 'tool-poisoning', 'skill-compromise'],
  },
  // eval / dynamic code execution
  {
    test: /\beval\s*\(|new\s+Function\s*\(|vm\.run/i,
    pattern: '(\\beval\\s*\\(|new\\s+Function\\s*\\(|vm\\.run)',
    description: 'Dynamic code execution',
    categories: ['tool-poisoning', 'skill-compromise'],
  },
  // Instruction override (prompt injection)
  {
    test: /ignore|disregard|forget|override|overwrite/i,
    pattern: '(override|overwrite|ignore|disregard|forget)\\s+(previous|prior|above|existing|all|any)\\s+(instructions?|rules?|constraints?|guidelines?|protocols?)',
    description: 'Instruction override attempt',
    categories: ['prompt-injection', 'agent-manipulation'],
  },
  // Role manipulation
  {
    test: /you are now|act as|pretend|new role|system prompt/i,
    pattern: '(you\\s+are\\s+now|act\\s+as\\s+(a|an|if)|pretend\\s+(to|you)|new\\s+role|system\\s+prompt)',
    description: 'Role/identity manipulation',
    categories: ['prompt-injection', 'agent-manipulation'],
  },
  // File system destructive operations
  {
    test: /rm\s+-rf|rmdir|unlink|deleteFile|fs\.rm/i,
    pattern: '(rm\\s+-rf|rmdir\\s|unlink\\s*\\(|deleteFile|fs\\.rm)',
    description: 'Destructive file system operation',
    categories: ['tool-poisoning', 'excessive-autonomy'],
  },
  // Base64/encoding evasion
  {
    test: /atob|btoa|base64|Buffer\.from.*encoding|fromCharCode/i,
    pattern: '(atob|btoa|base64|Buffer\\.from|fromCharCode)\\s*\\(',
    description: 'Encoding-based payload obfuscation',
    categories: ['tool-poisoning', 'skill-compromise'],
  },
  // Data exfiltration combo (credential + network)
  {
    test: /(password|secret|token|key).*(fetch|http|send|post|upload)/i,
    pattern: '(password|secret|token|api[_ ]?key).*(fetch|https?://|request|send|post|upload)',
    description: 'Credential access combined with network exfiltration',
    categories: ['context-exfiltration', 'tool-poisoning'],
  },
  // Download + execute combo
  {
    test: /(download|fetch|curl|wget).*(exec|eval|spawn)/i,
    pattern: '(download|fetch|curl|wget).*(exec|eval|spawn|child_process)',
    description: 'Download and execute pattern',
    categories: ['tool-poisoning', 'skill-compromise'],
  },
];

/**
 * Build detection regex from a payload string, using category-aware
 * attack pattern templates instead of naive keyword extraction.
 *
 * Priority:
 * 1. Match known attack indicators in the payload -> use behavioral regex
 * 2. Combine multiple indicators with alternation for multi-vector attacks
 * 3. Fall back to keyword extraction only for text that has no code patterns
 */
function buildRegexPattern(payload: string, category?: ATRCategory): string {
  const trimmed = payload.trim();

  // Step 1: Find all attack indicators present in this payload
  const matched = ATTACK_PATTERN_INDICATORS.filter((ind) => {
    const matchesPayload = ind.test.test(trimmed);
    const matchesCategory = !category || ind.categories.includes(category);
    return matchesPayload && matchesCategory;
  });

  // Step 2: If we found behavioral patterns, use them
  if (matched.length > 0) {
    if (matched.length === 1) {
      return `(?i)${matched[0]!.pattern}`;
    }
    // Combine multiple indicators — detect the most specific one
    // Sort by pattern length (longer = more specific) and take top 3
    const sorted = [...matched].sort((a, b) => b.pattern.length - a.pattern.length);
    const top = sorted.slice(0, 3);
    if (top.length === 1) {
      return `(?i)${top[0]!.pattern}`;
    }
    // Use alternation for multi-vector detection
    return `(?i)(${top.map((t) => t.pattern).join('|')})`;
  }

  // Step 3: Fallback — keyword extraction for text-based payloads
  // (e.g., prompt injection text without code patterns)
  const words = trimmed.split(/\s+/).filter((w) => w.length > 3);

  if (words.length === 0) {
    return `(?i).*${escapeRegex(trimmed)}.*`;
  }

  const keywords = words.slice(0, 4);
  return `(?i)${keywords.map((k) => `(?=.*${escapeRegex(k)})`).join('')}`;
}

function generateId(existingIds: ReadonlySet<string> = new Set()): string {
  const year = new Date().getFullYear();
  const maxAttempts = 800;
  for (let attempt = 0; attempt < maxAttempts; attempt++) {
    const seq = String(Math.floor(Math.random() * 900) + 100);
    const id = `ATR-${year}-${seq}`;
    if (!existingIds.has(id)) {
      return id;
    }
  }
  throw new Error('Unable to generate a unique ATR rule ID after maximum attempts');
}

function getCurrentDate(): string {
  const d = new Date();
  const yyyy = d.getFullYear();
  const mm = String(d.getMonth() + 1).padStart(2, '0');
  const dd = String(d.getDate()).padStart(2, '0');
  return `${yyyy}/${mm}/${dd}`;
}

export class RuleScaffolder {
  private readonly options: Required<ScaffoldOptions>;

  constructor(options: ScaffoldOptions = {}) {
    this.options = {
      author: options.author ?? 'ATR Community (auto-scaffolded)',
      schemaVersion: options.schemaVersion ?? '0.1',
    };
  }

  /**
   * Generate a complete ATR YAML rule from structured input.
   * Returns a ScaffoldResult with the YAML string, generated ID, and any warnings.
   */
  scaffold(input: ScaffoldInput, existingIds: ReadonlySet<string> = new Set()): ScaffoldResult {
    const warnings = this.validateInput(input);

    const severity = input.severity ?? 'medium';
    const sourceType = input.agentSourceType ?? CATEGORY_TO_SOURCE_TYPE[input.category];
    const field = CATEGORY_TO_FIELD[input.category];
    const id = generateId(existingIds);
    const date = getCurrentDate();

    const conditions: ATRArrayCondition[] = input.examplePayloads.map(
      (payload, idx) => ({
        field,
        operator: 'regex',
        value: buildRegexPattern(payload, input.category),
        description: `Pattern ${idx + 1}: detects "${payload.trim().slice(0, 80)}"`,
      }),
    );

    const truePositives = input.examplePayloads.map((payload) => ({
      input: payload.trim(),
      expected: 'trigger' as const,
    }));

    const trueNegatives = [
      {
        input: 'TODO: Add benign input that should not trigger this rule',
        expected: 'no_trigger' as const,
      },
    ];

    const references: Record<string, string[]> = {};
    if (input.owaspRefs && input.owaspRefs.length > 0) {
      references.owasp_llm = [...input.owaspRefs];
    }
    if (input.mitreRefs && input.mitreRefs.length > 0) {
      references.mitre_atlas = [...input.mitreRefs];
    }

    const conditionExpr = conditions.length > 1 ? 'any' : 'all';

    const rule: Record<string, unknown> = {
      title: input.title,
      id,
      schema_version: this.options.schemaVersion,
      status: 'draft',
      description: input.attackDescription,
      author: this.options.author,
      date,
      severity,
      detection_tier: 'pattern',
      maturity: 'draft',
      ...(Object.keys(references).length > 0 ? { references } : {}),
      tags: {
        category: input.category,
        confidence: severity === 'critical' || severity === 'high' ? 'high' : 'medium',
      },
      agent_source: {
        type: sourceType,
      },
      detection: {
        conditions,
        condition: conditionExpr,
        false_positives: [
          'TODO: Document known false positive scenarios',
        ],
      },
      response: {
        actions: [...SEVERITY_TO_ACTIONS[severity]],
        message_template: `Potential ${input.category} detected: {{matched_patterns}}`,
      },
      test_cases: {
        true_positives: truePositives,
        true_negatives: trueNegatives,
      },
    };

    const yamlStr = yaml.dump(rule, {
      indent: 2,
      lineWidth: 120,
      noRefs: true,
      sortKeys: false,
      quotingType: '"',
      forceQuotes: false,
    });

    return { yaml: yamlStr, id, warnings };
  }

  /**
   * Validate scaffold input, throwing on invalid required fields
   * and returning warnings for non-critical issues.
   */
  private validateInput(input: ScaffoldInput): string[] {
    const warnings: string[] = [];

    if (!input.title || input.title.trim().length === 0) {
      throw new Error('ScaffoldInput.title is required and must be non-empty');
    }
    if (!input.category) {
      throw new Error('ScaffoldInput.category is required');
    }
    if (!input.attackDescription || input.attackDescription.trim().length === 0) {
      throw new Error('ScaffoldInput.attackDescription is required and must be non-empty');
    }
    if (!input.examplePayloads || input.examplePayloads.length === 0) {
      throw new Error('ScaffoldInput.examplePayloads must contain at least one payload');
    }

    if (input.examplePayloads.length < 3) {
      warnings.push(
        'Fewer than 3 example payloads - consider adding more for better pattern coverage.',
      );
    }

    return warnings;
  }
}
