"""Set Kimi-related env vars as Hopsworks account env vars."""

from __future__ import annotations

import hopsworks

ENV_VARS: dict[str, str] = {
    "ANTHROPIC_AUTH_TOKEN": "sk-MCSyumUGcRGSxkvpixbEGz9z8vbLaOx5XicJyaPOBbbdLjO9",
    "ANTHROPIC_BASE_URL": "https://api.moonshot.ai/anthropic",
    "ANTHROPIC_DEFAULT_HAIK": "false",
    "ANTHROPIC_DEFAULT_OPUS": "kimi-k2.6",
    "ANTHROPIC_DEFAULT_SONI": "kimi-k2.6",
    "ANTHROPIC_MODEL": "kimi-k2.6",
    "CLAUDE_CODE_SUBAGENT_": "kimi-k2.6",
    "ENABLE_TOOL_SEARCH": "kimi-k2.6",
}


def main() -> None:
    hopsworks.login()
    api = hopsworks.get_env_vars_api()

    for name, value in ENV_VARS.items():
        api.set_env_var(name, value)
        print(f"  set {name}")

    existing = {v.name for v in api.get_env_vars(include_value=False)}
    missing = [n for n in ENV_VARS if n not in existing]
    if missing:
        raise SystemExit(f"Missing after upsert: {missing}")
    print(f"OK — {len(ENV_VARS)} env var(s) present on account.")


if __name__ == "__main__":
    main()