#!/usr/bin/env bash
# ── setup.sh — AI Video Upscaler environment setup ──────────────────────────
#
# Usage:
#     chmod +x setup.sh && ./setup.sh
#
# Automates first-time setup on a new system:
#   • Verifies required system tools (python3, ffmpeg, ffprobe)
#   • Warns if Vulkan is unavailable
#   • Creates a Python virtual environment and installs dependencies
#   • Makes the Real-ESRGAN binary executable
#   • Creates working directories (uploads, outputs, temp)
# ─────────────────────────────────────────────────────────────────────────────

set -euo pipefail

# ── Colors & Helpers ─────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
RESET='\033[0m'

info()  { echo -e "${CYAN}ℹ  $*${RESET}"; }
ok()    { echo -e "${GREEN}✅ $*${RESET}"; }
warn()  { echo -e "${YELLOW}⚠  $*${RESET}"; }
fail()  { echo -e "${RED}❌ $*${RESET}"; exit 1; }
step()  { echo -e "\n${BOLD}── $* ──${RESET}"; }

# ── Banner ───────────────────────────────────────────────────────────────────
echo ""
echo -e "${BOLD}🔧 AI Video Upscaler — Setup${RESET}"
echo ""

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

# ── 1. Check: python3 ───────────────────────────────────────────────────────
step "Checking dependencies"

if ! command -v python3 &>/dev/null; then
    fail "python3 is not installed.\n   Install it with your package manager (e.g. apt install python3, dnf install python3)."
fi
ok "python3 found — $(python3 --version)"

# ── 2. Check: ffmpeg ────────────────────────────────────────────────────────
if ! command -v ffmpeg &>/dev/null; then
    fail "ffmpeg is not installed.\n   Install it with your package manager:\n     • Debian/Ubuntu : sudo apt install ffmpeg\n     • Fedora/RHEL   : sudo dnf install ffmpeg\n     • Arch          : sudo pacman -S ffmpeg"
fi
ok "ffmpeg found — $(ffmpeg -version | head -1)"

# ── 3. Check: ffprobe ───────────────────────────────────────────────────────
if ! command -v ffprobe &>/dev/null; then
    fail "ffprobe is not installed (usually bundled with ffmpeg).\n   Reinstall ffmpeg to get it."
fi
ok "ffprobe found"

# ── 4. Check: vulkaninfo (optional) ─────────────────────────────────────────
if ! command -v vulkaninfo &>/dev/null; then
    warn "vulkaninfo not found — Vulkan GPU acceleration is recommended but not required."
    warn "Install vulkan-tools for GPU-accelerated upscaling."
else
    ok "vulkaninfo found — Vulkan support detected"
fi

# ── 5. Create Python venv ───────────────────────────────────────────────────
step "Setting up Python virtual environment"

if [ -d "$SCRIPT_DIR/venv" ]; then
    info "venv already exists at $SCRIPT_DIR/venv — skipping creation."
else
    info "Creating venv at $SCRIPT_DIR/venv …"
    python3 -m venv "$SCRIPT_DIR/venv"
    ok "Virtual environment created"
fi

# ── 6. Install pip dependencies ─────────────────────────────────────────────
step "Installing Python dependencies"

info "Running pip install -r requirements.txt …"
"$SCRIPT_DIR/venv/bin/pip" install --upgrade pip --quiet
"$SCRIPT_DIR/venv/bin/pip" install -r "$SCRIPT_DIR/requirements.txt"
ok "All Python dependencies installed"

# ── 7. Make Real-ESRGAN binary executable ────────────────────────────────────
step "Configuring Real-ESRGAN binary"

if [ -f "$SCRIPT_DIR/realesrgan-bin/realesrgan-ncnn-vulkan" ]; then
    chmod +x "$SCRIPT_DIR/realesrgan-bin/realesrgan-ncnn-vulkan"
    ok "realesrgan-ncnn-vulkan marked as executable"
else
    warn "realesrgan-ncnn-vulkan binary not found at $SCRIPT_DIR/realesrgan-bin/"
    warn "You may need to download it manually."
fi

# ── 8. Create working directories ───────────────────────────────────────────
step "Creating working directories"

for dir in uploads outputs temp; do
    mkdir -p "$SCRIPT_DIR/$dir"
done
ok "Directories ready: uploads/ outputs/ temp/"

# ── Done ─────────────────────────────────────────────────────────────────────
echo ""
echo -e "${GREEN}${BOLD}────────────────────────────────────────────${RESET}"
echo -e "${GREEN}${BOLD}  ✅ Setup complete!${RESET}"
echo -e "${GREEN}${BOLD}────────────────────────────────────────────${RESET}"
echo ""
echo -e "  Launch the server with:"
echo -e "    ${CYAN}python start.py${RESET}"
echo ""