#!/bin/sh
# Registers models into LiteLLM's database so they are UI-manageable.
# Idempotent: skips models that are already in the DB.

LITELLM_URL="${LITELLM_URL:-http://localhost:18804}"
LITELLM_MASTER_KEY="${LITELLM_MASTER_KEY:?LITELLM_MASTER_KEY is required}"

# Fetch existing DB model names once
EXISTING=$(curl -s "$LITELLM_URL/v2/model/info" \
  -H "Authorization: Bearer $LITELLM_MASTER_KEY" | \
  python3 -c "
import sys, json
data = json.load(sys.stdin)
for m in data.get('data', []):
    if m.get('model_info', {}).get('db_model'):
        print(m['model_name'])
" 2>/dev/null)

# Guard: skip entirely if already initialized (gpt-4o is a reliable sentinel).
# Set FORCE=1 to bypass this check and add any missing models.
if [ "${FORCE:-0}" != "1" ] && echo "$EXISTING" | grep -qx "gpt-4o"; then
  echo "Models already registered in DB. Set FORCE=1 to add missing models."
  exit 0
fi

add_model() {
  local model_name="$1"
  local litellm_model="$2"
  local api_key_env="$3"
  local api_base="${4:-}"

  if echo "$EXISTING" | grep -qx "$model_name"; then
    echo "  [skip] $model_name"
    return
  fi

  if [ -n "$api_base" ]; then
    base_json=",\"api_base\":\"$api_base\""
  else
    base_json=""
  fi

  local body="{\"model_name\":\"$model_name\",\"litellm_params\":{\"model\":\"$litellm_model\",\"api_key\":\"os.environ/$api_key_env\"$base_json}}"

  local status
  status=$(curl -s -o /tmp/model_resp.json -w "%{http_code}" \
    -X POST "$LITELLM_URL/model/new" \
    -H "Authorization: Bearer $LITELLM_MASTER_KEY" \
    -H "Content-Type: application/json" \
    -d "$body")

  if [ "$status" = "200" ] || [ "$status" = "201" ]; then
    echo "  [ok]   $model_name"
  else
    echo "  [fail] $model_name — HTTP $status: $(cat /tmp/model_resp.json)"
  fi
}

# GitHub Copilot models use token-file auth, no api_key needed.
add_copilot_model() {
  local model_name="$1"
  local copilot_model="$2"

  if echo "$EXISTING" | grep -qx "$model_name"; then
    echo "  [skip] $model_name"
    return
  fi

  local body="{\"model_name\":\"$model_name\",\"litellm_params\":{\"model\":\"github_copilot/$copilot_model\",\"extra_headers\":{\"editor-version\":\"vscode/1.85.1\",\"editor-plugin-version\":\"copilot/1.155.0\",\"Copilot-Integration-Id\":\"vscode-chat\",\"user-agent\":\"GithubCopilot/1.155.0\"}}}"

  local status
  status=$(curl -s -o /tmp/model_resp.json -w "%{http_code}" \
    -X POST "$LITELLM_URL/model/new" \
    -H "Authorization: Bearer $LITELLM_MASTER_KEY" \
    -H "Content-Type: application/json" \
    -d "$body")

  if [ "$status" = "200" ] || [ "$status" = "201" ]; then
    echo "  [ok]   $model_name"
  else
    echo "  [fail] $model_name — HTTP $status: $(cat /tmp/model_resp.json)"
  fi
}

echo "Registering models in LiteLLM DB..."

# OpenAI — GPT series
add_model "gpt-4o"           "openai/gpt-4o"           "OPENAI_API_KEY"
add_model "gpt-4o-mini"      "openai/gpt-4o-mini"      "OPENAI_API_KEY"
add_model "gpt-4.1"          "openai/gpt-4.1"          "OPENAI_API_KEY"
add_model "gpt-4.1-mini"     "openai/gpt-4.1-mini"     "OPENAI_API_KEY"
add_model "gpt-4.1-nano"     "openai/gpt-4.1-nano"     "OPENAI_API_KEY"
add_model "gpt-5"            "openai/gpt-5"            "OPENAI_API_KEY"
add_model "gpt-5-mini"       "openai/gpt-5-mini"       "OPENAI_API_KEY"
add_model "gpt-5-nano"       "openai/gpt-5-nano"       "OPENAI_API_KEY"
add_model "gpt-5-pro"        "openai/gpt-5-pro"        "OPENAI_API_KEY"
add_model "gpt-5.1"          "openai/gpt-5.1"          "OPENAI_API_KEY"
add_model "gpt-5.2"          "openai/gpt-5.2"          "OPENAI_API_KEY"
add_model "gpt-5.2-pro"      "openai/gpt-5.2-pro"      "OPENAI_API_KEY"
# OpenAI — o-series reasoning
add_model "o1"               "openai/o1"               "OPENAI_API_KEY"
add_model "o1-pro"           "openai/o1-pro"           "OPENAI_API_KEY"
add_model "o3"               "openai/o3"               "OPENAI_API_KEY"
add_model "o3-mini"          "openai/o3-mini"          "OPENAI_API_KEY"
add_model "o4-mini"          "openai/o4-mini"          "OPENAI_API_KEY"
# OpenAI — Codex series
add_model "gpt-5-codex"          "openai/gpt-5-codex"          "OPENAI_API_KEY"
add_model "gpt-5.1-codex"        "openai/gpt-5.1-codex"        "OPENAI_API_KEY"
add_model "gpt-5.1-codex-max"    "openai/gpt-5.1-codex-max"    "OPENAI_API_KEY"
add_model "gpt-5.1-codex-mini"   "openai/gpt-5.1-codex-mini"   "OPENAI_API_KEY"
add_model "gpt-5.2-codex"        "openai/gpt-5.2-codex"        "OPENAI_API_KEY"
add_model "gpt-5.3-codex"        "openai/gpt-5.3-codex"        "OPENAI_API_KEY"
add_model "gpt-5.3-codex-spark" "openai/gpt-5.3-codex-spark"  "OPENAI_API_KEY"
add_model "gpt-5.4"              "openai/gpt-5.4"              "OPENAI_API_KEY"

# Anthropic
add_model "claude-opus-4-6"    "anthropic/claude-opus-4-6"           "ANTHROPIC_API_KEY"
add_model "claude-sonnet-4-6"  "anthropic/claude-sonnet-4-6"         "ANTHROPIC_API_KEY"
add_model "claude-opus-4-5"    "anthropic/claude-opus-4-5-20251101"  "ANTHROPIC_API_KEY"
add_model "claude-opus-4-1"    "anthropic/claude-opus-4-1-20250805"  "ANTHROPIC_API_KEY"
add_model "claude-sonnet-4-5"  "anthropic/claude-sonnet-4-5-20250929" "ANTHROPIC_API_KEY"
add_model "claude-opus-4"      "anthropic/claude-opus-4-20250514"    "ANTHROPIC_API_KEY"
add_model "claude-sonnet-4"    "anthropic/claude-sonnet-4-20250514"  "ANTHROPIC_API_KEY"
add_model "claude-haiku-4-5"   "anthropic/claude-haiku-4-5-20251001" "ANTHROPIC_API_KEY"
add_model "claude-3-haiku"     "anthropic/claude-3-haiku-20240307"   "ANTHROPIC_API_KEY"

# Google Gemini
add_model "gemini-2.0-flash"       "gemini/gemini-2.0-flash"       "GEMINI_API_KEY"
add_model "gemini-2.0-flash-lite"  "gemini/gemini-2.0-flash-lite"  "GEMINI_API_KEY"
add_model "gemini-2.5-flash-lite"  "gemini/gemini-2.5-flash-lite"  "GEMINI_API_KEY"
add_model "gemini-2.5-flash"       "gemini/gemini-2.5-flash"       "GEMINI_API_KEY"
add_model "gemini-2.5-pro"         "gemini/gemini-2.5-pro"         "GEMINI_API_KEY"
add_model "gemini-3-flash-preview" "gemini/gemini-3-flash-preview"  "GEMINI_API_KEY"
add_model "gemini-3-pro-preview"   "gemini/gemini-3-pro-preview"    "GEMINI_API_KEY"
add_model "gemini-3.1-pro-preview" "gemini/gemini-3.1-pro-preview"  "GEMINI_API_KEY"
add_model "gemini-flash-latest"      "gemini/gemini-flash-latest"      "GEMINI_API_KEY"
add_model "gemini-flash-lite-latest" "gemini/gemini-flash-lite-latest" "GEMINI_API_KEY"
add_model "gemini-pro-latest"        "gemini/gemini-pro-latest"        "GEMINI_API_KEY"

# ZAI / GLM
add_model "zai-glm-4.5"     "openai/glm-4.5"     "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "zai-glm-4.5-air" "openai/glm-4.5-air" "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "zai-glm-4.6"     "openai/glm-4.6"     "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "zai-glm-4.7"     "openai/glm-4.7"       "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "zai-glm-5"       "openai/glm-5"         "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "glm-4.7-flash"   "openai/glm-4.7-flash" "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"
add_model "glm-5"           "openai/glm-5"         "ZAI_API_KEY" "https://api.z.ai/api/coding/paas/v4"

# GitHub Copilot (token-file auth, no API key)
add_copilot_model "copilot-gpt-4o"              "gpt-4o"
add_copilot_model "copilot-gpt-4.1"             "gpt-4.1"
add_copilot_model "copilot-gpt-5-mini"          "gpt-5-mini"
add_copilot_model "copilot-gpt-5.1"             "gpt-5.1"
add_copilot_model "copilot-gpt-5.2"             "gpt-5.2"
add_copilot_model "copilot-gpt-5.1-codex"       "gpt-5.1-codex"
add_copilot_model "copilot-gpt-5.1-codex-max"   "gpt-5.1-codex-max"
add_copilot_model "copilot-gpt-5.1-codex-mini"  "gpt-5.1-codex-mini"
add_copilot_model "copilot-gpt-5.2-codex"       "gpt-5.2-codex"
add_copilot_model "copilot-gpt-5.3-codex"       "gpt-5.3-codex"
add_copilot_model "copilot-claude-opus-4.6"      "claude-opus-4.6"
add_copilot_model "copilot-claude-opus-4.6-fast" "claude-opus-4.6-fast"
add_copilot_model "copilot-claude-sonnet-4.6"    "claude-sonnet-4.6"
add_copilot_model "copilot-claude-sonnet-4.5"    "claude-sonnet-4.5"
add_copilot_model "copilot-claude-sonnet-4"      "claude-sonnet-4"
add_copilot_model "copilot-claude-opus-4.5"      "claude-opus-4.5"
add_copilot_model "copilot-claude-haiku-4.5"     "claude-haiku-4.5"
add_copilot_model "copilot-gemini-2.5-pro"       "gemini-2.5-pro"
add_copilot_model "copilot-gemini-3-flash"       "gemini-3-flash-preview"
add_copilot_model "copilot-gemini-3-pro"         "gemini-3-pro-preview"
add_copilot_model "copilot-gemini-3.1-pro"       "gemini-3.1-pro-preview"
add_copilot_model "copilot-grok-code-fast"       "grok-code-fast-1"

# Local models (llama.cpp — no API key, custom model_info)
if ! echo "$EXISTING" | grep -qx "qwen2.5-14b-local"; then
  status=$(curl -s -o /tmp/model_resp.json -w "%{http_code}" \
    -X POST "$LITELLM_URL/model/new" \
    -H "Authorization: Bearer $LITELLM_MASTER_KEY" \
    -H "Content-Type: application/json" \
    -d '{
      "model_name": "qwen2.5-14b-local",
      "litellm_params": {
        "model": "openai/Qwen2.5-14B-Instruct-Q4_K_M.gguf",
        "api_base": "http://192.168.153.113:18806/v1"
      },
      "model_info": {
        "mode": "chat",
        "max_tokens": 8192,
        "max_input_tokens": 32768,
        "description": "Qwen2.5-14B-Instruct running locally via llama.cpp. Supports tool calling. Private/offline use.",
        "supports_function_calling": true,
        "supports_vision": false
      }
    }')
  if [ "$status" = "200" ] || [ "$status" = "201" ]; then
    echo "  [ok]   qwen2.5-14b-local"
  else
    echo "  [fail] qwen2.5-14b-local — HTTP $status: $(cat /tmp/model_resp.json)"
  fi
else
  echo "  [skip] qwen2.5-14b-local"
fi

echo "Done."