#!/bin/bash
# DupFinder first-time setup — configure paths, pull image, write override
set -e

CONF_FILE="/etc/dupfinder.conf"
COMPOSE_DIR="/opt/dupfinder"
OVERRIDE_YML="$COMPOSE_DIR/docker-compose.override.yml"
IMAGE_NAME="tocmo0nlord/dupfinder:latest"
DATA_DIR="/var/lib/dupfinder/data"
APP_PORT=8765

RED='\033[0;31m'; GREEN='\033[0;32m'; CYAN='\033[0;36m'; NC='\033[0m'
info()  { echo -e "${CYAN}==> $*${NC}"; }
ok()    { echo -e "${GREEN}    OK  $*${NC}"; }
err()   { echo -e "${RED}    !!  $*${NC}"; }

# ── Root check ────────────────────────────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    err "Please run as root: sudo dupfinder setup"
    exit 1
fi

echo ""
echo "  ╔══════════════════════════════════════╗"
echo "  ║     DupFinder Setup                  ║"
echo "  ╚══════════════════════════════════════╝"
echo ""

# ── Load existing config as defaults ─────────────────────────────────────────
[[ -f "$CONF_FILE" ]] && source "$CONF_FILE"
: "${PHOTOS_PATH:=/mnt/photos}"
: "${DATA_PATH:=$DATA_DIR}"
: "${APP_PORT:=8765}"

# ── Check Docker ──────────────────────────────────────────────────────────────
info "Checking Docker..."
if ! command -v docker &>/dev/null; then
    err "Docker is not installed."
    echo "    Install with:  curl -fsSL https://get.docker.com | sh"
    exit 1
fi
if ! docker info &>/dev/null; then
    err "Docker daemon is not running."
    echo "    Start with:  sudo systemctl start docker"
    exit 1
fi
ok "Docker is running"

# ── Check docker compose ──────────────────────────────────────────────────────
if ! docker compose version &>/dev/null; then
    err "docker compose (V2 plugin) not found. Update Docker or install docker-compose-plugin."
    exit 1
fi
ok "docker compose V2 available"

# ── Check NVIDIA GPU + container toolkit ─────────────────────────────────────
info "Checking GPU..."
if command -v nvidia-smi &>/dev/null && nvidia-smi &>/dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1)
    ok "NVIDIA GPU detected: $GPU_NAME"

    # nvidia-container-toolkit is required for Docker GPU passthrough
    if ! command -v nvidia-ctk &>/dev/null && ! dpkg -l nvidia-container-toolkit &>/dev/null 2>&1; then
        echo ""
        echo "    nvidia-container-toolkit is not installed."
        echo "    Without it Docker cannot pass the GPU to the container."
        echo "    Install with:"
        echo ""
        echo "      curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
        echo "      curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list"
        echo "      sudo apt update && sudo apt install -y nvidia-container-toolkit"
        echo "      sudo nvidia-ctk runtime configure --runtime=docker"
        echo "      sudo systemctl restart docker"
        echo ""
        read -rp "    Install nvidia-container-toolkit now? (Y/n): " INST_CTK
        if [[ "$INST_CTK" != "n" && "$INST_CTK" != "N" ]]; then
            curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
                | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
            curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
                | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
                | tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
            apt update -qq && apt install -y nvidia-container-toolkit
            nvidia-ctk runtime configure --runtime=docker
            systemctl restart docker
            ok "nvidia-container-toolkit installed and Docker restarted"
        else
            echo "    Skipping — GPU will not be available in Docker. You can re-run setup later."
            GPU_AVAILABLE=false
        fi
    else
        ok "nvidia-container-toolkit is present"
    fi

    [[ "$GPU_AVAILABLE" != "false" ]] && GPU_AVAILABLE=true
else
    echo "    No NVIDIA GPU detected — will use CPU for perceptual hashing"
    GPU_AVAILABLE=false
fi

# ── Photos path ───────────────────────────────────────────────────────────────
echo ""
info "Photos library path (mounted read-only):"
echo "    Current: $PHOTOS_PATH"
read -rp "    Path [Enter to keep]: " INPUT
INPUT="${INPUT%\"}" ; INPUT="${INPUT#\"}"   # strip quotes
[[ -n "$INPUT" ]] && PHOTOS_PATH="$INPUT"

if [[ ! -d "$PHOTOS_PATH" ]]; then
    err "Path not found: $PHOTOS_PATH"
    echo "    Create it or mount your drive first, then re-run setup."
    exit 1
fi
ok "Photos: $PHOTOS_PATH"

# ── Data path ─────────────────────────────────────────────────────────────────
echo ""
info "Database storage path:"
echo "    Current: $DATA_PATH"
read -rp "    Path [Enter to keep]: " INPUT
INPUT="${INPUT%\"}" ; INPUT="${INPUT#\"}"
[[ -n "$INPUT" ]] && DATA_PATH="$INPUT"
mkdir -p "$DATA_PATH"
ok "Data:   $DATA_PATH"

# ── Port ──────────────────────────────────────────────────────────────────────
echo ""
read -rp "    Web port [$APP_PORT]: " INPUT
[[ -n "$INPUT" ]] && APP_PORT="$INPUT"
ok "Port:   $APP_PORT"

# ── Build (or pull) Docker image ──────────────────────────────────────────────
# The .deb ships the full source tree, so building locally is the default.
# Registry pull is tried only as a quick path if the image happens to be
# published; failures are silent.
echo ""
info "Preparing Docker image ($IMAGE_NAME)..."

if docker image inspect "$IMAGE_NAME" >/dev/null 2>&1; then
    ok "Image already present locally"
elif docker pull "$IMAGE_NAME" >/dev/null 2>&1; then
    ok "Image pulled from registry"
elif [[ -f "$COMPOSE_DIR/source/Dockerfile" ]]; then
    echo "    Building image from bundled source (one-time, ~5-10 min)..."
    docker build -t "$IMAGE_NAME" "$COMPOSE_DIR/source"
    ok "Image built from source"
else
    err "No image available and no source bundled. Reinstall the .deb."
    exit 1
fi

# ── Write config + override ───────────────────────────────────────────────────
info "Writing configuration..."

cat > "$CONF_FILE" <<EOF
PHOTOS_PATH=$PHOTOS_PATH
DATA_PATH=$DATA_PATH
APP_PORT=$APP_PORT
GPU_AVAILABLE=$GPU_AVAILABLE
EOF
chmod 600 "$CONF_FILE"

# Docker requires forward slashes
PHOTOS_DOCKER="${PHOTOS_PATH//\\//}"
DATA_DOCKER="${DATA_PATH//\\//}"

cat > "$OVERRIDE_YML" <<EOF
services:
  dup-finder:
    image: $IMAGE_NAME
    ports:
      - "${APP_PORT}:8000"
    volumes:
      - "$PHOTOS_DOCKER:/photos:ro"
      - "$DATA_DOCKER:/data"
EOF

# Add GPU reservation if available
if [[ "$GPU_AVAILABLE" == "true" ]]; then
    cat >> "$OVERRIDE_YML" <<EOF
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
EOF
fi
ok "Config saved to $CONF_FILE"

# ── Start service ─────────────────────────────────────────────────────────────
info "Starting DupFinder..."
systemctl daemon-reload
systemctl enable --now dupfinder.service
ok "Service started"

echo ""
echo -e "${GREEN}  ╔══════════════════════════════════════════╗${NC}"
echo -e "${GREEN}  ║  DupFinder is running!                   ║${NC}"
echo -e "${GREEN}  ║  Open: http://localhost:$APP_PORT           ║${NC}"
echo -e "${GREEN}  ╚══════════════════════════════════════════╝${NC}"
echo ""