#!/bin/bash
set -eu

# ============================================================================
# LiteLLM Cloudron Startup Script
# Configures LiteLLM with Cloudron's PostgreSQL and Redis addons
# ============================================================================

echo "=== LiteLLM Cloudron Startup ==="

# --- Cloudron PostgreSQL addon provides:
#   CLOUDRON_POSTGRESQL_URL - postgresql://user:pass@host:port/dbname
#   CLOUDRON_POSTGRESQL_USERNAME
#   CLOUDRON_POSTGRESQL_PASSWORD
#   CLOUDRON_POSTGRESQL_HOST
#   CLOUDRON_POSTGRESQL_PORT
#   CLOUDRON_POSTGRESQL_DATABASE

# --- Cloudron Redis addon provides:
#   CLOUDRON_REDIS_URL - redis://user:pass@host:port
#   CLOUDRON_REDIS_USERNAME
#   CLOUDRON_REDIS_PASSWORD
#   CLOUDRON_REDIS_HOST
#   CLOUDRON_REDIS_PORT

# --- Cloudron provides:
#   CLOUDRON_APP_DOMAIN - the app's domain (e.g. gateway.inference.coop)
#   CLOUDRON_API_ORIGIN - origin for API calls
#   CLOUDRON_WEB_ORIGIN - origin for web UI

# --- LiteLLM required env vars ---

# Generate a master key if not already in /app/data
MASTER_KEY_FILE="/app/data/.master_key"
if [ ! -f "$MASTER_KEY_FILE" ]; then
    echo "Generating master key..."
    python3 -c "import secrets; print('sk-' + secrets.token_urlsafe(32))" > "$MASTER_KEY_FILE"
fi
export LITELLM_MASTER_KEY=$(cat "$MASTER_KEY_FILE")

# Generate a salt key if not already in /app/data
SALT_KEY_FILE="/app/data/.salt_key"
if [ ! -f "$SALT_KEY_FILE" ]; then
    echo "Generating salt key..."
    python3 -c "import secrets; print(secrets.token_urlsafe(32))" > "$SALT_KEY_FILE"
fi
export LITELLM_SALT_KEY=$(cat "$SALT_KEY_FILE")

# Use Cloudron's PostgreSQL
export DATABASE_URL="${CLOUDRON_POSTGRESQL_URL}"

# Use Cloudron's Redis (for rate limiting, caching)
export REDIS_HOST="${CLOUDRON_REDIS_HOST}"
export REDIS_PORT="${CLOUDRON_REDIS_PORT}"
export REDIS_PASSWORD="${CLOUDRON_REDIS_PASSWORD}"

# Store models in DB (manage from Admin UI)
export STORE_MODEL_IN_DB="True"

# Don't run schema migrations on every start — let LiteLLM handle it
# (for single-instance deployment this is fine)
export DISABLE_SCHEMA_UPDATE="false"

# UI configuration — use Cloudron's domain
export UI_USERNAME="admin"
export UI_PASSWORD="${LITELLM_MASTER_KEY}"

# Proxy settings for Cloudron's reverse proxy
export LITELLM_PROXY_BASE_URL="https://${CLOUDRON_APP_DOMAIN}"

# Allow requests from Cloudron's Open WebUI, LobeChat, etc.
export LITELLM_CORS_ALLOWED_ORIGINS="${CLOUDRON_WEB_ORIGIN:-*}"

echo "PostgreSQL: ${CLOUDRON_POSTGRESQL_HOST}:${CLOUDRON_POSTGRESQL_PORT}"
echo "Redis: ${CLOUDRON_REDIS_HOST}:${CLOUDRON_REDIS_PORT}"
echo "Domain: ${CLOUDRON_APP_DOMAIN}"

# --- Write config.yaml with Cloudron provider keys if set ---

CONFIG_FILE="/app/data/config.yaml"

if [ ! -f "$CONFIG_FILE" ] || [ ! -s "$CONFIG_FILE" ]; then
    echo "Writing default config.yaml..."
    cat > "$CONFIG_FILE" << 'YAML'
# LiteLLM config for inference.coop
# MVP config: Tinfoil only (TEE-protected, privacy-first)
# Two models: DeepSeek V4 Flash (default, best agents) + GPT-OSS 120B (cheapest)
# Provider key set via environment variable or Admin UI

model_list:
  # DeepSeek V4 Flash — default model
  - model_name: deepseek-v4-flash
    litellm_params:
      model: openai/deepseek-v4-flash
      api_base: https://api.tinfoil.sh/v1
      api_key: os.environ/TINFOIL_API_KEY

  # GPT-OSS 120B — lightweight fallback
  - model_name: gpt-oss-120b
    litellm_params:
      model: openai/gpt-oss-120b
      api_base: https://api.tinfoil.sh/v1
      api_key: os.environ/TINFOIL_API_KEY

# Fallback: if DeepSeek is down, use GPT-OSS
router_settings:
  num_retries: 2
  timeout: 30
  fallbacks:
    - deepseek-v4-flash:
        - gpt-oss-120b

general_settings:
  master_key: os.environ/LITELLM_MASTER_KEY
  database_url: os.environ/DATABASE_URL
  store_model_in_db: true

litellm_settings:
  salt_key: os.environ/LITELLM_SALT_KEY
  drop_params: true
  num_threads: 4
  request_timeout: 30
YAML
    echo "Config written to $CONFIG_FILE"
fi

# --- Start LiteLLM ---
echo "Starting LiteLLM on port 4000..."

# Start the Tinfoil proxy as a local sidecar. It verifies the enclave
# attestation and encrypts request/response bodies with EHBP (HPKE), so
# LiteLLM can talk plaintext OpenAI to it while the proxy handles the
# end-to-end encryption to the Tinfoil enclave.
echo "Starting Tinfoil proxy on 127.0.0.1:3301..."
/app/code/tinfoil-proxy -b 127.0.0.1 -p 3301 &
TINFOIL_PROXY_PID=$!
echo "Tinfoil proxy PID: $TINFOIL_PROXY_PID"

# Give the proxy a moment to verify the enclave attestation
sleep 3

exec litellm --config "$CONFIG_FILE" --port 4000 --host 0.0.0.0