import { BaseGuardrailPlugin } from '../base-plugin.js';
import { GuardrailPhase, GuardrailContext, GuardrailResult } from '../../types.js';
import { PIIRedactorConfig, ContentPart } from '../../../config/types.js';
import { PIIDetector, PIIDetectorConfig } from './detectors.js';
import { Pseudonymizer } from './pseudonymizer.js';
import { logger } from '../../../utils/logger.js';

/**
 * PII Redactor plugin - detects and redacts sensitive data
 */
export class PIIRedactorPlugin extends BaseGuardrailPlugin {
  name = 'pii_redactor';
  phases: GuardrailPhase[] = ['pre_request', 'post_response', 'pre_tool_input', 'post_tool_output'];

  private detector!: PIIDetector;
  private pseudonymizer!: Pseudonymizer;
  private restoreOnResponse: boolean = false;
  private logDetections: boolean = true;

  async initialize(config: Record<string, unknown>): Promise<void> {
    await super.initialize(config);

    const typedConfig = config as Partial<PIIRedactorConfig>;

    const detectorConfig: PIIDetectorConfig = {
      detectEmails: typedConfig.detect_emails ?? true,
      detectPhones: typedConfig.detect_phones ?? true,
      detectSSN: typedConfig.detect_ssn ?? true,
      detectAPIKeys: typedConfig.detect_api_keys ?? true,
      detectCreditCards: typedConfig.detect_credit_cards ?? true,
      detectIPAddresses: typedConfig.detect_ip_addresses ?? false,
      customPatterns: typedConfig.custom_patterns ?? [],
      allowlist: typedConfig.allowlist ?? [],
      allowlistDomains: typedConfig.allowlist_domains ?? [],
    };

    this.detector = new PIIDetector(detectorConfig);
    this.pseudonymizer = new Pseudonymizer();
    this.restoreOnResponse = typedConfig.restore_on_response ?? false;
    this.logDetections = typedConfig.log_detections ?? true;
    this.priority = typedConfig.priority ?? 25;
  }

  async execute(phase: GuardrailPhase, context: GuardrailContext): Promise<GuardrailResult> {
    switch (phase) {
      case 'pre_request':
      case 'pre_tool_input':
        return this.redactPII(context, phase);

      case 'post_response':
      case 'post_tool_output':
        if (this.restoreOnResponse) {
          return this.restorePII(context, phase);
        }
        return this.allow(context);

      default:
        return this.allow(context);
    }
  }

  private redactPII(context: GuardrailContext, phase: GuardrailPhase): Promise<GuardrailResult> {
    let textToScan: string;
    let field: string;

    if (phase === 'pre_request') {
      textToScan = context.prompt || '';
      field = 'prompt';
    } else {
      textToScan = JSON.stringify(context.toolArgs || {});
      field = 'toolArgs';
    }

    if (!textToScan) {
      return Promise.resolve(this.allow(context));
    }

    const detections = this.detector.detect(textToScan);

    if (detections.length === 0) {
      return Promise.resolve(this.allow(context));
    }

    // Log detections
    if (this.logDetections) {
      logger.info(`PII detected in ${field}:`, {
        requestId: context.requestId,
        types: [...new Set(detections.map((d) => d.type))],
        count: detections.length,
      });
    }

    // Pseudonymize
    const { text: redactedText, mappings } = this.pseudonymizer.pseudonymize(
      textToScan,
      detections
    );

    // Store mappings for potential restoration
    context.metadata.set('pii_mappings', mappings);

    // Record modification
    this.addModification(
      context,
      phase,
      field,
      `Redacted ${detections.length} PII items: ${[...new Set(detections.map((d) => d.type))].join(', ')}`,
      undefined, // Don't store original (contains PII)
      undefined // Don't store new (would expose placeholder patterns)
    );

    // Apply modification
    if (phase === 'pre_request') {
      context.prompt = redactedText;
      // Also update the last message if present
      if (context.messages.length > 0) {
        const lastIndex = context.messages.length - 1;
        const originalContent = context.messages[lastIndex].content;
        context.messages[lastIndex] = {
          ...context.messages[lastIndex],
          content:
            typeof originalContent === 'string'
              ? redactedText
              : originalContent.map((part: ContentPart) =>
                  part.type === 'text'
                    ? { ...part, text: this.redactTextPart(part.text) }
                    : part
                ),
        };
      }
    } else {
      try {
        context.toolArgs = JSON.parse(redactedText) as Record<string, unknown>;
      } catch {
        // If parse fails, store as string
        context.toolArgs = { _redacted: redactedText };
      }
    }

    return Promise.resolve(this.modify(context));
  }

  private restorePII(context: GuardrailContext, phase: GuardrailPhase): Promise<GuardrailResult> {
    const mappings = context.metadata.get('pii_mappings') as Map<string, string> | undefined;

    if (!mappings || mappings.size === 0) {
      return Promise.resolve(this.allow(context));
    }

    let textToRestore: string;

    if (phase === 'post_response') {
      textToRestore = context.response || '';
    } else {
      textToRestore =
        typeof context.toolResult === 'string'
          ? context.toolResult
          : JSON.stringify(context.toolResult);
    }

    if (!textToRestore) {
      return Promise.resolve(this.allow(context));
    }

    const restoredText = this.pseudonymizer.restore(textToRestore, mappings);

    // Only modify if something changed
    if (restoredText === textToRestore) {
      return Promise.resolve(this.allow(context));
    }

    this.addModification(
      context,
      phase,
      phase === 'post_response' ? 'response' : 'toolResult',
      `Restored ${mappings.size} PII placeholders`
    );

    if (phase === 'post_response') {
      context.response = restoredText;
    } else {
      try {
        context.toolResult = JSON.parse(restoredText) as unknown;
      } catch {
        context.toolResult = restoredText;
      }
    }

    return Promise.resolve(this.modify(context));
  }

  /**
   * Redact PII from an individual text string
   */
  private redactTextPart(text: string): string {
    const detections = this.detector.detect(text);
    if (detections.length === 0) return text;
    const { text: redacted } = this.pseudonymizer.pseudonymize(text, detections);
    return redacted;
  }

  /**
   * Get detector for testing
   */
  getDetector(): PIIDetector {
    return this.detector;
  }

  /**
   * Get pseudonymizer for testing
   */
  getPseudonymizer(): Pseudonymizer {
    return this.pseudonymizer;
  }
}