from typing import Any

from mmar_llm.airi_endpoint import AiriChatEndpoint
from mmar_llm.llm_endpoint import LLMEndpoint
from mmar_llm.fusion_brain_endpoint import FusionBrainEndpoint
from mmar_llm.gigachat_endpoint import (
    GigaChatCensoredEndpoint,
    GigaChatEndpoint,
    GigaMax2Endpoint,
    GigaMax2SberdevicesEndpoint,
    GigaMaxEndpoint,
    GigaPlusEndpoint,
)
from mmar_llm.openrouter_endpoint import OpenRouterEndpoint
from mmar_llm.yandex_gpt_endpoint import YandexGPTEndpoint

ENDPOINTS: dict[str, type[LLMEndpoint]] = {
    "airi": AiriChatEndpoint,
    "giga": GigaChatEndpoint,
    "giga-max": GigaMaxEndpoint,
    "giga-max-2": GigaMax2Endpoint,
    "giga-plus": GigaPlusEndpoint,
    "giga-censored": GigaChatCensoredEndpoint,
    "open-router": OpenRouterEndpoint,
    "yandex": YandexGPTEndpoint,
    "fusion-brain": FusionBrainEndpoint,
    "giga-max-2-sberdevices": GigaMax2SberdevicesEndpoint,
}


def create_endpoint(endpoint_kind: str, endpoint_args: dict[str, Any]):
    endpoint_class = ENDPOINTS.get(endpoint_kind)
    if endpoint_class is None:
        err = f"Not found endpoint for endpoint_kind={endpoint_kind}"
        raise ValueError(err)
    return endpoint_class(**endpoint_args)
