import os
from openai import AsyncOpenAI
from dotenv import load_dotenv

load_dotenv()

PROVIDER_CONFIG = {
    "local": {
        "base_url": "http://localhost:11434/v1",
        "api_key": "OLLAMA_API_KEY",
        "models": [
            "qwen2.5:7b",
            "qwen2.5:3b",
            "llama3.1:8b",
        ],
    },
    "modal": {
        "base_url": os.getenv("MODAL_BASE_URL"),
        "api_key": "MODAL_API_KEY",
        "models": [
            "google/gemma-4-26B-A4B-it",
        ],
    },
    "nvidia": {
        "base_url": "https://integrate.api.nvidia.com/v1",
        "api_key": "NVIDIA_NIM_API_KEY",
        "models": [
            "minimaxai/minimax-m2.7",
            "stepfun-ai/step-3.5-flash",
            "z-ai/glm4.7",
            "deepseek-ai/deepseek-v3.2",
            "moonshotai/kimi-k2-thinking",
            "mistralai/devstral-2-123b-instruct-2512",
            "mistralai/mistral-large-3-675b-instruct-2512",
            "bytedance/seed-oss-36b-instruct",
            "qwen/qwen3-coder-480b-a35b-instruct",
            "nvidia/nemotron-3-ultra-550b-a55b",
        ],
    },
    "zai": {
        "base_url": "https://api.z.ai/api/paas/v4/",
        "api_key": "ZAI_API_KEY",
        "models": [
            "glm-4.7-flash",
            "glm-4.6v-flash",
        ],
    },
}

DEFAULT_PROVIDER = os.getenv("LLM_PROVIDER", "modal")


def get_llm_client(provider: str | None = None) -> AsyncOpenAI:
    provider = provider or DEFAULT_PROVIDER
    cfg = PROVIDER_CONFIG.get(provider) or PROVIDER_CONFIG[DEFAULT_PROVIDER]
    api_key = os.getenv(cfg["api_key"]) or "ollama"
    return AsyncOpenAI(base_url=cfg["base_url"], api_key=api_key)