from typing import Optional

from rhodecode.apps.ai_agents.ai_settings import AISettings, AIModelName
from rhodecode.apps.ai_agents.models.base import wrap_ai_exceptions, AIService
from rhodecode.apps.ai_agents.models.claude import ClaudeService
from rhodecode.apps.ai_agents.models.gemini import GeminiService
from rhodecode.apps.ai_agents.models.gpt import GPTService


@wrap_ai_exceptions
def get_ai_service(
    api_key: str, model_name: str = AIModelName.GPT.name, api_model_version_name: str = "5-nano"
) -> Optional[AIService]:
    assert api_key, "API key is required"
    try:
        if AIModelName[model_name] is AIModelName.GPT:
            return GPTService(
                AISettings(
                    model_name=AIModelName.GPT,
                    model_version=api_model_version_name,
                    api_key=api_key,
                )
            )
        if AIModelName[model_name] is AIModelName.Claude:
            return ClaudeService(
                AISettings(
                    model_name=AIModelName.Claude,
                    model_version=api_model_version_name,
                    api_key=api_key,
                )
            )
        if AIModelName[model_name] is AIModelName.Gemini:
            return GeminiService(
                AISettings(
                    model_name=AIModelName.Gemini,
                    model_version=api_model_version_name,
                    api_key=api_key,
                )
            )
    except KeyError as ke:
        raise ValueError(f"Unknown model: {ke}")
