import { Model } from 'mongoose';
import { ConfigDocument } from '../config.model';
import { Migration } from '../registry/migration.interface';

const GROUP_ID_TO_HANDLER_KEY_MAP: Record<string, string> = {
	'chatgpt': 'chatgpt',
	'claude': 'claude',
	'gemini': 'gemini',
	'openrouter': 'openrouter',
	'flux': 'replica',
	'minimax': 'replica',
	'stability.ai': 'stability.ai',
	'stable-diffusion': 'stable-diffusion'
};

const DEFAULT_HANDLER_WAYS: Record<string, Record<string, string>> = {
	chatgpt: {
		text: 'openai-standard',
		imageReading: 'openai-vision',
		audio: 'openai-whisper-v1'
	},
	claude: {
		text: 'anthropic-standard'
	},
	gemini: {
		text: 'gemini-standard',
		image: 'vertex-ai-platform',
		imageReading: 'gemini-vision-v1'
	},
	openrouter: {
		text: 'openrouter-standard'
	},
	replica: {
		image: 'replicate-standard'
	},
	'stability.ai': {
		image: 'stability-v3'
	},
	'stable-diffusion': {
		image: 'stable-diffusion-api'
	}
};

function getHandlerKeyForGroup(groupId: string): string | null {
	const normalizedGroupId = groupId.toLowerCase();
	return GROUP_ID_TO_HANDLER_KEY_MAP[normalizedGroupId] || null;
}

function getDefaultHandlerWay(handlerKey: string, capability: 'text' | 'image' | 'imageReading' | 'audio'): string | null {
	const handlerKeyLower = handlerKey.toLowerCase();
	const ways = DEFAULT_HANDLER_WAYS[handlerKeyLower];
	if (!ways) {
		return null;
	}
	return ways[capability] || null;
}

export const migration_v4: Migration = {
	version: 4,
	name: 'Migrate AI models from group-based service selection to handler config system',
	up: async (configModel: Model<ConfigDocument>) => {
		console.log('Migration 4: up - starting handler config migration');
		
		const connection = configModel.db;
		const aiModelGroupCollection = connection.collection('aimodelgroups');
		
		const groups = await aiModelGroupCollection.find({}).toArray();
		console.log(`Migration 4: up - found ${groups.length} AI model groups to process`);
		
		let totalModelsUpdated = 0;
		let totalHandlersAdded = 0;
		
		for (const group of groups) {
			if (!group.models || !Array.isArray(group.models)) {
				continue;
			}
			
			const groupId = group.id || group._id?.toString();
			if (!groupId) {
				console.log(`Migration 4: up - skipping group without id: ${group.name || 'unknown'}`);
				continue;
			}
			
			const handlerKey = getHandlerKeyForGroup(groupId);
			if (!handlerKey) {
				console.log(`Migration 4: up - no handler key mapping found for group: ${groupId}, skipping`);
				continue;
			}
			
			console.log(`Migration 4: up - processing group: ${groupId} -> handlerKey: ${handlerKey}`);
			
			let groupUpdated = false;
			let modelsInGroupUpdated = 0;
			let handlersInGroupAdded = 0;
			
			const updatedModels = group.models.map((model: any) => {
				let modelUpdated = false;
				const updatedModel = { ...model };
				
				if (model.canGenerateText === true && !model.textHandler) {
					const handlerWay = getDefaultHandlerWay(handlerKey, 'text');
					if (handlerWay) {
						updatedModel.textHandler = {
							handlerKey: handlerKey,
							handlerWay: handlerWay
						};
						modelUpdated = true;
						handlersInGroupAdded++;
						console.log(`Migration 4: up - added textHandler to model ${model.id || model._id}: ${handlerKey}/${handlerWay}`);
					} else {
						console.log(`Migration 4: up - warning: no default handler way for text capability with handlerKey: ${handlerKey}`);
					}
				}
				
				if (model.canGenerateImage === true && !model.imageHandler) {
					const handlerWay = getDefaultHandlerWay(handlerKey, 'image');
					if (handlerWay) {
						updatedModel.imageHandler = {
							handlerKey: handlerKey,
							handlerWay: handlerWay
						};
						modelUpdated = true;
						handlersInGroupAdded++;
						console.log(`Migration 4: up - added imageHandler to model ${model.id || model._id}: ${handlerKey}/${handlerWay}`);
					} else {
						console.log(`Migration 4: up - warning: no default handler way for image capability with handlerKey: ${handlerKey}`);
					}
				}
				
				if (model.canReadImage === true && !model.imageReadingHandler) {
					const handlerWay = getDefaultHandlerWay(handlerKey, 'imageReading');
					if (handlerWay) {
						updatedModel.imageReadingHandler = {
							handlerKey: handlerKey,
							handlerWay: handlerWay
						};
						modelUpdated = true;
						handlersInGroupAdded++;
						console.log(`Migration 4: up - added imageReadingHandler to model ${model.id || model._id}: ${handlerKey}/${handlerWay}`);
					} else {
						console.log(`Migration 4: up - warning: no default handler way for imageReading capability with handlerKey: ${handlerKey}`);
					}
				}
				
				if (modelUpdated) {
					groupUpdated = true;
					modelsInGroupUpdated++;
				}
				
				return updatedModel;
			});
			
			if (groupUpdated) {
				await aiModelGroupCollection.updateOne(
					{ _id: group._id },
					{ $set: { models: updatedModels } }
				);
				totalModelsUpdated += modelsInGroupUpdated;
				totalHandlersAdded += handlersInGroupAdded;
				console.log(`Migration 4: up - updated ${modelsInGroupUpdated} model(s) with ${handlersInGroupAdded} handler(s) in group: ${group.name || groupId}`);
			}
		}
		
		console.log(`Migration 4: up - handler config migration completed. Updated ${totalModelsUpdated} models with ${totalHandlersAdded} handlers across ${groups.length} groups`);
	}
};