/**
 * OutputSanitizer — Post-LLM output validation and cleaning.
 *
 * Strips leaked system prompt fragments, tool call instructions,
 * HTML/script injection, and canary tokens from LLM output before
 * it reaches the end user.
 */

import type { ShieldXConfig } from '../types/detection.js'

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/** Result of output sanitization */
export interface OutputSanitizationResult {
  readonly sanitized: string
  readonly originalLength: number
  readonly sanitizedLength: number
  readonly systemPromptLeakDetected: boolean
  readonly toolCallsRemoved: number
  readonly scriptInjectionsCleaned: number
  readonly canaryTokensRemoved: number
  readonly issues: readonly string[]
}

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/**
 * Patterns indicating leaked system prompt fragments.
 * If these appear in LLM output, the system prompt has leaked.
 */
const SYSTEM_PROMPT_LEAK_PATTERNS: readonly RegExp[] = Object.freeze([
  /\[?system\s*(?:prompt|message|instruction)\]?\s*:/i,
  /you\s+are\s+(?:a|an)\s+(?:helpful|AI|language\s+model)\b/i,
  /\bmy\s+system\s+prompt\s+is\b/i,
  /\bhere\s+(?:is|are)\s+my\s+(?:instructions?|rules?|guidelines?)\b/i,
  /\b(?:system|hidden)\s+(?:instructions?|directives?)\s*:/i,
  /\bI\s+was\s+(?:given|told|instructed)\s+(?:to|that)\b.*\brole\b/i,
  /\bmy\s+(?:original|initial|base)\s+(?:prompt|instructions?)\b/i,
  /\bconfidential\s+(?:system\s+)?instructions?\b/i,
  /```\s*(?:system|hidden|internal)\b[^`]*```/gis,
  /<(system|internal|hidden)_?(?:prompt|instruction)[^>]*>[\s\S]*?<\/\1>/gi,
])

/**
 * Tool call patterns that should never appear in user-facing output.
 */
const TOOL_CALL_PATTERNS: readonly RegExp[] = Object.freeze([
  /\b(?:function|tool)_call\s*\(/gi,
  /\{"(?:function|tool|action)":\s*"/gi,
  /<(?:function|tool)_call[^>]*>[\s\S]*?<\/(?:function|tool)_call>/gi,
  /\b(?:execute|run|call)\s*:\s*\{[\s\S]*?\}/gi,
  /\[(?:TOOL|ACTION|FUNCTION)\s*:\s*[^\]]+\]/gi,
  /```(?:tool|function|action)\b[\s\S]*?```/gi,
])

/**
 * HTML/script injection patterns to strip from output.
 */
const SCRIPT_INJECTION_PATTERNS: readonly RegExp[] = Object.freeze([
  /<script\b[^>]*>[\s\S]*?<\/script>/gi,
  /\bon\w+\s*=\s*["'][^"']*["']/gi,
  /<iframe\b[^>]*>[\s\S]*?<\/iframe>/gi,
  /<object\b[^>]*>[\s\S]*?<\/object>/gi,
  /<embed\b[^>]*\/?>/gi,
  /<link\b[^>]*\/?>/gi,
  /<meta\b[^>]*\/?>/gi,
  /javascript\s*:/gi,
  /data\s*:\s*text\/html/gi,
  /vbscript\s*:/gi,
  /<svg\b[^>]*\bon\w+\s*=/gi,
  /<img\b[^>]*\bon\w+\s*=/gi,
])

/**
 * Canary token pattern — UUIDs and hex strings that look like tracking tokens.
 * Matches the format generated by PolymorphicAssembler and other ShieldX modules.
 */
const CANARY_TOKEN_PATTERNS: readonly RegExp[] = Object.freeze([
  /\[CANARY:[a-f0-9-]{8,}\]/gi,
  /\bCANARY_[A-Z0-9]{8,32}\b/g,
  /\bshieldx_canary_[a-f0-9]{8,}\b/gi,
  /<!--\s*canary:[^-]*-->/gi,
])

// ---------------------------------------------------------------------------
// Implementation
// ---------------------------------------------------------------------------

/**
 * Post-LLM output sanitizer.
 *
 * Validates and cleans LLM responses before they reach the user.
 * Strips system prompt leaks, tool call instructions, script injection,
 * and canary tokens from visible output.
 */
export class OutputSanitizer {
  private readonly _config: ShieldXConfig

  /** Access the active configuration */
  get config(): ShieldXConfig { return this._config }

  constructor(config: ShieldXConfig) {
    this._config = config
  }

  /**
   * Sanitize LLM output by removing dangerous or leaked content.
   *
   * @param output - Raw LLM output string
   * @returns Sanitized output with metadata about what was cleaned
   */
  async sanitize(output: string): Promise<OutputSanitizationResult> {
    if (!output || output.trim().length === 0) {
      return this.buildResult(output, '', false, 0, 0, 0, [])
    }

    const issues: string[] = []
    let current = output

    // Phase 1: Detect and strip system prompt leaks
    const { cleaned: afterPromptStrip, detected: promptLeakDetected } =
      this.stripSystemPromptLeaks(current)
    if (promptLeakDetected) {
      issues.push('system_prompt_leak_detected')
    }
    current = afterPromptStrip

    // Phase 2: Remove tool call instructions
    const { cleaned: afterToolStrip, count: toolCallsRemoved } =
      this.stripToolCalls(current)
    if (toolCallsRemoved > 0) {
      issues.push(`tool_calls_removed:${toolCallsRemoved}`)
    }
    current = afterToolStrip

    // Phase 3: Sanitize HTML/script injection
    const { cleaned: afterScriptStrip, count: scriptsCleaned } =
      this.stripScriptInjection(current)
    if (scriptsCleaned > 0) {
      issues.push(`script_injections_cleaned:${scriptsCleaned}`)
    }
    current = afterScriptStrip

    // Phase 4: Remove canary tokens from visible output
    const { cleaned: afterCanaryStrip, count: canaryTokensRemoved } =
      this.stripCanaryTokens(current)
    if (canaryTokensRemoved > 0) {
      issues.push(`canary_tokens_removed:${canaryTokensRemoved}`)
    }
    current = afterCanaryStrip

    // Phase 5: Normalize whitespace
    const normalized = current.replace(/\n{3,}/g, '\n\n').trim()

    return this.buildResult(
      output,
      normalized,
      promptLeakDetected,
      toolCallsRemoved,
      scriptsCleaned,
      canaryTokensRemoved,
      issues,
    )
  }

  /** Strip system prompt leak patterns from output */
  private stripSystemPromptLeaks(
    input: string,
  ): { readonly cleaned: string; readonly detected: boolean } {
    let detected = false
    let cleaned = input

    for (const pattern of SYSTEM_PROMPT_LEAK_PATTERNS) {
      const resetPattern = new RegExp(pattern.source, pattern.flags)
      if (resetPattern.test(cleaned)) {
        detected = true
        cleaned = cleaned.replace(resetPattern, '[REDACTED]')
      }
    }

    return { cleaned, detected }
  }

  /** Strip tool call patterns from output */
  private stripToolCalls(
    input: string,
  ): { readonly cleaned: string; readonly count: number } {
    let count = 0
    let cleaned = input

    for (const pattern of TOOL_CALL_PATTERNS) {
      const resetPattern = new RegExp(pattern.source, pattern.flags)
      const matches = cleaned.match(resetPattern)
      if (matches) {
        count += matches.length
        cleaned = cleaned.replace(resetPattern, '')
      }
    }

    return { cleaned, count }
  }

  /** Strip HTML/script injection patterns */
  private stripScriptInjection(
    input: string,
  ): { readonly cleaned: string; readonly count: number } {
    let count = 0
    let cleaned = input

    for (const pattern of SCRIPT_INJECTION_PATTERNS) {
      const resetPattern = new RegExp(pattern.source, pattern.flags)
      const matches = cleaned.match(resetPattern)
      if (matches) {
        count += matches.length
        cleaned = cleaned.replace(resetPattern, '')
      }
    }

    return { cleaned, count }
  }

  /** Strip canary tokens from visible output */
  private stripCanaryTokens(
    input: string,
  ): { readonly cleaned: string; readonly count: number } {
    let count = 0
    let cleaned = input

    for (const pattern of CANARY_TOKEN_PATTERNS) {
      const resetPattern = new RegExp(pattern.source, pattern.flags)
      const matches = cleaned.match(resetPattern)
      if (matches) {
        count += matches.length
        cleaned = cleaned.replace(resetPattern, '')
      }
    }

    return { cleaned, count }
  }

  /** Build an immutable output sanitization result */
  private buildResult(
    original: string,
    sanitized: string,
    systemPromptLeakDetected: boolean,
    toolCallsRemoved: number,
    scriptInjectionsCleaned: number,
    canaryTokensRemoved: number,
    issues: readonly string[],
  ): OutputSanitizationResult {
    return Object.freeze({
      sanitized,
      originalLength: original.length,
      sanitizedLength: sanitized.length,
      systemPromptLeakDetected,
      toolCallsRemoved,
      scriptInjectionsCleaned,
      canaryTokensRemoved,
      issues: Object.freeze([...issues]),
    })
  }
}
