from __future__ import annotations

from analogai.foundation.infrastructure.embeddings.base import EmbeddingProvider
from analogai.foundation.common.enums.embedding_provider import EmbeddingProvider as EmbeddingProviderEnum
from analogai.foundation.infrastructure.embeddings.local_provider import LocalSentenceTransformerProvider
from analogai.foundation.infrastructure.embeddings.openai_provider import OpenAIEmbeddingProvider
from analogai.foundation.infrastructure.embeddings.azure_openai_provider import AzureOpenAIEmbeddingProvider
from analogai.foundation.infrastructure.embeddings.azure_foundry_provider import AzureFoundryEmbeddingProvider


DEFAULT_LOCAL_MODEL = "all-MiniLM-L6-v2"


def create_embedding_provider(
    provider: EmbeddingProviderEnum,
    model_name: str | None = None,
) -> EmbeddingProvider:
    resolved_model = model_name or DEFAULT_LOCAL_MODEL
    provider_factories = {
        EmbeddingProviderEnum.LOCAL: lambda: LocalSentenceTransformerProvider(
            model_name=resolved_model,
        ),
        EmbeddingProviderEnum.AZURE: lambda: AzureFoundryEmbeddingProvider(
            deployment_name=resolved_model,
        ),
        EmbeddingProviderEnum.AZURE_OPENAI: lambda: AzureOpenAIEmbeddingProvider(
            deployment_name=resolved_model,
        ),
        EmbeddingProviderEnum.AZURE_FOUNDRY: lambda: AzureFoundryEmbeddingProvider(
            deployment_name=resolved_model,
        ),
        EmbeddingProviderEnum.OPENAI: lambda: OpenAIEmbeddingProvider(
            model_name=resolved_model,
        ),
    }

    return provider_factories[provider]()
