#!/usr/bin/env bash
#
# Docker Compose Startup Script for LLM Inference
#
# Usage: ./deploy/docker-start.sh [OPTIONS]
#
# Options:
#   --profile <api|vllm|llamacpp|full>  Docker compose profile
#   --detach                             Run in detached mode
#   --down                               Stop and remove containers
#   --logs                               Show logs
#   --help                               Show this help message
#
# Required: Set environment variables in .env file or export them before running.
# See .env.example for the full list of required variables.

set -euo pipefail

# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Load .env file if it exists (check deploy/ first, then module root)
ENV_FILE=""
if [[ -f "$SCRIPT_DIR/.env" ]]; then
    ENV_FILE="$SCRIPT_DIR/.env"
elif [[ -f "$SCRIPT_DIR/../.env" ]]; then
    ENV_FILE="$SCRIPT_DIR/../.env"
fi

if [[ -n "$ENV_FILE" ]]; then
    echo "Loading environment from: $ENV_FILE"
    set -a
    source "$ENV_FILE"
    set +a
fi

PROFILE=""
DETACH=""
ACTION="up"

show_help() {
    sed -n '2,16p' "$0" | sed 's/^# //' | sed 's/^#//'
    exit 0
}

check_required_var() {
    local var_name="$1"
    if [[ -z "${!var_name:-}" ]]; then
        echo "ERROR: Required environment variable $var_name is not set"
        echo "Set it in .env file or export it before running this script"
        exit 1
    fi
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --profile)
            PROFILE="$2"
            shift 2
            ;;
        --detach|-d)
            DETACH="-d"
            shift
            ;;
        --down)
            ACTION="down"
            shift
            ;;
        --logs)
            ACTION="logs"
            shift
            ;;
        --help|-h)
            show_help
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# Require profile to be specified
if [[ -z "$PROFILE" && "$ACTION" == "up" ]]; then
    echo "ERROR: --profile is required"
    echo "Options: api, vllm, llamacpp, full"
    exit 1
fi

# Check required variables based on profile
check_required_var "LLM_DEFAULT_BACKEND"
check_required_var "LLM_ENABLE_VLLM"
check_required_var "LLM_ENABLE_LLAMACPP"
check_required_var "LLM_EXTERNAL_URL"

# API keys - at least one MUST be set for litellm backend (fail-fast)
if [[ "$LLM_DEFAULT_BACKEND" == "litellm" ]]; then
    if [[ -z "${OPENROUTER_API_KEY:-}" && -z "${OPENAI_API_KEY:-}" && -z "${ANTHROPIC_API_KEY:-}" ]]; then
        echo "ERROR: litellm backend requires at least one API key"
        echo "Set one of: OPENROUTER_API_KEY, OPENAI_API_KEY, or ANTHROPIC_API_KEY"
        exit 1
    fi
fi

# Check model-specific variables for vllm/llamacpp profiles
case $PROFILE in
    vllm|full)
        check_required_var "MODELS_DIR"
        check_required_var "VLLM_MODEL"
        ;;
    llamacpp|full)
        check_required_var "MODELS_DIR"
        check_required_var "LLAMACPP_MODEL"
        check_required_var "LLAMACPP_THREADS"
        check_required_var "LLAMACPP_PARALLEL"
        ;;
esac

cd "$SCRIPT_DIR"

case $ACTION in
    up)
        echo "Starting LLM Inference with profile: $PROFILE"
        echo "  Default backend: $LLM_DEFAULT_BACKEND"
        echo "  vLLM enabled: $LLM_ENABLE_VLLM"
        echo "  llama.cpp enabled: $LLM_ENABLE_LLAMACPP"
        if [[ "$PROFILE" == "vllm" || "$PROFILE" == "full" ]]; then
            echo "  vLLM model: $VLLM_MODEL"
        fi
        if [[ "$PROFILE" == "llamacpp" || "$PROFILE" == "full" ]]; then
            echo "  llama.cpp model: $LLAMACPP_MODEL"
        fi
        echo ""
        # shellcheck disable=SC2086
        exec docker compose --profile "$PROFILE" up $DETACH
        ;;
    down)
        if [[ -z "$PROFILE" ]]; then
            echo "ERROR: --profile is required with --down"
            exit 1
        fi
        echo "Stopping LLM Inference containers..."
        exec docker compose --profile "$PROFILE" down
        ;;
    logs)
        if [[ -z "$PROFILE" ]]; then
            echo "ERROR: --profile is required with --logs"
            exit 1
        fi
        exec docker compose --profile "$PROFILE" logs -f
        ;;
esac