/**
 * Prompt Module System
 *
 * Provides category-aware prompt assembly for AI validation.
 * Only includes relevant prompt sections based on the categories in each batch,
 * reducing token usage by ~40-60% while maintaining identical behavior.
 *
 * Modules are assembled in fixed order to maximize Anthropic prefix cache hits.
 */

import type { VulnerabilityCategory } from '../../../shared/types'
import { COMMON_PROMPT } from './common'
import { AUTH_ACCESS_MODULE } from './auth-access'
import { XSS_PROMPT_MODULE } from './xss-prompt'
import { SECRETS_CRYPTO_MODULE } from './secrets-crypto'
import { AI_PATTERNS_MODULE } from './ai-patterns'
import { OWASP_CLASSIC_MODULE } from './owasp-classic'

// ============================================================================
// Types
// ============================================================================

export type PromptModuleName =
  | 'auth_access'
  | 'xss_prompt'
  | 'secrets_crypto'
  | 'ai_patterns'
  | 'owasp_classic'

export interface PromptModule {
  name: PromptModuleName
  content: string
  categories: ReadonlySet<VulnerabilityCategory>
}

// ============================================================================
// Module Definitions (fixed order for prefix cache hits)
// ============================================================================

const MODULES: readonly PromptModule[] = [
  {
    name: 'auth_access',
    content: AUTH_ACCESS_MODULE,
    categories: new Set<VulnerabilityCategory>(['missing_auth', 'security_bypass']),
  },
  {
    name: 'xss_prompt',
    content: XSS_PROMPT_MODULE,
    categories: new Set<VulnerabilityCategory>(['xss', 'ai_prompt_injection']),
  },
  {
    name: 'secrets_crypto',
    content: SECRETS_CRYPTO_MODULE,
    categories: new Set<VulnerabilityCategory>([
      'hardcoded_secret',
      'high_entropy_string',
      'sensitive_variable',
      'weak_crypto',
    ]),
  },
  {
    name: 'ai_patterns',
    content: AI_PATTERNS_MODULE,
    categories: new Set<VulnerabilityCategory>([
      'ai_pattern',
      'ai_prompt_injection',
      'ai_unsafe_execution',
      'ai_overpermissive_tool',
      'suspicious_package',
      'ai_rag_exfiltration',
      'ai_endpoint_unprotected',
      'ai_schema_mismatch',
      'ai_package_hallucination',
      'ai_rag_corpus_poisoning',
      'ai_rag_pii_leakage',
      'ai_mcp_tool_poisoning',
      'ai_mcp_credential_issue',
      'ai_mcp_confused_deputy',
      'ai_mcp_description_injection',
      'ai_mcp_server_shadowing',
      'ai_mcp_config_secrets',
      'ai_mcp_config_permissions',
      'ai_rag_query_injection',
      'ai_rag_embedding_poisoning',
      'ai_rag_chunk_injection',
      'ai_package_typosquat',
      'ai_package_malicious',
      'ai_unsafe_model_load',
      'ai_unverified_model',
      'ai_unsafe_finetuning',
      'ai_excessive_agency',
      'ai_skill_injection',
    ]),
  },
  {
    name: 'owasp_classic',
    content: OWASP_CLASSIC_MODULE,
    categories: new Set<VulnerabilityCategory>([
      'missing_security_headers',
      'ssrf',
      'log_injection',
      'xxe',
    ]),
  },
] as const

// ============================================================================
// Category-to-Module Mapping
// ============================================================================

/**
 * Maps each VulnerabilityCategory to the module(s) it requires.
 * Categories not in this map get COMMON only (intentionally unmapped).
 */
export const CATEGORY_TO_MODULE: ReadonlyMap<VulnerabilityCategory, PromptModuleName[]> = (() => {
  const map = new Map<VulnerabilityCategory, PromptModuleName[]>()
  for (const mod of MODULES) {
    for (const cat of mod.categories) {
      const existing = map.get(cat) || []
      existing.push(mod.name)
      map.set(cat, existing)
    }
  }
  return map
})()

/**
 * Categories that are intentionally unmapped to any specific module.
 * They get the COMMON prompt only (which includes condensed heuristic reminders).
 * If a new category is added to VulnerabilityCategory but not here or in CATEGORY_TO_MODULE,
 * the category completeness test will fail.
 */
export const INTENTIONALLY_UNMAPPED: ReadonlySet<VulnerabilityCategory> = new Set([
  'dangerous_function',
  'sql_injection',
  'command_injection',
  'insecure_config',
  'cors_misconfiguration',
  'root_container',
  'dangerous_file',
  'sensitive_url',
  'data_exposure',
])

// ============================================================================
// Assembler Functions
// ============================================================================

/**
 * Assemble a validation prompt containing only the modules relevant to
 * the given set of vulnerability categories.
 *
 * Always starts with COMMON, then adds modules in fixed order:
 * auth_access -> xss_prompt -> secrets_crypto -> ai_patterns -> owasp_classic
 *
 * Deterministic ordering ensures Anthropic prefix cache hits across batches.
 *
 * @param categories - The vulnerability categories present in the current batch
 * @returns The assembled prompt string
 */
export function assembleValidationPrompt(categories: VulnerabilityCategory[]): string {
  // Determine which modules are needed
  const neededModules = new Set<PromptModuleName>()
  for (const cat of categories) {
    const modules = CATEGORY_TO_MODULE.get(cat)
    if (modules) {
      for (const mod of modules) {
        neededModules.add(mod)
      }
    }
    // Unmapped categories just get COMMON (no additional module)
  }

  // Build prompt: COMMON first, then modules in fixed order
  const parts: string[] = [COMMON_PROMPT]
  for (const mod of MODULES) {
    if (neededModules.has(mod.name)) {
      parts.push(mod.content)
    }
  }

  return parts.join('\n')
}

/**
 * Get the full validation prompt with all modules included.
 * Equivalent to the old monolithic HIGH_CONTEXT_VALIDATION_PROMPT.
 * Used for legacy compatibility and the completeness test.
 */
export function getFullValidationPrompt(): string {
  const parts: string[] = [COMMON_PROMPT]
  for (const mod of MODULES) {
    parts.push(mod.content)
  }
  return parts.join('\n')
}

// Re-export module constants for direct access in tests
export { COMMON_PROMPT } from './common'
export { AUTH_ACCESS_MODULE } from './auth-access'
export { XSS_PROMPT_MODULE } from './xss-prompt'
export { SECRETS_CRYPTO_MODULE } from './secrets-crypto'
export { AI_PATTERNS_MODULE } from './ai-patterns'
export { OWASP_CLASSIC_MODULE } from './owasp-classic'
