#!/usr/bin/env bash
# ============================================================================
# Cloud Backup Restore Script
#
# Restores the full media-downloader + Immich stack from a B2 cloud backup.
# Run on a fresh Ubuntu 24.04 server (or the same machine after failure).
#
# Usage:
#   sudo bash cloud_backup_restore.sh [--rclone-conf /path/to/rclone.conf]
#
# Prerequisites on a fresh machine:
#   apt update && apt install -y rclone
#   Then place your rclone.conf at /root/.config/rclone/rclone.conf
#   (contains cloud-backup-remote + cloud-backup-crypt sections)
#
# The script is interactive — it will ask before each destructive step.
# ============================================================================

set -euo pipefail

# ── Configuration ──────────────────────────────────────────────────────────

IMMICH_BASE="/opt/immich"
APP_DIR="/opt/media-downloader"
RCLONE_CONF="${1:---rclone-conf}"
RESTORE_TMP="/tmp/cloud-backup-restore"
LOG_FILE="/tmp/cloud_backup_restore.log"

# If --rclone-conf was passed, grab the value
if [[ "${1:-}" == "--rclone-conf" ]]; then
    RCLONE_CONF_PATH="${2:-/root/.config/rclone/rclone.conf}"
else
    RCLONE_CONF_PATH="/root/.config/rclone/rclone.conf"
fi

RCLONE_CRYPT="cloud-backup-crypt"

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

# ── Helpers ────────────────────────────────────────────────────────────────

log() { echo -e "${GREEN}[$(date '+%H:%M:%S')]${NC} $*" | tee -a "$LOG_FILE"; }
warn() { echo -e "${YELLOW}[$(date '+%H:%M:%S')] WARNING:${NC} $*" | tee -a "$LOG_FILE"; }
err() { echo -e "${RED}[$(date '+%H:%M:%S')] ERROR:${NC} $*" | tee -a "$LOG_FILE"; }
step() { echo -e "\n${BLUE}━━━ $* ━━━${NC}" | tee -a "$LOG_FILE"; }

confirm() {
    local msg="$1"
    echo -en "${YELLOW}$msg [y/N]: ${NC}"
    read -r answer
    [[ "$answer" =~ ^[Yy]$ ]]
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        err "This script must be run as root"
        exit 1
    fi
}

# ── Pre-flight checks ─────────────────────────────────────────────────────

preflight() {
    step "Pre-flight checks"

    check_root

    # Check rclone
    if ! command -v rclone &>/dev/null; then
        err "rclone not installed. Install with: apt install -y rclone"
        exit 1
    fi
    log "rclone: $(rclone --version | head -1)"

    # Check rclone config
    if [[ ! -f "$RCLONE_CONF_PATH" ]]; then
        err "rclone config not found at $RCLONE_CONF_PATH"
        echo "You need the rclone.conf with [cloud-backup-remote] and [cloud-backup-crypt] sections."
        echo "If restoring to a new machine, copy rclone.conf from your backup records."
        exit 1
    fi
    log "rclone config: $RCLONE_CONF_PATH"

    # Test remote connection
    log "Testing remote connection..."
    if rclone lsd "${RCLONE_CRYPT}:" --config "$RCLONE_CONF_PATH" --max-depth 1 &>/dev/null; then
        log "Remote connection: OK"
    else
        err "Cannot connect to remote. Check your rclone config and encryption passwords."
        exit 1
    fi

    # Show what's available on remote
    log "Remote directories:"
    rclone lsd "${RCLONE_CRYPT}:" --config "$RCLONE_CONF_PATH" --max-depth 1 2>/dev/null | tee -a "$LOG_FILE"

    mkdir -p "$RESTORE_TMP"
}

# ── Step 1: Download app_backup and db_dumps first ────────────────────────

download_configs() {
    step "Step 1: Downloading app_backup and db_dumps from remote"

    mkdir -p "$RESTORE_TMP/app_backup" "$RESTORE_TMP/db_dumps"

    log "Downloading app_backup..."
    rclone copy "${RCLONE_CRYPT}:app_backup" "$RESTORE_TMP/app_backup" \
        --config "$RCLONE_CONF_PATH" --progress 2>&1 | tee -a "$LOG_FILE"

    log "Downloading db_dumps..."
    rclone copy "${RCLONE_CRYPT}:db_dumps" "$RESTORE_TMP/db_dumps" \
        --config "$RCLONE_CONF_PATH" --progress 2>&1 | tee -a "$LOG_FILE"

    # Verify we got the essentials
    if [[ ! -f "$RESTORE_TMP/app_backup/media-downloader-app.tar.gz" ]]; then
        err "media-downloader-app.tar.gz not found in backup!"
        exit 1
    fi
    log "App archive size: $(du -sh "$RESTORE_TMP/app_backup/media-downloader-app.tar.gz" | cut -f1)"

    ls -la "$RESTORE_TMP/db_dumps/" | tee -a "$LOG_FILE"
    ls -la "$RESTORE_TMP/app_backup/" | tee -a "$LOG_FILE"
    ls -la "$RESTORE_TMP/app_backup/systemd/" 2>/dev/null | tee -a "$LOG_FILE"
}

# ── Step 2: Install system dependencies ───────────────────────────────────

install_dependencies() {
    step "Step 2: Install system dependencies"

    if ! confirm "Install system packages (python3, postgresql, docker, node, etc.)?"; then
        warn "Skipping dependency installation"
        return
    fi

    log "Updating package lists..."
    apt update

    log "Installing core packages..."
    apt install -y \
        python3 python3-venv python3-pip python3-dev \
        postgresql postgresql-client \
        docker.io docker-compose-v2 \
        nodejs npm \
        rclone \
        xvfb \
        python3-pyinotify \
        nginx \
        git curl wget jq \
        build-essential libffi-dev libssl-dev \
        libgl1-mesa-glx libglib2.0-0 \
        ffmpeg \
        2>&1 | tee -a "$LOG_FILE"

    # Enable and start essential services
    systemctl enable --now docker
    systemctl enable --now postgresql

    log "System dependencies installed"
}

# ── Step 3: Restore media-downloader application ──────────────────────────

restore_app() {
    step "Step 3: Restore media-downloader application"

    if [[ -d "$APP_DIR" ]]; then
        if confirm "$APP_DIR already exists. Back it up and replace?"; then
            local backup_name="${APP_DIR}.bak.$(date +%Y%m%d_%H%M%S)"
            log "Moving existing app to $backup_name"
            mv "$APP_DIR" "$backup_name"
        else
            warn "Skipping app restore"
            return
        fi
    fi

    log "Extracting media-downloader app..."
    mkdir -p /opt
    tar xzf "$RESTORE_TMP/app_backup/media-downloader-app.tar.gz" -C /opt
    log "App extracted to $APP_DIR"

    # Recreate venv
    log "Creating Python virtual environment..."
    python3 -m venv "$APP_DIR/venv"

    log "Installing Python dependencies (this may take a while)..."
    "$APP_DIR/venv/bin/pip" install --upgrade pip wheel 2>&1 | tail -3 | tee -a "$LOG_FILE"
    "$APP_DIR/venv/bin/pip" install -r "$APP_DIR/requirements.txt" 2>&1 | tail -10 | tee -a "$LOG_FILE"
    log "Python dependencies installed"

    # Rebuild frontend
    log "Installing frontend dependencies..."
    cd "$APP_DIR/web/frontend"
    npm install 2>&1 | tail -5 | tee -a "$LOG_FILE"
    log "Building frontend..."
    npx tsc && npx vite build 2>&1 | tail -5 | tee -a "$LOG_FILE"
    log "Frontend built"

    # Install Playwright browsers
    log "Installing Playwright browsers..."
    "$APP_DIR/venv/bin/python3" -m playwright install chromium firefox 2>&1 | tail -5 | tee -a "$LOG_FILE"
    "$APP_DIR/venv/bin/python3" -m playwright install-deps 2>&1 | tail -5 | tee -a "$LOG_FILE"

    # Create required directories
    mkdir -p "$APP_DIR/logs" "$APP_DIR/temp" "$APP_DIR/cache/thumbnails"

    log "Application restored"
}

# ── Step 4: Restore Immich ────────────────────────────────────────────────

restore_immich() {
    step "Step 4: Restore Immich configuration"

    mkdir -p "$IMMICH_BASE"

    # Restore docker-compose and .env
    if [[ -f "$RESTORE_TMP/app_backup/immich-docker-compose.yml" ]]; then
        cp "$RESTORE_TMP/app_backup/immich-docker-compose.yml" "$IMMICH_BASE/docker-compose.yml"
        log "Restored Immich docker-compose.yml"
    fi
    if [[ -f "$RESTORE_TMP/app_backup/immich-env" ]]; then
        cp "$RESTORE_TMP/app_backup/immich-env" "$IMMICH_BASE/.env"
        log "Restored Immich .env"
    fi

    # Create required directories
    mkdir -p "$IMMICH_BASE/upload" "$IMMICH_BASE/db" "$IMMICH_BASE/db_dumps" "$IMMICH_BASE/app_backup"

    log "Immich config restored. Media files will be synced in Step 7."
}

# ── Step 5: Restore databases ─────────────────────────────────────────────

restore_databases() {
    step "Step 5: Restore databases"

    # Media Downloader PostgreSQL (supports both .dump and legacy .sql)
    # Media Downloader PostgreSQL (supports directory dump, .dump, and legacy .sql)
    local md_dir="$RESTORE_TMP/db_dumps/media_downloader_dump"
    local md_dump="$RESTORE_TMP/db_dumps/media_downloader.dump"
    local md_sql="$RESTORE_TMP/db_dumps/media_downloader.sql"
    if [[ -d "$md_dir" || -f "$md_dump" || -f "$md_sql" ]]; then
        if confirm "Restore media_downloader PostgreSQL database?"; then
            log "Creating media_downloader database and user..."
            sudo -u postgres psql -c "CREATE USER media_downloader WITH PASSWORD 'PNsihOXvvuPwWiIvGlsc9Fh2YmMmB';" 2>/dev/null || true
            sudo -u postgres psql -c "DROP DATABASE IF EXISTS media_downloader;" 2>/dev/null || true
            sudo -u postgres psql -c "CREATE DATABASE media_downloader OWNER media_downloader;" 2>/dev/null || true

            if [[ -d "$md_dir" ]]; then
                log "Importing media_downloader dump (parallel directory format)..."
                PGPASSWORD=PNsihOXvvuPwWiIvGlsc9Fh2YmMmB pg_restore -h localhost -U media_downloader \
                    -d media_downloader --no-owner --no-acl -j 4 "$md_dir" 2>&1 | tail -5 | tee -a "$LOG_FILE"
            elif [[ -f "$md_dump" ]]; then
                log "Importing media_downloader dump (custom format)..."
                PGPASSWORD=PNsihOXvvuPwWiIvGlsc9Fh2YmMmB pg_restore -h localhost -U media_downloader \
                    -d media_downloader --no-owner --no-acl "$md_dump" 2>&1 | tail -5 | tee -a "$LOG_FILE"
            else
                log "Importing media_downloader dump (SQL format)..."
                PGPASSWORD=PNsihOXvvuPwWiIvGlsc9Fh2YmMmB psql -h localhost -U media_downloader \
                    -d media_downloader < "$md_sql" 2>&1 | tail -5 | tee -a "$LOG_FILE"
            fi
            log "media_downloader database restored"
        fi
    else
        warn "media_downloader dump not found in backup"
    fi

    # Immich PostgreSQL (supports .tar directory dump, .dump, and legacy .sql)
    local im_tar="$RESTORE_TMP/db_dumps/immich_dump.tar"
    local im_dump="$RESTORE_TMP/db_dumps/immich.dump"
    local im_sql="$RESTORE_TMP/db_dumps/immich.sql"
    if [[ -f "$im_tar" || -f "$im_dump" || -f "$im_sql" ]]; then
        if confirm "Restore Immich PostgreSQL database? (starts Immich containers first)"; then
            log "Starting Immich database container..."
            cd "$IMMICH_BASE"
            docker compose up -d database 2>&1 | tee -a "$LOG_FILE"

            log "Waiting for Immich PostgreSQL to be ready..."
            for i in $(seq 1 30); do
                if docker exec immich_postgres pg_isready -U postgres &>/dev/null; then
                    break
                fi
                sleep 2
            done

            if [[ -f "$im_tar" ]]; then
                log "Importing Immich dump (parallel directory format)..."
                docker cp "$im_tar" immich_postgres:/tmp/immich_dump.tar
                docker exec immich_postgres sh -c "cd /tmp && tar xf immich_dump.tar"
                docker exec immich_postgres pg_restore -U postgres -d immich \
                    --no-owner --no-acl -j 4 /tmp/immich_dump 2>&1 | tail -5 | tee -a "$LOG_FILE"
                docker exec immich_postgres sh -c "rm -rf /tmp/immich_dump /tmp/immich_dump.tar"
            elif [[ -f "$im_dump" ]]; then
                log "Importing Immich dump (custom format)..."
                docker cp "$im_dump" immich_postgres:/tmp/immich.dump
                docker exec immich_postgres pg_restore -U postgres -d immich \
                    --no-owner --no-acl /tmp/immich.dump 2>&1 | tail -5 | tee -a "$LOG_FILE"
                docker exec immich_postgres rm -f /tmp/immich.dump
            else
                log "Importing Immich dump (SQL format)..."
                docker exec -i immich_postgres psql -U postgres -d immich \
                    < "$im_sql" 2>&1 | tail -5 | tee -a "$LOG_FILE"
            fi
            log "Immich database restored"
        fi
    else
        warn "immich dump not found in backup"
    fi
}

# ── Step 6: Restore systemd services & rclone config ─────────────────────

restore_services() {
    step "Step 6: Restore systemd services and configs"

    # Systemd service files
    if [[ -d "$RESTORE_TMP/app_backup/systemd" ]]; then
        if confirm "Install systemd service files?"; then
            for svc in "$RESTORE_TMP/app_backup/systemd/"*; do
                local name=$(basename "$svc")
                cp "$svc" "/etc/systemd/system/$name"
                log "Installed $name"
            done
            systemctl daemon-reload
            log "systemd reloaded"
        fi
    fi

    # rclone config
    if [[ -f "$RESTORE_TMP/app_backup/rclone.conf" ]]; then
        if confirm "Restore rclone.conf?"; then
            mkdir -p /root/.config/rclone
            cp "$RESTORE_TMP/app_backup/rclone.conf" /root/.config/rclone/rclone.conf
            chmod 600 /root/.config/rclone/rclone.conf
            log "rclone.conf restored"
        fi
    fi
}

# ── Step 7: Download media files ──────────────────────────────────────────

restore_media() {
    step "Step 7: Download media files from remote"

    local dirs=$(rclone lsd "${RCLONE_CRYPT}:" --config "$RCLONE_CONF_PATH" --max-depth 1 2>/dev/null | awk '{print $NF}')

    echo ""
    log "Available remote directories:"
    echo "$dirs" | while read -r d; do echo "  - $d"; done

    if ! confirm "Download all media directories? (This may take a long time for large backups)"; then
        warn "Skipping media download. You can sync manually later with:"
        echo "  rclone copy ${RCLONE_CRYPT}:<dir> ${IMMICH_BASE}/<dir> --config $RCLONE_CONF_PATH --progress --transfers 4"
        return
    fi

    echo "$dirs" | while read -r dir_name; do
        # Skip dirs we already downloaded
        [[ "$dir_name" == "app_backup" || "$dir_name" == "db_dumps" ]] && continue
        [[ -z "$dir_name" ]] && continue

        log "Syncing $dir_name..."
        mkdir -p "$IMMICH_BASE/$dir_name"
        rclone copy "${RCLONE_CRYPT}:${dir_name}" "$IMMICH_BASE/$dir_name" \
            --config "$RCLONE_CONF_PATH" \
            --progress \
            --transfers 4 \
            --checkers 8 \
            2>&1 | tee -a "$LOG_FILE"
        log "$dir_name done"
    done

    log "Media files restored"
}

# ── Step 8: Start services ───────────────────────────────────────────────

start_services() {
    step "Step 8: Start services"

    if confirm "Start all services?"; then
        # Start Immich stack
        log "Starting Immich..."
        cd "$IMMICH_BASE"
        docker compose up -d 2>&1 | tee -a "$LOG_FILE"

        # Enable and start media-downloader services
        log "Starting media-downloader services..."
        systemctl enable --now xvfb-media-downloader.service 2>/dev/null || true
        systemctl enable --now media-downloader-api.service
        systemctl enable --now media-downloader.service
        systemctl enable --now media-downloader-frontend.service 2>/dev/null || true
        systemctl enable --now media-downloader-db-cleanup.timer 2>/dev/null || true
        systemctl enable --now cloud-backup-sync.service

        sleep 5

        # Status check
        log "Service status:"
        for svc in media-downloader-api media-downloader cloud-backup-sync; do
            local status=$(systemctl is-active "$svc" 2>/dev/null || echo "not found")
            if [[ "$status" == "active" ]]; then
                echo -e "  ${GREEN}●${NC} $svc: $status"
            else
                echo -e "  ${RED}●${NC} $svc: $status"
            fi
        done

        # Docker containers
        log "Docker containers:"
        docker ps --format "table {{.Names}}\t{{.Status}}" | tee -a "$LOG_FILE"
    fi
}

# ── Step 9: Post-restore verification ─────────────────────────────────────

verify() {
    step "Step 9: Post-restore verification"

    local issues=0

    # Check API
    if curl -sf http://localhost:8000/api/health &>/dev/null; then
        log "API health check: OK"
    else
        warn "API health check: FAILED (may still be starting)"
        ((issues++))
    fi

    # Check Immich
    if curl -sf http://localhost:2283/api/server-info/ping &>/dev/null; then
        log "Immich health check: OK"
    else
        warn "Immich health check: FAILED (may still be starting)"
        ((issues++))
    fi

    # Check database
    if PGPASSWORD=PNsihOXvvuPwWiIvGlsc9Fh2YmMmB psql -h localhost -U media_downloader -d media_downloader -c "SELECT 1" &>/dev/null; then
        log "Media Downloader DB: OK"
    else
        warn "Media Downloader DB: FAILED"
        ((issues++))
    fi

    # Disk usage
    log "Disk usage:"
    df -h /opt/immich /opt/media-downloader 2>/dev/null | tee -a "$LOG_FILE"

    echo ""
    if [[ $issues -eq 0 ]]; then
        log "${GREEN}Restore completed successfully!${NC}"
    else
        warn "Restore completed with $issues issue(s). Check the log: $LOG_FILE"
    fi

    echo ""
    log "Restore log saved to: $LOG_FILE"
}

# ── Main ──────────────────────────────────────────────────────────────────

main() {
    echo -e "${BLUE}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║         Cloud Backup Restore — Media Downloader            ║"
    echo "║                                                            ║"
    echo "║  Restores: App, Databases, Media, Configs, Services        ║"
    echo "║  Source:   Backblaze B2 (rclone crypt encrypted)           ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"

    echo "This script will restore your media-downloader + Immich stack"
    echo "from an encrypted B2 cloud backup. Each step asks for confirmation."
    echo ""
    echo "Log: $LOG_FILE"
    echo ""

    if ! confirm "Ready to begin restore?"; then
        echo "Aborted."
        exit 0
    fi

    echo "" > "$LOG_FILE"

    preflight
    download_configs
    install_dependencies
    restore_app
    restore_immich
    restore_databases
    restore_services
    restore_media
    start_services
    verify
}

main "$@"