import { IPromptOptions, ModuleOptions } from '../types';
import { OpenAiHandler } from './handlers/OpenAiHandler';
import { AnthropicHandler } from './handlers/AnthropicHandler';
import { GroqHandler } from './handlers/GroqHandler';
import { GoogleHandler } from './handlers/GoogleHandler';
import { MistralHandler } from './handlers/MistralHandler';
import { OpenRouterHandler } from './handlers/OpenRouterHandler';
import { CohereHandler } from './handlers/CohereHandler';
import { XAIHandler } from './handlers/XAIHandler';
import { DeepSeekHandler } from './handlers/DeepSeekHandler';
import { AI21Handler } from './handlers/AI21Handler';
import { CloudflareHandler } from './handlers/CloudflareHandler';
import { getRequestLogger } from './RequestLogger';
import { createLlmShield } from './LlmShield';
import { applyBlackoutToStream } from './BlackoutManager';
import { libx } from 'libx.js/build/bundles/essentials.js';

export class LlmManager {
	private handlers: Map<string, any> = new Map();
	private options: ModuleOptions;

	constructor(options: ModuleOptions) {
		this.options = options;
		this.initializeHandlers();
	}

	private initializeHandlers() {
		// Initialize available handlers
		this.handlers.set('openai', new OpenAiHandler(this.options));
		this.handlers.set('anthropic', new AnthropicHandler(this.options));
		this.handlers.set('groq', new GroqHandler(this.options));
		this.handlers.set('google', new GoogleHandler(this.options));
		this.handlers.set('mistral', new MistralHandler(this.options));
		this.handlers.set('openrouter', new OpenRouterHandler(this.options));
		this.handlers.set('cohere', new CohereHandler(this.options));
		this.handlers.set('xai', new XAIHandler(this.options));
		this.handlers.set('deepseek', new DeepSeekHandler(this.options));
		this.handlers.set('ai21', new AI21Handler(this.options));
		this.handlers.set('cloudflare', new CloudflareHandler(this.options));

		libx.log.v('LlmManager: initialized handlers:', Array.from(this.handlers.keys()));
	}

	async executePrompt(options: IPromptOptions): Promise<ReadableStream> {
		const provider = options.config.provider?.toLowerCase();

		if (!provider) {
			throw new Error('Provider not specified in config');
		}

		const handler = this.handlers.get(provider);

		if (!handler) {
			throw new Error(`Provider '${provider}' not supported. Available: ${Array.from(this.handlers.keys()).join(', ')}`);
		}

		// Apply LLM Shield if enabled
		if (options.useShield) {
			const shield = createLlmShield(true);

			// Scan and sanitize input
			if (options.systemPrompt) {
				const scanResult = shield.scanInput(options.systemPrompt);
				if (!scanResult.safe) {
					libx.log.w('LlmManager: Shield detected issues in system prompt', scanResult.issues);
				}
				options.systemPrompt = shield.sanitizeInput(options.systemPrompt);
			}

			if (options.messages) {
				for (const msg of options.messages) {
					const scanResult = shield.scanInput(msg.content);
					if (!scanResult.safe) {
						libx.log.w('LlmManager: Shield detected issues in message', scanResult.issues);
					}
					msg.content = shield.sanitizeInput(msg.content);
				}
			}
		}

		libx.log.v('LlmManager: executing prompt with provider:', provider, 'model:', options.config.model);

		const isStream = options.config.stream ?? true;

		// Start request tracking
		const logger = getRequestLogger(!!options.meta?.enableLogging);
		const tracker = logger.startRequest(provider, options.config.model || 'unknown', options.meta);

		try {
			let stream = await handler.execute(isStream, options);

			// Apply blackout if enabled
			if (options.useBlackout) {
				stream = await applyBlackoutToStream(stream, true);
			}

			// Log successful request (note: we don't have token count at this point for streaming)
			logger.logRequest(tracker.end(true));

			return stream;
		} catch (error) {
			libx.log.e('LlmManager: execution error:', error);

			// Log failed request
			logger.logRequest(tracker.end(false, undefined, (error as Error).message));

			throw error;
		}
	}

	getSupportedProviders(): string[] {
		return Array.from(this.handlers.keys());
	}

	getHandler(provider: string) {
		return this.handlers.get(provider.toLowerCase());
	}
}