#!/usr/bin/env bash
# detect_capabilities.sh
# Detects NVIDIA GPU presence and configures .env for Docker Compose capabilities.

gpu_detected=false

# 1. Check if nvidia-smi works
if command -v nvidia-smi &> /dev/null; then
    if nvidia-smi &> /dev/null; then
        gpu_detected=true
    fi
fi

# 2. Check using lspci if nvidia-smi is not found
if [ "$gpu_detected" = false ] && command -v lspci &> /dev/null; then
    if lspci | grep -qi nvidia; then
        gpu_detected=true
    fi
fi

# 3. Check macOS System Profiler (for metal/Apple silicon or AMD/Nvidia eGPUs)
if [ "$gpu_detected" = false ] && command -v system_profiler &> /dev/null; then
    if system_profiler SPDisplaysDataType | grep -qi nvidia; then
        gpu_detected=true
    fi
fi

echo "--- CoReason Hardware Capability Detection ---"
if [ "$gpu_detected" = true ]; then
    echo -e "\033[0;32m[INFO] NVIDIA GPU detected. Configuring environment for GPU acceleration.\033[0m"
    profiles="gpu"
    extras="inference"
    sglang_url="http://sglang:30000"
else
    echo -e "\033[0;33m[INFO] No NVIDIA GPU detected. Configuring environment for CPU-only execution.\033[0m"
    profiles=""
    extras=""
    sglang_url=""
fi

# Paths
script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
env_example_path="${script_dir}/../.env.example"
env_path="${script_dir}/../.env"

# Load base content
if [ -f "$env_path" ]; then
    echo "[INFO] Updating existing .env file."
    content_source="$env_path"
elif [ -f "$env_example_path" ]; then
    echo "[INFO] Creating .env file from .env.example."
    content_source="$env_example_path"
else
    echo -e "\033[0;31m[ERROR] Neither .env nor .env.example found!\033[0m"
    exit 1
fi

# Read and write updated file using awk/sed or a simple loop
temp_env=$(mktemp)

# Keys to set
keys=("COMPOSE_PROFILES" "EXTRAS" "SGLANG_URL")
declare -A values
values["COMPOSE_PROFILES"]="$profiles"
values["EXTRAS"]="$extras"
values["SGLANG_URL"]="$sglang_url"

# Process source file
while IFS= read -r line || [[ -n "$line" ]]; do
    matched=false
    for key in "${keys[@]}"; do
        if [[ "$line" =~ ^$key[[:space:]]*= ]]; then
            echo "$key=${values[$key]}" >> "$temp_env"
            values["$key"]="__PROCESSED__"
            matched=true
            break
        fi
    done
    if [ "$matched" = false ]; then
        echo "$line" >> "$temp_env"
    fi
done < "$content_source"

# Append any keys that weren't present/updated
for key in "${keys[@]}"; do
    if [ "${values[$key]}" != "__PROCESSED__" ]; then
        echo "$key=${values[$key]}" >> "$temp_env"
    fi
done

mv "$temp_env" "$env_path"
chmod 644 "$env_path"
echo -e "\033[0;32m[SUCCESS] Updated .env file at $env_path successfully.\033[0m"