import { IRequest } from 'itty-router';
import { LlmManager } from '../modules/LlmManager';
import { Env } from '../types';
import { libx } from 'libx.js/build/bundles/essentials.js';
import { getAuthContext } from '../middleware/auth';
import { ApiKeyManager } from '../modules/ApiKeyManager';
import { resolveModel } from 'ai.libx.js';
import { json } from 'itty-router';

export const completionController = {
	async handleRequest(req: IRequest, env: Env, ctx: any) {
		try {
			const body = await req.json();
			const { messages, systemPrompt, config } = body;

			if (!config || !config.model) {
				return json({
					error: 'Missing required field: config.model is required'
			}, { status: 400 });
			}

			// Auto-resolve model if provider is not specified
			// This allows using model names like "gpt4o", "claude", "gemini-pro" etc.
			if (!config.provider && config.model) {
				const resolvedModel = resolveModel(config.model);
				libx.log.v('Completion: Model resolution:', config.model, '→', resolvedModel);

				// Check if resolution resulted in "provider/model" format
				if (resolvedModel.includes('/')) {
					const slashIndex = resolvedModel.indexOf('/');
					const provider = resolvedModel.slice(0, slashIndex);
					const model = resolvedModel.slice(slashIndex + 1);
					config.provider = provider;
					config.model = model;
					libx.log.v('Completion: Auto-resolved provider:', provider, 'model:', model);
				} else {
					// If still no provider after resolution, return error
					return json({
						error: 'Could not resolve model. Please specify both config.provider and config.model, or use a valid model identifier.'
				}, { status: 400 });
				}
			}

			// Check for authenticated request with Ask API key
			const authContext = getAuthContext(req);
			let apiKeys: Record<string, string>;
			let apiKeyData: any = null;

			if (authContext && authContext.providerKeys) {
				// Use provider keys from Ask API key
				apiKeys = authContext.providerKeys;
				apiKeyData = authContext.apiKeyData;
				libx.log.v('Completion: Using provider keys from Ask API key');
			} else {
				// Fallback to direct provider keys from headers or environment
				apiKeys = {
					openai: req.headers.get('x-openai-api-key') || env.OPENAI_API_KEY,
					anthropic: req.headers.get('x-claude-api-key') || req.headers.get('x-anthropic-api-key') || env.CLAUDE_API_KEY,
					google: req.headers.get('x-google-api-key') || env.GOOGLE_AI_API_KEY,
					groq: req.headers.get('x-groq-api-key') || env.GROQ_API_KEY,
					mistral: req.headers.get('x-mistral-api-key') || env.MISTRAL_API_KEY,
					openrouter: req.headers.get('x-openrouter-api-key') || env.OPENROUTER_API_KEY,
					cohere: req.headers.get('x-cohere-api-key') || env.COHERE_API_KEY,
					ai21: req.headers.get('x-ai21-api-key') || env.AI21_API_KEY,
					xai: req.headers.get('x-xai-api-key') || env.XAI_API_KEY,
					deepseek: req.headers.get('x-deepseek-api-key') || env.DEEPSEEK_API_KEY,
					cloudflare: req.headers.get('x-cloudflare-api-key') || env.CLOUDFLARE_API_KEY,
					stability: req.headers.get('x-stability-api-key') || env.STABILITY_AI_API_KEY,
					falai: req.headers.get('x-falai-api-key') || env.FALAI_API_KEY,
				};
				libx.log.v('Completion: Using direct provider keys from headers/environment');
			}

			// Get Cloudflare Account ID
			const cloudflareAccountId = req.headers.get('x-cloudflare-account-id') || env.CLOUDFLARE_ACCOUNT_ID;

			// Initialize LlmManager with API keys and options
			const llmManager = new LlmManager({
				apiKeys,
				cloudflareAccountId
			});

			// Execute prompt
			const startTime = Date.now();
			const readable = await llmManager.executePrompt({
				messages,
				systemPrompt,
				config,
			});
			const latencyMs = Date.now() - startTime;

			// Log usage if using Ask API key (fire and forget)
			if (apiKeyData && env.ASK_DB) {
				const encryptionKey = env.API_KEY_ENCRYPTION_KEY || 'default-encryption-key-change-in-production';
				const keyManager = new ApiKeyManager(env.ASK_DB, encryptionKey);
				keyManager.logUsage({
					apiKeyId: apiKeyData.id,
					provider: config.provider,
					model: config.model,
					latencyMs,
					success: true
				}).catch(err => {
					libx.log.e('Completion: Error logging usage', err);
				});
			}

			// Determine content type based on plain and stream options
			let contentType = 'application/json';
			if (config?.plain) {
				contentType = 'text/plain; charset=utf-8';
			} else if (config?.stream) {
				contentType = 'text/event-stream';
			}

			// Return stream (CORS headers will be added automatically by server adapter)
			return new Response(readable, {
				headers: {
					'Content-Type': contentType,
				'Cache-Control': 'no-cache',
				'Connection': 'keep-alive',
			}
		});

		} catch (error: any) {
			libx.log.e('Completion controller error:', error);

			// Log error usage if using Ask API key (fire and forget)
			const authContext = getAuthContext(req);
			if (authContext && authContext.apiKeyData && env.ASK_DB) {
				const encryptionKey = env.API_KEY_ENCRYPTION_KEY || 'default-encryption-key-change-in-production';
				const keyManager = new ApiKeyManager(env.ASK_DB, encryptionKey);
				const body = await req.json().catch(() => ({}));
				keyManager.logUsage({
					apiKeyId: authContext.apiKeyData.id,
					provider: body.config?.provider || 'unknown',
					model: body.config?.model || 'unknown',
					success: false,
					errorMessage: error.message
				}).catch(err => {
					libx.log.e('Completion: Error logging failed usage', err);
				});
			}

			return json({
				error: error.message || 'Internal server error',
				timestamp: new Date().toISOString()
		}, { status: 500 });
		}
	},

	async handleStream(req: IRequest, env: Env, ctx: any) {
		// Explicitly set stream to true
		try {
			const body = await req.json();
			body.config = { ...body.config, stream: true };

			// Reuse handleRequest with modified body
			return this.handleRequest({ ...req, json: async () => body } as any, env, ctx);
		} catch (error: any) {
			return json({ error: error.message }, { status: 500 });
		}
	}
};