"""
Model registry and setup helper.

This module maps model aliases to liteLLM provider/model routes and exposes
`setup(...)` to build a configured `LiteLLMChat` instance.
"""

from __future__ import annotations

import os

from .adapter import LiteLLMChat


MODEL_REGISTRY = {
    # ---------------- DeepSeek ----------------
    "deepseek-reasoner": {
        "model": "deepseek/deepseek-reasoner",
        "env_key": "DEEPSEEK_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "DeepSeek Reasoner",
    },
    "deepseek-chat": {
        "model": "deepseek/deepseek-chat",
        "env_key": "DEEPSEEK_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "DeepSeek Chat",
    },
    # ---------------- OpenAI ----------------
    "gpt-5.3": {
        "model": "openai/gpt-5.3",
        "env_key": "OPENAI_API_KEY",
        "max_context_tokens": 400000,
        "supports_native_tool_calling": True,
        "description": "GPT-5.3",
    },
    "gpt-5.2": {
        "model": "openai/gpt-5.2",
        "env_key": "OPENAI_API_KEY",
        "max_context_tokens": 400000,
        "supports_native_tool_calling": True,
        "description": "GPT-5.2",
    },
    "o3-mini": {
        "model": "openai/o3-mini",
        "env_key": "OPENAI_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "OpenAI o3-mini",
    },
    "gpt-4o": {
        "model": "openai/gpt-4o",
        "env_key": "OPENAI_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "GPT-4o",
    },
    "o1-preview": {
        "model": "openai/o1-preview",
        "env_key": "OPENAI_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "OpenAI o1-preview",
    },
    # ---------------- Anthropic ----------------
    "claude-4.6-sonnet": {
        "model": "anthropic/claude-sonnet-4-6",
        "env_key": "ANTHROPIC_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "Claude 4.6 Sonnet",
    },
    "claude-4.5-opus": {
        "model": "anthropic/claude-opus-4-5-20251101",
        "env_key": "ANTHROPIC_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "Claude 4.5 Opus",
    },
    "claude-3.5-sonnet": {
        "model": "anthropic/claude-3-5-sonnet-20241022",
        "env_key": "ANTHROPIC_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "Claude 3.5 Sonnet",
    },
    # ---------------- Gemini ----------------
    "gemini-3.1-pro": {
        "model": "gemini/gemini-3.1-pro",
        "env_key": "GEMINI_API_KEY",
        "max_context_tokens": 1048576,
        "supports_native_tool_calling": True,
        "description": "Gemini 3.1 Pro",
    },
    "gemini-2.5-flash": {
        "model": "gemini/gemini-2.5-flash",
        "env_key": "GEMINI_API_KEY",
        "max_context_tokens": 1048576,
        "supports_native_tool_calling": True,
        "description": "Gemini 2.5 Flash",
    },
    "gemini-1.5-pro": {
        "model": "gemini/gemini-1.5-pro",
        "env_key": "GEMINI_API_KEY",
        "max_context_tokens": 2097152,
        "supports_native_tool_calling": True,
        "description": "Gemini 1.5 Pro",
    },
    # ---------------- GLM ----------------
    "glm-5": {
        "model": "openai/glm-5",
        "env_key": "GLM_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "GLM-5",
        "default_kwargs": {"api_base": "https://open.bigmodel.cn/api/paas/v4/"},
    },
    "glm-4-plus": {
        "model": "openai/glm-4-plus",
        "env_key": "GLM_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "GLM-4 Plus",
        "default_kwargs": {"api_base": "https://open.bigmodel.cn/api/paas/v4/"},
    },
    "glm-4-long": {
        "model": "openai/glm-4-long",
        "env_key": "GLM_API_KEY",
        "max_context_tokens": 1000000,
        "supports_native_tool_calling": True,
        "description": "GLM-4 Long",
        "default_kwargs": {"api_base": "https://open.bigmodel.cn/api/paas/v4/"},
    },
    "glm-4-flash": {
        "model": "openai/glm-4-flash",
        "env_key": "GLM_API_KEY",
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "GLM-4 Flash",
        "default_kwargs": {"api_base": "https://open.bigmodel.cn/api/paas/v4/"},
    },
    # ---------------- Moonshot ----------------
    "kimi-k2.5": {
        "model": "openai/kimi-k2.5",
        "env_key": ["KIMI_API_KEY", "MOONSHOT_API_KEY"],
        "max_context_tokens": 256000,
        "supports_native_tool_calling": True,
        "description": "Kimi k2.5",
        "default_kwargs": {
            "api_base": "https://api.moonshot.cn/v1",
            "extra_body": {"enable_thinking": True},
        },
    },
    "kimi": {
        "model": "openai/moonshot-v1-auto",
        "env_key": ["KIMI_API_KEY", "MOONSHOT_API_KEY"],
        "max_context_tokens": 128000,
        "supports_native_tool_calling": True,
        "description": "Kimi",
        "default_kwargs": {"api_base": "https://api.moonshot.cn/v1"},
    },
    # ---------------- MiniMax ----------------
    "minimax-m2.5": {
        "model": "minimax/MiniMax-M2.5",
        "env_key": "MINIMAX_API_KEY",
        "max_context_tokens": 204800,
        "supports_native_tool_calling": True,
        "description": "MiniMax M2.5",
        "default_kwargs": {
            "api_base": os.environ.get("MINIMAX_API_BASE", "https://api.minimaxi.com/v1"),
        },
    },
    "minimax-text-01": {
        "model": "minimax/MiniMax-M2.1",
        "env_key": "MINIMAX_API_KEY",
        "max_context_tokens": 200000,
        "supports_native_tool_calling": True,
        "description": "MiniMax M2.1",
        "default_kwargs": {
            "api_base": os.environ.get("MINIMAX_API_BASE", "https://api.minimaxi.com/v1"),
        },
    },
    # ---------------- Qwen ----------------
    "qwen-max-latest": {
        "model": "dashscope/qwen-max-latest",
        "env_key": "DASHSCOPE_API_KEY",
        "max_context_tokens": 131072,
        "supports_native_tool_calling": True,
        "description": "Qwen Max Latest",
    },
    "qwen-plus-latest": {
        "model": "dashscope/qwen-plus-latest",
        "env_key": "DASHSCOPE_API_KEY",
        "max_context_tokens": 1000000,
        "supports_native_tool_calling": True,
        "description": "Qwen Plus Latest",
    },
}


def setup(provider: str, model_name: str = None, api_key: str = None, **kwargs) -> LiteLLMChat:
    """
    Build a model instance by provider/model or alias.

    Supported forms:
    1. setup("deepseek", "deepseek-reasoner")
    2. setup("gpt-4o")
    3. setup("openai", "gpt-4o", "sk-xxx")
    """

    if model_name is None:
        alias = provider
        if alias in MODEL_REGISTRY:
            config = MODEL_REGISTRY[alias]
            actual_model_string = config["model"]
            desc = config["description"]
            env_keys = config.get("env_key")
            default_kwargs = config.get("default_kwargs", {})
            registry_max_context_tokens = config.get("max_context_tokens", 128000)
            supports_native_tool_calling = config.get("supports_native_tool_calling", True)
        else:
            actual_model_string = alias
            desc = f"Custom Model: {alias}"
            env_keys = None
            default_kwargs = {}
            registry_max_context_tokens = 128000
            supports_native_tool_calling = True
    else:
        if model_name in MODEL_REGISTRY:
            config = MODEL_REGISTRY[model_name]
            actual_model_string = config["model"]
            desc = config["description"]
            env_keys = config.get("env_key")
            default_kwargs = config.get("default_kwargs", {})
            registry_max_context_tokens = config.get("max_context_tokens", 128000)
            supports_native_tool_calling = config.get("supports_native_tool_calling", True)
        else:
            actual_model_string = f"{provider}/{model_name}"
            desc = f"Custom Model: {actual_model_string}"
            env_keys = None
            default_kwargs = {}
            registry_max_context_tokens = 128000
            supports_native_tool_calling = True

    final_api_key = api_key
    if not final_api_key and env_keys:
        if isinstance(env_keys, list):
            for key in env_keys:
                final_api_key = os.environ.get(key)
                if final_api_key:
                    break
        else:
            final_api_key = os.environ.get(env_keys)

    if not final_api_key:
        raise ValueError(
            f"Missing API key for {desc}.\n"
            "Pass `api_key` explicitly, for example:\n"
            "setup_model('provider', 'model_name', 'sk-xxx')\n"
            f"Or set environment variable(s): {env_keys}"
        )

    final_kwargs = default_kwargs.copy()
    final_kwargs.update(kwargs)

    max_tokens = final_kwargs.pop("max_context_tokens", registry_max_context_tokens)
    api_base = final_kwargs.pop("api_base", None)
    temperature = final_kwargs.pop("temperature", 0.1)
    supports_native_tool_calling = bool(
        final_kwargs.pop("supports_native_tool_calling", supports_native_tool_calling)
    )
    model_backend = str(
        final_kwargs.pop("model_backend", final_kwargs.pop("backend", "litellm")) or "litellm"
    ).strip().lower()

    print(f"[Model] {desc} (liteLLM route: {actual_model_string}, backend: {model_backend})")

    model = LiteLLMChat(
        model=actual_model_string,
        api_key=final_api_key,
        api_base=api_base,
        temperature=temperature,
        backend=model_backend,
        **final_kwargs,
    )
    model.max_context_tokens = max_tokens
    model.supports_native_tool_calling = supports_native_tool_calling
    return model


def list_models() -> None:
    """Print all registered aliases and API key readiness."""
    print("\nRegistered model aliases:")
    print("-" * 75)
    for name, config in MODEL_REGISTRY.items():
        env_keys = config.get("env_key", "")
        if isinstance(env_keys, list):
            env_list = [str(item).strip() for item in env_keys if str(item or "").strip()]
            env_str = " / ".join(env_list)
            has_key = any(bool(os.getenv(item)) for item in env_list)
        else:
            env_str = str(env_keys or "N/A")
            has_key = bool(os.getenv(env_str)) if env_str != "N/A" else False

        status = "OK" if has_key else "MISSING"
        print(f"  [{status:7s}] {name:20s} | {config['description']:30s} | Key: {env_str}")
    print("-" * 75)
    print('Tip: you can also call setup("provider", "model_name") for models not listed here.')
