#!/usr/bin/env bash
# Updates vllm/lmstudio/ollama model listings in opencode.json
# by querying their /v1/models endpoints.
#
# Usage:
#   ./update-models.sh                 # update configured defaults
#   ./update-models.sh vllm            # update only vllm
#   ./update-models.sh lmstudio        # update only lmstudio
#   ./update-models.sh ollama          # update only ollama
#
# Requires: jq, curl

set -uo pipefail

CONFIG="${OPENCODE_CONFIG:-${HOME}/.config/opencode/opencode.json}"

if [[ ! -f "$CONFIG" ]]; then
  echo "Error: config not found at $CONFIG" >&2
  exit 1
fi

if ! command -v jq &>/dev/null; then
  echo "Error: jq is required. Install with: brew install jq" >&2
  exit 1
fi

# Map provider name -> base URL (read from config)
get_base_url() {
  jq -r ".provider.${1}.options.baseURL // empty" "$CONFIG"
}

# Map provider name -> display label (read from config)
get_display_name() {
  jq -r ".provider.${1}.name // \"${1}\"" "$CONFIG"
}

# Emit the default opencode limits for locally served OpenAI-compatible models.
# The output limit must stay below the context window so auto-compaction can
# reserve space before the upstream API rejects input_tokens + max_tokens.
model_limit_json() {
  local provider="$1"
  local model_id="$2"

  case "${provider}:${model_id}" in
    *) jq -cn '{context: 262144, output: 32000}' ;;
  esac
}

# Fetch models from API endpoint and build the opencode models object.
# Preserves any existing entries that have extra metadata (limits, modalities, variants)
# beyond just a name — so your manual overrides aren't lost.
build_models_json() {
  local provider="$1"
  local base_url="$2"
  local display_name="$3"

  local api_response
  api_response=$(curl -sf --connect-timeout 5 "${base_url}/models" 2>/dev/null) || {
    echo "Error: could not reach ${base_url}/models — is ${provider} running?" >&2
    return 1
  }

  # Get list of model IDs from API
  local model_ids
  model_ids=$(echo "$api_response" | jq -r '.data[].id' | sort)

  if [[ -z "$model_ids" ]]; then
    echo "Warning: no models returned from ${provider}" >&2
    return 1
  fi

  # Get existing models object (to preserve manual overrides)
  local existing
  existing=$(jq ".provider.${provider}.models // {}" "$CONFIG")

  # Build new models object
  local new_models="{}"
  while IFS= read -r model_id; do
    [[ -z "$model_id" ]] && continue

    # Check if this model has extra config beyond just "name"
    local existing_entry
    existing_entry=$(echo "$existing" | jq -r --arg id "$model_id" '.[$id] // empty')

    local default_limit
    default_limit=$(model_limit_json "$provider" "$model_id")

    local has_extras="false"
    if [[ -n "$existing_entry" ]]; then
      local key_count
      key_count=$(echo "$existing_entry" | jq 'keys | length')
      if [[ "$key_count" -gt 1 ]]; then
        has_extras="true"
      fi
    fi

    if [[ "$has_extras" == "true" ]]; then
      # Preserve existing entry with overrides, just update the name
      new_models=$(echo "$new_models" | jq \
        --arg id "$model_id" \
        --argjson entry "$existing_entry" \
        --argjson limit "$default_limit" \
        --arg name "${model_id} (${display_name})" \
        '.[$id] = ($entry | .name = $name | .limit = (.limit // $limit))')
    else
      # Simple entry with just a name
      new_models=$(echo "$new_models" | jq \
        --arg id "$model_id" \
        --argjson limit "$default_limit" \
        --arg name "${model_id} (${display_name})" \
        '.[$id] = { "name": $name, "limit": $limit }')
    fi
  done <<< "$model_ids"

  echo "$new_models"
}

update_provider() {
  local provider="$1"

  # Check provider exists in config
  if ! jq -e ".provider.${provider}" "$CONFIG" &>/dev/null; then
    echo "Skipping ${provider}: not configured in opencode.json" >&2
    return 0
  fi

  local base_url
  base_url=$(get_base_url "$provider")
  if [[ -z "$base_url" ]]; then
    echo "Skipping ${provider}: no baseURL configured" >&2
    return 0
  fi

  local display_name
  display_name=$(get_display_name "$provider")

  echo "Updating ${provider} models from ${base_url}..."

  local new_models
  new_models=$(build_models_json "$provider" "$base_url" "$display_name") || return 1

  local count
  count=$(echo "$new_models" | jq 'keys | length')

  # Update config in-place
  local tmp
  tmp=$(mktemp)
  jq --argjson models "$new_models" \
    ".provider.${provider}.models = \$models" \
    "$CONFIG" > "$tmp" && mv "$tmp" "$CONFIG"

  echo "  ✓ ${provider}: ${count} models synced"
}

# --- Main ---

targets=("${@:-vllm ollama}")
errors=0

for target in $targets; do
  update_provider "$target" || ((errors++))
done

if [[ "$errors" -gt 0 ]]; then
  echo "Done with ${errors} error(s). Run 'opencode models' to verify."
  exit 1
else
  echo "Done. Run 'opencode models' to verify."
fi