#!/bin/bash

# ==============================================================================
# Cluster Presentation Setup Script
# ==============================================================================

set -e

# Configuration
CLUSTER_HOME="$HOME/cluster"
COMPOSE_FILE="docker-compose.yml"
ENV_FILE=".env"
SAMPLE_ENV="sample.env"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${GREEN}Starting Cluster Setup Script...${NC}"

# 1. Root/Sudo Check
if [ "$EUID" -ne 0 ]; then
  echo -e "${YELLOW}This script may need to install packages. Please run with sudo or be prepared to enter your password.${NC}"
fi

# 2. Dependency Check
echo -e "${GREEN}Checking system dependencies...${NC}"
deps=(curl wget jq whiptail lspci python3 python3-yaml)
for dep in "${deps[@]}"; do
  if ! command -v "$dep" &> /dev/null && [ "$dep" != "python3-yaml" ]; then
    echo -e "${YELLOW}Installing missing dependency: $dep${NC}"
    sudo apt-get update -y -qq && sudo apt-get install -y -qq "$dep"
  fi
done

# Special check for python3-yaml which doesn't have a command
if ! python3 -c "import yaml" &> /dev/null; then
    echo -e "${YELLOW}Installing missing dependency: python3-yaml${NC}"
    sudo apt-get update -y -qq && sudo apt-get install -y -qq python3-yaml
fi

# 3. Docker Installation
if ! command -v docker &> /dev/null; then
  echo -e "${YELLOW}Docker not found. Installing Docker...${NC}"
  curl -fsSL https://get.docker.com -o get-docker.sh
  sudo sh get-docker.sh
  sudo usermod -aG docker "$USER"
  echo -e "${GREEN}Docker installed. You may need to log out and back in for group changes to take effect.${NC}"
else
  echo -e "${GREEN}Docker is already installed.${NC}"
fi

# 4. GPU Detection
GPU_TYPE="NONE"
if lspci | grep -iq nvidia; then
  GPU_TYPE="NVIDIA"
  echo -e "${GREEN}NVIDIA GPU detected.${NC}"
elif lspci | grep -iq "amd/ati"; then
  GPU_TYPE="AMD"
  echo -e "${GREEN}AMD GPU detected.${NC}"
elif lspci | grep -iq "intel.*graphics"; then
  GPU_TYPE="INTEL"
  echo -e "${GREEN}Intel GPU detected.${NC}"
else
  echo -e "${YELLOW}No dedicated GPU detected. Defaulting to CPU mode.${NC}"
fi

# 5. Hardware-specific Toolkit Setup
if [ "$GPU_TYPE" == "NVIDIA" ]; then
  if ! command -v nvidia-smi &> /dev/null || ! docker info | grep -q "Runtimes:.*nvidia"; then
    if whiptail --title "GPU Toolkit" --yesno "NVIDIA GPU detected but drivers or container toolkit missing. Attempt automatic installation?" 10 60; then
       echo -e "${GREEN}Installing NVIDIA Drivers and Container Toolkit...${NC}"
       sudo apt-get install -y nvidia-driver-535 nvidia-container-toolkit
       sudo nvidia-ctk runtime configure --runtime=docker
       sudo systemctl restart docker
    fi
  fi
fi

# 6. Service Selection
echo -e "${GREEN}Opening service selection menu...${NC}"
CHOICES=$(whiptail --title "Service Selection" --checklist \
"Select the stacks/services you want to install:" 20 75 10 \
"NEXTCLOUD" "Full Nextcloud Suite" ON \
"AI" "Ollama, Whisper, OpenWebUI" ON \
"AUTOMATION" "Home Assistant, ESPHome" OFF \
"PRODUCTIVITY" "Obsidian Indexer, n8n, Syncthing" OFF \
"MEDIA" "Immich Photo Library" OFF \
"DOCUMENTS" "Paperless-ngx (with AI)" OFF \
"MAPS" "OpenStreetMap Stack" OFF \
"MONITORING" "Uptime Kuma, Cloudflared" ON 3>&1 1>&2 2>&3)

# 7. Create Cluster Directory
mkdir -p "$CLUSTER_HOME"
cp "$COMPOSE_FILE" "$CLUSTER_HOME/$COMPOSE_FILE"
cp "$SAMPLE_ENV" "$CLUSTER_HOME/$ENV_FILE"

# 8. Process Compose File (Comment/Uncomment based on selection and hardware)
cd "$CLUSTER_HOME"

echo -e "${GREEN}Configuring hardware and service selection in $COMPOSE_FILE...${NC}"

# Python helper for cleaner YAML processing
python3 - <<EOF
import sys
import yaml

choices = "$CHOICES"
gpu_type = "$GPU_TYPE"

with open("$COMPOSE_FILE", 'r') as f:
    data = yaml.safe_load(f)

# Service Stack Mapping
stacks = {
    "NEXTCLOUD": ["nc_db", "nc_redis", "nextcloud", "nc_cron", "nc_collabora", "nc_coturn", "nc_janus", "nc_nats", "nextcloud-talk-signaling", "nextcloud-push", "nextcloud-imaginary", "nextcloud-search", "nextcloud-clamav", "nextcloud-appapi-harp"],
    "AI": ["ollama", "whisper", "open-webui", "qdrant", "comfyui", "kokoro"],
    "AUTOMATION": ["homeassistant", "wyze-bridge", "esphome"],
    "PRODUCTIVITY": ["homebox", "kiwix", "n8n", "obsidian", "obsidian-indexer", "syncthing"],
    "MEDIA": ["immich-server", "immich-redis", "immich-database", "immich-machine-learning"],
    "DOCUMENTS": ["paperless-webserver", "paperless-db", "paperless-broker", "paperless-ai", "gotenberg", "tika"],
    "MAPS": ["map-osrm-prep", "map-osrm", "map-nominatim", "map-tiles"],
    "MONITORING": ["uptime-kuma", "cloudflared", "cloudflare-ddns"]
}

# Remove non-selected services
services_to_keep = []
for stack, services in stacks.items():
    if stack in choices:
        services_to_keep.extend(services)

all_services = list(data['services'].keys())
for service in all_services:
    if service not in services_to_keep:
        del data['services'][service]

# Hardware adjustments for remaining services
for name, svc in data['services'].items():
    if gpu_type != "NVIDIA":
        if 'runtime' in svc and svc['runtime'] == 'nvidia':
            del svc['runtime']
        if 'deploy' in svc and 'resources' in svc['deploy']:
            # Potentially keep but adjust for AMD/CPU
            # For simplicity in this script, we'll just strip the NVIDIA-specific reservation
            del svc['deploy']

with open("$COMPOSE_FILE", 'w') as f:
    yaml.dump(data, f, sort_keys=False)
EOF

# 9. Directory Building
echo -e "${GREEN}Building directory structure...${NC}"
# Extract volumes from compose file and create them
grep -oP '(?<=\- \$\{PROJECT_ROOT:-\.\}/)[^:]+' "$COMPOSE_FILE" | xargs -I {} mkdir -p "$CLUSTER_HOME/{}"

# 10. Maps Initialization
if [[ $CHOICES == *"MAPS"* ]]; then
  MAP_URL=$(whiptail --title "Maps Setup" --inputbox "Enter the URL for the OSM .pbf file (e.g. from Geofabrik):" 10 60 3>&1 1>&2 2>&3)
  if [ ! -z "$MAP_URL" ]; then
    mkdir -p "$CLUSTER_HOME/osm-stack/data"
    echo -e "${GREEN}Downloading map data...${NC}"
    wget -O "$CLUSTER_HOME/osm-stack/data/region.osm.pbf" "$MAP_URL"
  fi
fi

# 11. Ollama Models
if [[ $CHOICES == *"AI"* ]]; then
  MODELS=$(whiptail --title "Ollama Models" --inputbox "Enter LLM models to pull (space separated, e.g. llama3 mistral):" 10 60 "llama3 nomic-embed-text" 3>&1 1>&2 2>&3)
  # We'll handle pulling these after start
fi

# 12. Permissions
echo -e "${GREEN}Setting permissions...${NC}"
sudo chown -R "$USER:$USER" "$CLUSTER_HOME"
# Ensure docker socket is accessible
sudo chmod 666 /var/run/docker.sock || true

echo -e "${GREEN}Setup Complete!${NC}"
echo -e "Your cluster is ready in: ${YELLOW}$CLUSTER_HOME${NC}"
echo -e "1. Edit ${YELLOW}$CLUSTER_HOME/$ENV_FILE${NC} with your secrets."
if [ ! -z "$MODELS" ]; then
  echo -e "2. After starting, run: ${YELLOW}docker exec -it ollama ollama pull $MODELS${NC}"
fi
echo -e "3. Run: ${YELLOW}cd $CLUSTER_HOME && docker compose up -d${NC}"