import { libx } from 'libx.js/build/bundles/essentials.js';

/**
 * LlmShield - Input protection and leak detection for LLM interactions
 * Helps prevent prompt injection and data leakage
 */
export class LlmShield {
	private enabled: boolean;

	constructor(enabled: boolean = false) {
		this.enabled = enabled;
	}

	/**
	 * Scan input for potential security issues
	 */
	public scanInput(input: string): ShieldScanResult {
		if (!this.enabled) {
			return { safe: true, issues: [], score: 1.0 };
		}

		const issues: string[] = [];
		let score = 1.0;

		// Check for prompt injection patterns
		const injectionPatterns = [
			/ignore\s+(previous|all|above)\s+(instructions|rules|prompts)/i,
			/disregard\s+(previous|all|above)\s+(instructions|rules|prompts)/i,
			/forget\s+(previous|all|above)\s+(instructions|rules|prompts)/i,
			/you\s+are\s+now\s+a/i,
			/system\s*:\s*new\s+instructions/i,
			/\[SYSTEM\]/i,
			/\<\|system\|\>/i,
			/roleplay\s+as/i,
			/act\s+as\s+if/i,
		];

		for (const pattern of injectionPatterns) {
			if (pattern.test(input)) {
				issues.push(`Potential prompt injection detected: ${pattern.source}`);
				score -= 0.2;
			}
		}

		// Check for excessive special characters (obfuscation attempts)
		const specialCharRatio = (input.match(/[^a-zA-Z0-9\s.,!?]/g) || []).length / input.length;
		if (specialCharRatio > 0.3) {
			issues.push('High ratio of special characters detected');
			score -= 0.15;
		}

		// Check for data exfiltration patterns
		const exfiltrationPatterns = [
			/print\s+(all|entire)\s+(conversation|history|context)/i,
			/show\s+me\s+(your|the)\s+(system|initial)\s+prompt/i,
			/what\s+(are|were)\s+your\s+(initial|original)\s+instructions/i,
			/reveal\s+your\s+(system|initial)\s+prompt/i,
			/output\s+your\s+(system|initial)\s+prompt/i,
		];

		for (const pattern of exfiltrationPatterns) {
			if (pattern.test(input)) {
				issues.push(`Potential data exfiltration attempt: ${pattern.source}`);
				score -= 0.25;
			}
		}

		// Check for excessive length (potential DoS)
		if (input.length > 50000) {
			issues.push('Input length exceeds recommended limit');
			score -= 0.1;
		}

		// Check for repeated patterns (potential attack)
		const words = input.split(/\s+/);
		const uniqueWords = new Set(words);
		const repetitionRatio = uniqueWords.size / words.length;
		if (words.length > 100 && repetitionRatio < 0.3) {
			issues.push('High repetition ratio detected');
			score -= 0.1;
		}

		score = Math.max(0, score);
		const safe = score >= 0.6 && issues.length === 0;

		libx.log.v('LlmShield: input scan complete', {
			safe,
			score,
			issuesCount: issues.length,
			inputLength: input.length
		});

		return { safe, issues, score };
	}

	/**
	 * Sanitize input by removing or escaping suspicious content
	 */
	public sanitizeInput(input: string): string {
		if (!this.enabled) return input;

		let sanitized = input;

		// Remove system-level markers
		sanitized = sanitized.replace(/\[SYSTEM\]/gi, '[REMOVED]');
		sanitized = sanitized.replace(/\<\|system\|\>/gi, '');

		// Escape markdown code blocks that might contain injection attempts
		sanitized = sanitized.replace(/```(system|admin|root)/gi, '```text');

		// Normalize excessive whitespace
		sanitized = sanitized.replace(/\s{10,}/g, ' ');

		// Limit repeated characters
		sanitized = sanitized.replace(/(.)\1{20,}/g, '$1$1$1');

		libx.log.v('LlmShield: input sanitized', {
			originalLength: input.length,
			sanitizedLength: sanitized.length,
			changed: input !== sanitized
		});

		return sanitized;
	}

	/**
	 * Scan output for potential data leaks
	 */
	public scanOutput(output: string): ShieldScanResult {
		if (!this.enabled) {
			return { safe: true, issues: [], score: 1.0 };
		}

		const issues: string[] = [];
		let score = 1.0;

		// Check for leaked API keys or tokens
		const secretPatterns = [
			/sk-[a-zA-Z0-9]{48}/,  // OpenAI-style keys
			/sk-ant-[a-zA-Z0-9-]{48,}/,  // Anthropic keys
			/gsk_[a-zA-Z0-9]{52}/,  // Groq keys
			/AIza[a-zA-Z0-9_-]{35}/,  // Google API keys
			/xai-[a-zA-Z0-9]{40,}/,  // XAI keys
			/[a-f0-9]{64}/,  // Generic 64-char hex tokens
		];

		for (const pattern of secretPatterns) {
			if (pattern.test(output)) {
				issues.push(`Potential API key leak detected`);
				score -= 0.5;
			}
		}

		// Check for PII patterns
		const piiPatterns = [
			/\b\d{3}-\d{2}-\d{4}\b/,  // SSN
			/\b\d{16}\b/,  // Credit card
			/\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/,  // Email (be lenient)
		];

		for (const pattern of piiPatterns) {
			const matches = output.match(pattern);
			if (matches && matches.length > 2) {
				issues.push(`Multiple PII-like patterns detected`);
				score -= 0.2;
			}
		}

		score = Math.max(0, score);
		const safe = score >= 0.7;

		libx.log.v('LlmShield: output scan complete', {
			safe,
			score,
			issuesCount: issues.length,
			outputLength: output.length
		});

		return { safe, issues, score };
	}
}

export interface ShieldScanResult {
	safe: boolean;
	issues: string[];
	score: number;  // 0.0 - 1.0, higher is safer
}

/**
 * Factory function to create a shield instance
 */
export function createLlmShield(enabled: boolean = true): LlmShield {
	return new LlmShield(enabled);
}