#!/bin/bash
# =============================================================================
# AYON Backup Restore & Database Repair Script
# Author: Daniel van Westen
# Date: 2026-02-04
# License: MIT - Use at your own risk, no warranty!
# =============================================================================
set -euo pipefail

# === Configuration ===
BACKUP_DIR="$HOME/ayon-docker/backups"
LOG_DIR="$HOME/ayon-docker/logs"
AYON_DIR="$HOME/ayon-docker"

TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOGFILE="$LOG_DIR/restore_${TIMESTAMP}.log"

# === Logging ===
mkdir -p "$LOG_DIR"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$LOGFILE"
}

log_error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $*" | tee -a "$LOGFILE" >&2
}

log_success() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$LOGFILE"
}

# === Help ===
usage() {
    echo ""
    echo "AYON Backup Restore & Database Repair Script"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --repair                    Repair database only (no restore)"
    echo "  --full <backup_file>        Restore full backup (.sql.gz)"
    echo "  --project <backup_file>     Restore project backup (.sql.gz)"
    echo "  --list                      Show available backups"
    echo "  --dry-run                   Show what would be done without making changes"
    echo "  -h, --help                  Show this help"
    echo ""
    echo "Examples:"
    echo "  $0 --repair                                    # Repair current database"
    echo "  $0 --full backups/20250201_full_backup.sql.gz  # Restore full backup"
    echo "  $0 --project backups/20250201_myproject.sql.gz # Restore project"
    echo "  $0 --list                                      # List backups"
    echo ""
    exit 0
}

# === Check if container is running ===
check_postgres() {
    cd "$AYON_DIR"
    if ! docker compose ps postgres --format '{{.State}}' 2>/dev/null | grep -q running; then
        log_error "PostgreSQL container is not running!"
        log_error "Start with: cd $AYON_DIR && docker compose up -d postgres"
        exit 1
    fi
}

# === Execute SQL ===
run_sql() {
    docker compose exec -T postgres psql -U ayon ayon -c "$1" 2>&1
}

run_sql_quiet() {
    docker compose exec -T postgres psql -U ayon ayon -t -c "$1" 2>/dev/null | tr -d ' '
}

# === List backups ===
list_backups() {
    echo ""
    echo "=== Available backups in $BACKUP_DIR ==="
    echo ""

    echo "Full backups:"
    echo "-------------"
    find "$BACKUP_DIR" -name "*_full_backup.sql.gz" -type f -printf "%T+ %p\n" 2>/dev/null | sort -r | while read -r line; do
        file=$(echo "$line" | cut -d' ' -f2)
        size=$(du -h "$file" | cut -f1)
        echo "  $file ($size)"
    done

    echo ""
    echo "Project backups:"
    echo "----------------"
    find "$BACKUP_DIR" -name "*.sql.gz" ! -name "*_full_backup.sql.gz" -type f -printf "%T+ %p\n" 2>/dev/null | sort -r | head -20 | while read -r line; do
        file=$(echo "$line" | cut -d' ' -f2)
        size=$(du -h "$file" | cut -f1)
        echo "  $file ($size)"
    done

    echo ""
    echo "(Showing max. 20 project backups, sorted by date)"
    echo ""
}

# === Get all project schemas ===
get_project_schemas() {
    run_sql_quiet "SELECT schema_name FROM information_schema.schemata WHERE schema_name LIKE 'project_%';" | grep -v '^$'
}

# === Schema migration: add product_base_type column ===
migrate_product_base_type() {
    local schema="$1"
    local dry_run="${2:-false}"

    # Check if column already exists
    local exists
    exists=$(run_sql_quiet "SELECT column_name FROM information_schema.columns WHERE table_schema = '$schema' AND table_name = 'products' AND column_name = 'product_base_type';")

    if [ -z "$exists" ]; then
        if [ "$dry_run" = "true" ]; then
            log "  [DRY-RUN] Would add product_base_type to $schema.products"
        else
            log "  Adding product_base_type to $schema.products..."
            run_sql "ALTER TABLE $schema.products ADD COLUMN IF NOT EXISTS product_base_type VARCHAR;" >> "$LOGFILE" 2>&1

            # Set product_base_type = product_type for existing entries
            run_sql "UPDATE $schema.products SET product_base_type = product_type WHERE product_base_type IS NULL;" >> "$LOGFILE" 2>&1
            log_success "Added product_base_type to $schema.products"
        fi
        return 0
    else
        log "  product_base_type already exists in $schema.products"
        return 1
    fi
}

# === Reset sequence ===
reset_sequence() {
    local schema="$1"
    local table="$2"
    local seq_name="${table}_creation_order_seq"
    local dry_run="${3:-false}"

    # Check if sequence exists
    local seq_exists
    seq_exists=$(run_sql_quiet "SELECT sequence_name FROM information_schema.sequences WHERE sequence_schema = '$schema' AND sequence_name = '$seq_name';")

    if [ -z "$seq_exists" ]; then
        return 0
    fi

    # Get max creation_order
    local max_val
    max_val=$(run_sql_quiet "SELECT COALESCE(MAX(creation_order), 0) FROM $schema.$table;" 2>/dev/null || echo "0")
    max_val=${max_val:-0}

    # Get current sequence value
    local current_val
    current_val=$(run_sql_quiet "SELECT last_value FROM $schema.$seq_name;" 2>/dev/null || echo "0")
    current_val=${current_val:-0}

    if [ "$max_val" -ge "$current_val" ]; then
        local new_val=$((max_val + 1))
        if [ "$dry_run" = "true" ]; then
            log "  [DRY-RUN] Would set $schema.$seq_name to $new_val (current: $current_val, max: $max_val)"
        else
            run_sql "SELECT setval('$schema.$seq_name', $new_val, false);" >> "$LOGFILE" 2>&1
            log "  $schema.$seq_name: $current_val -> $new_val"
        fi
    fi
}

# === Reset all sequences for a schema ===
reset_all_sequences() {
    local schema="$1"
    local dry_run="${2:-false}"

    local tables="folders products versions representations tasks workfiles links activities activity_references entity_lists"

    for table in $tables; do
        reset_sequence "$schema" "$table" "$dry_run"
    done
}

# === Reset public sequences ===
reset_public_sequences() {
    local dry_run="${1:-false}"

    log "Resetting public sequences..."

    # events_creation_order_seq
    local max_events
    max_events=$(run_sql_quiet "SELECT COALESCE(MAX(creation_order), 0) FROM public.events;" 2>/dev/null || echo "0")
    max_events=${max_events:-0}

    local current_events
    current_events=$(run_sql_quiet "SELECT last_value FROM public.events_creation_order_seq;" 2>/dev/null || echo "0")
    current_events=${current_events:-0}

    if [ "$max_events" -ge "$current_events" ]; then
        local new_val=$((max_events + 1))
        if [ "$dry_run" = "true" ]; then
            log "  [DRY-RUN] Would set public.events_creation_order_seq to $new_val"
        else
            run_sql "SELECT setval('public.events_creation_order_seq', $new_val, false);" >> "$LOGFILE" 2>&1
            log "  public.events_creation_order_seq: $current_events -> $new_val"
        fi
    fi
}

# === Repair database ===
repair_database() {
    local dry_run="${1:-false}"

    log "=== Starting database repair ==="

    if [ "$dry_run" = "true" ]; then
        log "[DRY-RUN mode - no changes will be made]"
    fi

    # Get all project schemas
    local schemas
    schemas=$(get_project_schemas)

    if [ -z "$schemas" ]; then
        log "No project schemas found"
        return 0
    fi

    local schema_count
    schema_count=$(echo "$schemas" | wc -l)
    log "Found: $schema_count project schema(s)"

    # 1. Schema migrations
    log ""
    log "--- Schema migrations ---"
    local migrations_done=0

    for schema in $schemas; do
        log "Checking schema: $schema"
        if migrate_product_base_type "$schema" "$dry_run"; then
            ((migrations_done++)) || true
        fi
    done

    log ""
    log "Schema migrations completed: $migrations_done change(s)"

    # 2. Reset sequences
    log ""
    log "--- Resetting sequences ---"

    for schema in $schemas; do
        log "Resetting sequences for $schema..."
        reset_all_sequences "$schema" "$dry_run"
    done

    # Public sequences
    reset_public_sequences "$dry_run"

    log ""
    log_success "Database repair completed"

    # 3. Recommend AYON server reload
    if [ "$dry_run" = "false" ]; then
        log ""
        log "RECOMMENDATION: Reload AYON server with:"
        log "  cd $AYON_DIR && docker compose restart server"
        log "  or: make reload"
    fi
}

# === Restore full backup ===
restore_full_backup() {
    local backup_file="$1"
    local dry_run="${2:-false}"

    if [ ! -f "$backup_file" ]; then
        log_error "Backup file not found: $backup_file"
        exit 1
    fi

    log "=== Restoring full backup: $backup_file ==="

    if [ "$dry_run" = "true" ]; then
        log "[DRY-RUN mode]"
        log "Would perform the following steps:"
        log "  1. Drop all project_* schemas"
        log "  2. Import backup"
        log "  3. Repair database"
        return 0
    fi

    # Warning
    echo ""
    echo "WARNING: This will OVERWRITE ALL existing project data!"
    echo "Backup: $backup_file"
    echo ""
    read -p "Continue? (yes/no): " confirm

    if [ "$confirm" != "yes" ]; then
        log "Cancelled by user"
        exit 0
    fi

    log "Restoring backup..."

    # Import backup
    if [[ "$backup_file" == *.gz ]]; then
        gunzip -c "$backup_file" | docker compose exec -T postgres psql -U ayon ayon >> "$LOGFILE" 2>&1
    else
        docker compose exec -T postgres psql -U ayon ayon < "$backup_file" >> "$LOGFILE" 2>&1
    fi

    log_success "Backup imported"

    # Perform repair
    repair_database "false"
}

# === Restore project backup ===
restore_project_backup() {
    local backup_file="$1"
    local dry_run="${2:-false}"

    if [ ! -f "$backup_file" ]; then
        log_error "Backup file not found: $backup_file"
        exit 1
    fi

    # Extract project name from filename (format: YYYYMMDD_projectname.sql.gz)
    local filename
    filename=$(basename "$backup_file")
    local projectname
    projectname=$(echo "$filename" | sed 's/^[0-9]*_//; s/\.sql\.gz$//; s/\.sql$//')

    log "=== Restoring project backup ==="
    log "File: $backup_file"
    log "Project: $projectname"

    if [ "$dry_run" = "true" ]; then
        log "[DRY-RUN mode]"
        log "Would perform the following steps:"
        log "  1. Drop schema project_$projectname (if exists)"
        log "  2. Import project backup"
        log "  3. Repair schema project_$projectname"
        return 0
    fi

    # Warning
    echo ""
    echo "WARNING: This will OVERWRITE project '$projectname' (if it exists)!"
    echo ""
    read -p "Continue? (yes/no): " confirm

    if [ "$confirm" != "yes" ]; then
        log "Cancelled by user"
        exit 0
    fi

    log "Restoring project backup..."

    # Import backup (already contains DROP SCHEMA)
    if [[ "$backup_file" == *.gz ]]; then
        gunzip -c "$backup_file" | docker compose exec -T postgres psql -U ayon ayon >> "$LOGFILE" 2>&1
    else
        docker compose exec -T postgres psql -U ayon ayon < "$backup_file" >> "$LOGFILE" 2>&1
    fi

    log_success "Project backup imported"

    # Repair only this schema
    log ""
    log "--- Repairing schema project_$projectname ---"
    migrate_product_base_type "project_$projectname" "false"
    reset_all_sequences "project_$projectname" "false"

    log_success "Project $projectname restored and repaired"

    log ""
    log "RECOMMENDATION: Reload AYON server with:"
    log "  cd $AYON_DIR && make reload"
}

# === Main ===
main() {
    cd "$AYON_DIR"

    # Parse arguments
    local mode=""
    local backup_file=""
    local dry_run="false"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --repair)
                mode="repair"
                shift
                ;;
            --full)
                mode="full"
                backup_file="$2"
                shift 2
                ;;
            --project)
                mode="project"
                backup_file="$2"
                shift 2
                ;;
            --list)
                list_backups
                exit 0
                ;;
            --dry-run)
                dry_run="true"
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                log_error "Unknown option: $1"
                usage
                ;;
        esac
    done

    if [ -z "$mode" ]; then
        usage
    fi

    # Check PostgreSQL
    check_postgres

    log "=== AYON Restore Script started ==="
    log "Mode: $mode"
    log "Log: $LOGFILE"

    case "$mode" in
        repair)
            repair_database "$dry_run"
            ;;
        full)
            restore_full_backup "$backup_file" "$dry_run"
            ;;
        project)
            restore_project_backup "$backup_file" "$dry_run"
            ;;
    esac

    log ""
    log "=== Done ==="
}

main "$@"