/**
 * Defines profiles for different embedding models, including their dimensions.
 */

import type { EmbedderProvider, EmbeddingModelProfiles } from "@roo-code/types"

// Example profiles - expand this list as needed
export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
	openai: {
		"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
		"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
		"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
	},
	ollama: {
		"nomic-embed-text": { dimension: 768, scoreThreshold: 0.4 },
		"nomic-embed-code": {
			dimension: 3584,
			scoreThreshold: 0.15,
			queryPrefix: "Represent this query for searching relevant code: ",
		},
		"mxbai-embed-large": { dimension: 1024, scoreThreshold: 0.4 },
		"all-minilm": { dimension: 384, scoreThreshold: 0.4 },
		// Add default Ollama model if applicable, e.g.:
		// 'default': { dimension: 768 } // Assuming a default dimension
	},
	"openai-compatible": {
		"text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
		"text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
		"text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
		"nomic-embed-code": {
			dimension: 3584,
			scoreThreshold: 0.15,
			queryPrefix: "Represent this query for searching relevant code: ",
		},
	},
	gemini: {
		"gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
		// Deprecated: text-embedding-004 is migrated to gemini-embedding-001 in GeminiEmbedder
		// Kept here for backward-compatible dimension lookup in createVectorStore()
		"text-embedding-004": { dimension: 3072, scoreThreshold: 0.4 },
	},
	mistral: {
		"codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
	},
	"vercel-ai-gateway": {
		// OpenAI models
		"openai/text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
		"openai/text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
		"openai/text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
		// Cohere models
		"cohere/embed-v4.0": { dimension: 1024, scoreThreshold: 0.4 },
		// Google models
		"google/gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
		"google/text-embedding-005": { dimension: 768, scoreThreshold: 0.4 },
		"google/text-multilingual-embedding-002": { dimension: 768, scoreThreshold: 0.4 },
		// Amazon models
		"amazon/titan-embed-text-v2": { dimension: 1024, scoreThreshold: 0.4 },
		// Mistral models
		"mistral/codestral-embed": { dimension: 1536, scoreThreshold: 0.4 },
		"mistral/mistral-embed": { dimension: 1024, scoreThreshold: 0.4 },
	},
	bedrock: {
		// Amazon Titan Embed models
		"amazon.titan-embed-text-v1": { dimension: 1536, scoreThreshold: 0.4 },
		"amazon.titan-embed-text-v2:0": { dimension: 1024, scoreThreshold: 0.4 },
		"amazon.titan-embed-image-v1": { dimension: 1024, scoreThreshold: 0.4 },
		// Amazon Nova Embed models
		"amazon.nova-2-multimodal-embeddings-v1:0": { dimension: 1024, scoreThreshold: 0.4 },
		// Cohere Embed v4 (supports only text for now; multimodal image support planned)
		"cohere.embed-v4:0": { dimension: 1536, scoreThreshold: 0.4 },
		// Cohere Embed v3 models available through Bedrock
		"cohere.embed-english-v3": { dimension: 1024, scoreThreshold: 0.4 },
		"cohere.embed-multilingual-v3": { dimension: 1024, scoreThreshold: 0.4 },
	},
	openrouter: {
		// OpenAI models via OpenRouter
		"openai/text-embedding-3-small": { dimension: 1536, scoreThreshold: 0.4 },
		"openai/text-embedding-3-large": { dimension: 3072, scoreThreshold: 0.4 },
		"openai/text-embedding-ada-002": { dimension: 1536, scoreThreshold: 0.4 },
		// Google models via OpenRouter
		"google/gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
		// Mistral models via OpenRouter
		"mistralai/mistral-embed-2312": { dimension: 1024, scoreThreshold: 0.4 },
		"mistralai/codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },
		// Qwen models via OpenRouter
		"qwen/qwen3-embedding-0.6b": { dimension: 1024, scoreThreshold: 0.4 },
		"qwen/qwen3-embedding-4b": { dimension: 2560, scoreThreshold: 0.4 },
		"qwen/qwen3-embedding-8b": { dimension: 4096, scoreThreshold: 0.4 },
	},
}

/**
 * Retrieves the embedding dimension for a given provider and model ID.
 * @param provider The embedder provider (e.g., "openai").
 * @param modelId The specific model ID (e.g., "text-embedding-3-small").
 * @returns The dimension size or undefined if the model is not found.
 */
export function getModelDimension(provider: EmbedderProvider, modelId: string): number | undefined {
	const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
	if (!providerProfiles) {
		console.warn(`Provider not found in profiles: ${provider}`)
		return undefined
	}

	const modelProfile = providerProfiles[modelId]
	if (!modelProfile) {
		// Don't warn here, as it might be a custom model ID not in our profiles
		// console.warn(`Model not found for provider ${provider}: ${modelId}`)
		return undefined // Or potentially return a default/fallback dimension?
	}

	return modelProfile.dimension
}

/**
 * Retrieves the score threshold for a given provider and model ID.
 * @param provider The embedder provider (e.g., "openai").
 * @param modelId The specific model ID (e.g., "text-embedding-3-small").
 * @returns The score threshold or undefined if the model is not found.
 */
export function getModelScoreThreshold(provider: EmbedderProvider, modelId: string): number | undefined {
	const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
	if (!providerProfiles) {
		return undefined
	}

	const modelProfile = providerProfiles[modelId]
	return modelProfile?.scoreThreshold
}

/**
 * Retrieves the query prefix for a given provider and model ID.
 * @param provider The embedder provider (e.g., "openai").
 * @param modelId The specific model ID (e.g., "nomic-embed-code").
 * @returns The query prefix or undefined if the model doesn't require one.
 */
export function getModelQueryPrefix(provider: EmbedderProvider, modelId: string): string | undefined {
	const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
	if (!providerProfiles) {
		return undefined
	}

	const modelProfile = providerProfiles[modelId]
	return modelProfile?.queryPrefix
}

/**
 * Gets the default *specific* embedding model ID based on the provider.
 * Does not include the provider prefix.
 * Currently defaults to OpenAI's 'text-embedding-3-small'.
 * TODO: Make this configurable or more sophisticated.
 * @param provider The embedder provider.
 * @returns The default specific model ID for the provider (e.g., "text-embedding-3-small").
 */
export function getDefaultModelId(provider: EmbedderProvider): string {
	switch (provider) {
		case "openai":
		case "openai-compatible":
			return "text-embedding-3-small"

		case "ollama": {
			// Choose a sensible default for Ollama, e.g., the first one listed or a specific one
			const ollamaModels = EMBEDDING_MODEL_PROFILES.ollama
			const defaultOllamaModel = ollamaModels && Object.keys(ollamaModels)[0]
			if (defaultOllamaModel) {
				return defaultOllamaModel
			}
			// Fallback if no Ollama models are defined (shouldn't happen with the constant)
			console.warn("No default Ollama model found in profiles.")
			// Return a placeholder or throw an error, depending on desired behavior
			return "unknown-default" // Placeholder specific model ID
		}

		case "gemini":
			return "gemini-embedding-001"

		case "mistral":
			return "codestral-embed-2505"

		case "vercel-ai-gateway":
			return "openai/text-embedding-3-large"

		case "bedrock":
			return "amazon.titan-embed-text-v2:0"
		case "openrouter":
			return "openai/text-embedding-3-large"

		default:
			// Fallback for unknown providers
			console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)
			return "text-embedding-3-small"
	}
}