from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any

import yaml

from core.config.initial.init_data import load_llm_catalog


def _provider_key_env_name(factory_name: str) -> str:
    normalized = factory_name.strip().upper()
    if normalized == "VOLCENGINE":
        normalized = "ARK"
    return f"SOCIAL_LLM__PROVIDER_KEYS__{normalized}"


def build_proxy_config() -> dict[str, Any]:
    catalog = load_llm_catalog()

    factories = catalog.get("factories", [])
    llms = catalog.get("llms", [])
    if not isinstance(factories, list) or not isinstance(llms, list):
        raise ValueError("invalid llm catalog format")

    factory_url_map: dict[str, str] = {}
    for factory in factories:
        if not isinstance(factory, dict):
            continue
        name = str(factory.get("name", "")).strip().lower()
        request_url = str(factory.get("request_url", "")).strip()
        if name and request_url:
            factory_url_map[name] = request_url

    model_list: list[dict[str, Any]] = []
    for llm in llms:
        if not isinstance(llm, dict):
            continue
        model_code = str(llm.get("model_code", "")).strip()
        factory_name = str(llm.get("factory_name", "")).strip()
        litellm_model = str(llm.get("litellm_model", "")).strip()
        if not model_code or not factory_name or not litellm_model:
            continue

        api_base = factory_url_map.get(factory_name.lower())
        if not api_base:
            raise ValueError(
                f"factory request_url missing for model {model_code}: {factory_name}"
            )

        env_key_name = _provider_key_env_name(factory_name)
        provider_model = (
            litellm_model.split("/", 1)[1] if "/" in litellm_model else litellm_model
        )

        model_list.append(
            {
                "model_name": model_code,
                "litellm_params": {
                    "model": f"openai/{provider_model}",
                    "api_base": api_base,
                    "api_key": f"os.environ/{env_key_name}",
                },
            }
        )

    if not model_list:
        raise ValueError("no models found in llm catalog")

    return {"model_list": model_list}


def main() -> int:
    parser = argparse.ArgumentParser(description="Build LiteLLM proxy config")
    parser.add_argument("--output", required=True, help="Output YAML file path")
    args = parser.parse_args()

    output_path = Path(args.output).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    config = build_proxy_config()
    with output_path.open("w", encoding="utf-8") as file:
        yaml.safe_dump(config, file, sort_keys=False, allow_unicode=False)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())