#!/bin/bash

################################################################################
# Advanced Plex Database Recovery Script
################################################################################
#
# Author: Peter Wood <peter@peterwood.dev>
# Description: Advanced database recovery script with multiple repair strategies
#              for corrupted Plex databases. Implements progressive recovery
#              techniques from gentle repairs to aggressive reconstruction
#              methods, with comprehensive logging and rollback capabilities.
#
# Features:
#   - Progressive recovery strategy (gentle to aggressive)
#   - Multiple repair techniques (VACUUM, dump/restore, rebuild)
#   - Automatic backup before any recovery attempts
#   - Database integrity verification at each step
#   - Rollback capability if recovery fails
#   - Dry-run mode for safe testing
#   - Comprehensive logging and reporting
#
# Related Scripts:
#   - backup-plex.sh: Creates backups for recovery scenarios
#   - icu-aware-recovery.sh: ICU-specific recovery methods
#   - nuclear-plex-recovery.sh: Last-resort complete replacement
#   - validate-plex-recovery.sh: Validates recovery results
#   - restore-plex.sh: Standard restoration from backups
#   - plex.sh: General Plex service management
#
# Usage:
#   ./recover-plex-database.sh                 # Interactive recovery
#   ./recover-plex-database.sh --auto          # Automated recovery
#   ./recover-plex-database.sh --dry-run       # Show recovery plan
#   ./recover-plex-database.sh --gentle        # Gentle repair only
#   ./recover-plex-database.sh --aggressive    # Aggressive repair methods
#
# Dependencies:
#   - sqlite3 or Plex SQLite binary
#   - systemctl (for service management)
#   - Sufficient disk space for backups and temp files
#
# Exit Codes:
#   0 - Recovery successful
#   1 - General error
#   2 - Database corruption beyond repair
#   3 - Service management failure
#   4 - Insufficient disk space
#   5 - Recovery partially successful (manual intervention needed)
#
################################################################################

# Advanced Plex Database Recovery Script
# Usage: ./recover-plex-database.sh [--auto] [--dry-run]

set -e

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Configuration
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
PLEX_DB_DIR="/var/lib/plexmediaserver/Library/Application Support/Plex Media Server/Plug-in Support/Databases"
MAIN_DB="com.plexapp.plugins.library.db"
BLOBS_DB="com.plexapp.plugins.library.blobs.db"
PLEX_SQLITE="/usr/lib/plexmediaserver/Plex SQLite"
BACKUP_SUFFIX="recovery-$(date +%Y%m%d_%H%M%S)"
RECOVERY_LOG="$SCRIPT_DIR/logs/database-recovery-$(date +%Y%m%d_%H%M%S).log"

# Script options
AUTO_MODE=false
DRY_RUN=false

# Ensure logs directory exists
mkdir -p "$SCRIPT_DIR/logs"

# Logging function
log_message() {
    local message="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "$message"
    echo "$message" >> "$RECOVERY_LOG"
}

log_success() {
    log_message "${GREEN}SUCCESS: $1${NC}"
}

log_error() {
    log_message "${RED}ERROR: $1${NC}"
}

log_warning() {
    log_message "${YELLOW}WARNING: $1${NC}"
}

log_info() {
    log_message "${BLUE}INFO: $1${NC}"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --auto)
            AUTO_MODE=true
            shift
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [--auto] [--dry-run] [--help]"
            echo ""
            echo "Options:"
            echo "  --auto      Automatically attempt all recovery methods without prompts"
            echo "  --dry-run   Show what would be done without making changes"
            echo "  --help      Show this help message"
            echo ""
            echo "Recovery Methods (in order):"
            echo "  1. SQLite .recover command (modern SQLite recovery)"
            echo "  2. Partial table extraction with LIMIT"
            echo "  3. Emergency data extraction"
            echo "  4. Backup restoration from most recent good backup"
            echo ""
            exit 0
            ;;
        *)
            log_error "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Check dependencies
check_dependencies() {
    log_info "Checking dependencies..."

    if [ ! -f "$PLEX_SQLITE" ]; then
        log_error "Plex SQLite binary not found at: $PLEX_SQLITE"
        return 1
    fi

    if ! command -v sqlite3 >/dev/null 2>&1; then
        log_error "Standard sqlite3 command not found"
        return 1
    fi

    # Make Plex SQLite executable
    sudo chmod +x "$PLEX_SQLITE" 2>/dev/null || true

    log_success "Dependencies check passed"
    return 0
}

# Stop Plex service safely
stop_plex_service() {
    log_info "Stopping Plex Media Server..."

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would stop Plex service"
        return 0
    fi

    if sudo systemctl is-active --quiet plexmediaserver; then
        sudo systemctl stop plexmediaserver

        # Wait for service to fully stop
        local timeout=30
        while sudo systemctl is-active --quiet plexmediaserver && [ $timeout -gt 0 ]; do
            sleep 1
            timeout=$((timeout - 1))
        done

        if sudo systemctl is-active --quiet plexmediaserver; then
            log_error "Failed to stop Plex service within timeout"
            return 1
        fi

        log_success "Plex service stopped successfully"
    else
        log_info "Plex service was already stopped"
    fi

    return 0
}

# Start Plex service
start_plex_service() {
    log_info "Starting Plex Media Server..."

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would start Plex service"
        return 0
    fi

    sudo systemctl start plexmediaserver

    # Wait for service to start
    local timeout=30
    while ! sudo systemctl is-active --quiet plexmediaserver && [ $timeout -gt 0 ]; do
        sleep 1
        timeout=$((timeout - 1))
    done

    if sudo systemctl is-active --quiet plexmediaserver; then
        log_success "Plex service started successfully"
    else
        log_warning "Plex service may not have started properly"
    fi
}

# Check database integrity
check_database_integrity() {
    local db_file="$1"
    local db_name=$(basename "$db_file")

    log_info "Checking integrity of $db_name..."

    if [ ! -f "$db_file" ]; then
        log_error "Database file not found: $db_file"
        return 1
    fi

    local integrity_result
    integrity_result=$(sudo "$PLEX_SQLITE" "$db_file" "PRAGMA integrity_check;" 2>&1)
    local check_exit_code=$?

    if [ $check_exit_code -ne 0 ]; then
        log_error "Failed to run integrity check on $db_name"
        return 1
    fi

    if echo "$integrity_result" | grep -q "^ok$"; then
        log_success "Database integrity check passed: $db_name"
        return 0
    else
        log_warning "Database integrity issues detected in $db_name:"
        echo "$integrity_result" | while IFS= read -r line; do
            log_warning "  $line"
        done
        return 1
    fi
}

# Recovery Method 1: SQLite .recover command
recovery_method_sqlite_recover() {
    local db_file="$1"
    local db_name=$(basename "$db_file")
    local recovered_sql="${db_file}.recovered.sql"
    local new_db="${db_file}.recovered"

    log_info "Recovery Method 1: SQLite .recover command for $db_name"

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would attempt SQLite .recover method"
        return 0
    fi

    # Check if .recover is available (SQLite 3.37.0+)
    if ! echo ".help" | sqlite3 2>/dev/null | grep -q "\.recover"; then
        log_warning "SQLite .recover command not available in this version"
        return 1
    fi

    log_info "Attempting SQLite .recover method..."

    # Use standard sqlite3 for .recover as it's more reliable
    if sqlite3 "$db_file" ".recover" > "$recovered_sql" 2>/dev/null; then
        log_success "Recovery SQL generated successfully"

        # Create new database from recovered data
        if [ -f "$recovered_sql" ] && [ -s "$recovered_sql" ]; then
            if sqlite3 "$new_db" < "$recovered_sql" 2>/dev/null; then
                log_success "New database created from recovered data"

                # Verify new database integrity
                if sqlite3 "$new_db" "PRAGMA integrity_check;" | grep -q "ok"; then
                    log_success "Recovered database integrity verified"

                    # Replace original with recovered database
                    if sudo mv "$db_file" "${db_file}.corrupted" && sudo mv "$new_db" "$db_file"; then
                        sudo chown plex:plex "$db_file"
                        sudo chmod 644 "$db_file"
                        log_success "Database successfully recovered using .recover method"

                        # Clean up
                        rm -f "$recovered_sql"

                        return 0
                    else
                        log_error "Failed to replace original database"
                    fi
                else
                    log_error "Recovered database failed integrity check"
                fi
            else
                log_error "Failed to create database from recovered SQL"
            fi
        else
            log_error "Recovery SQL file is empty or not generated"
        fi
    else
        log_error "SQLite .recover command failed"
    fi

    # Clean up on failure
    rm -f "$recovered_sql" "$new_db"
    return 1
}

# Recovery Method 2: Partial table extraction
recovery_method_partial_extraction() {
    local db_file="$1"
    local db_name=$(basename "$db_file")
    local partial_sql="${db_file}.partial.sql"
    local new_db="${db_file}.partial"

    log_info "Recovery Method 2: Partial table extraction for $db_name"

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would attempt partial extraction method"
        return 0
    fi

    log_info "Extracting schema and partial data..."

    # Start the SQL file with schema
    {
        echo "-- Partial recovery of $db_name"
        echo "-- Generated on $(date)"
        echo ""
    } > "$partial_sql"

    # Extract schema
    if sudo "$PLEX_SQLITE" "$db_file" ".schema" >> "$partial_sql" 2>/dev/null; then
        log_success "Schema extracted successfully"
    else
        log_warning "Schema extraction failed, trying alternative method"
        # Try with standard sqlite3
        if sqlite3 "$db_file" ".schema" >> "$partial_sql" 2>/dev/null; then
            log_success "Schema extracted with standard sqlite3"
        else
            log_error "Schema extraction failed completely"
            rm -f "$partial_sql"
            return 1
        fi
    fi

    # Critical tables to extract (in order of importance)
    local critical_tables=(
        "accounts"
        "library_sections"
        "directories"
        "metadata_items"
        "media_items"
        "media_parts"
        "media_streams"
        "taggings"
        "tags"
    )

    log_info "Attempting to extract critical tables..."

    for table in "${critical_tables[@]}"; do
        log_info "Extracting table: $table"

        # Try to extract with LIMIT to avoid hitting corrupted data
        local extract_success=false
        local limit=10000

        while [ $limit -le 100000 ] && [ "$extract_success" = false ]; do
            if sudo "$PLEX_SQLITE" "$db_file" "SELECT COUNT(*) FROM $table;" >/dev/null 2>&1; then
                # Table exists and is readable
                {
                    echo ""
                    echo "-- Data for table $table (limited to $limit rows)"
                    echo "DELETE FROM $table;"
                } >> "$partial_sql"

                if sudo "$PLEX_SQLITE" "$db_file" ".mode insert $table" >>/dev/null 2>&1 && \
                   sudo "$PLEX_SQLITE" "$db_file" "SELECT * FROM $table LIMIT $limit;" >> "$partial_sql" 2>/dev/null; then
                    local row_count=$(tail -n +3 "$partial_sql" | grep "INSERT INTO $table" | wc -l)
                    log_success "Extracted $row_count rows from $table"
                    extract_success=true
                else
                    log_warning "Failed to extract $table with limit $limit, trying smaller limit"
                    limit=$((limit / 2))
                fi
            else
                log_warning "Table $table is not accessible or doesn't exist"
                break
            fi
        done

        if [ "$extract_success" = false ]; then
            log_warning "Could not extract any data from table $table"
        fi
    done

    # Create new database from partial data
    if [ -f "$partial_sql" ] && [ -s "$partial_sql" ]; then
        log_info "Creating database from partial extraction..."

        if sqlite3 "$new_db" < "$partial_sql" 2>/dev/null; then
            log_success "Partial database created successfully"

            # Verify basic functionality
            if sqlite3 "$new_db" "PRAGMA integrity_check;" | grep -q "ok"; then
                log_success "Partial database integrity verified"

                # Replace original with partial database
                if sudo mv "$db_file" "${db_file}.corrupted" && sudo mv "$new_db" "$db_file"; then
                    sudo chown plex:plex "$db_file"
                    sudo chmod 644 "$db_file"
                    log_success "Database partially recovered - some data may be lost"
                    log_warning "Please verify your Plex library after recovery"

                    # Clean up
                    rm -f "$partial_sql"

                    return 0
                else
                    log_error "Failed to replace original database"
                fi
            else
                log_error "Partial database failed integrity check"
            fi
        else
            log_error "Failed to create database from partial extraction"
        fi
    else
        log_error "Partial extraction SQL file is empty"
    fi

    # Clean up on failure
    rm -f "$partial_sql" "$new_db"
    return 1
}

# Recovery Method 3: Emergency data extraction
recovery_method_emergency_extraction() {
    local db_file="$1"
    local db_name=$(basename "$db_file")

    log_info "Recovery Method 3: Emergency data extraction for $db_name"

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would attempt emergency extraction method"
        return 0
    fi

    log_warning "This method will create a minimal database with basic library structure"
    log_warning "You will likely need to re-scan your media libraries"

    if [ "$AUTO_MODE" = false ]; then
        read -p "Continue with emergency extraction? This will lose most metadata [y/N]: " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Emergency extraction cancelled by user"
            return 1
        fi
    fi

    local emergency_db="${db_file}.emergency"

    # Create a minimal database with essential tables
    log_info "Creating minimal emergency database..."

    cat > "/tmp/emergency_schema.sql" << 'EOF'
-- Emergency Plex database schema (minimal)
CREATE TABLE accounts (
    id INTEGER PRIMARY KEY,
    name TEXT,
    hashed_password TEXT,
    salt TEXT,
    created_at DATETIME,
    updated_at DATETIME
);

CREATE TABLE library_sections (
    id INTEGER PRIMARY KEY,
    name TEXT,
    section_type INTEGER,
    agent TEXT,
    scanner TEXT,
    language TEXT,
    created_at DATETIME,
    updated_at DATETIME
);

CREATE TABLE directories (
    id INTEGER PRIMARY KEY,
    library_section_id INTEGER,
    path TEXT,
    created_at DATETIME,
    updated_at DATETIME
);

-- Insert default admin account
INSERT INTO accounts (id, name, created_at, updated_at)
VALUES (1, 'plex', datetime('now'), datetime('now'));
EOF

    if sqlite3 "$emergency_db" < "/tmp/emergency_schema.sql" 2>/dev/null; then
        log_success "Emergency database created"

        # Replace original with emergency database
        if sudo mv "$db_file" "${db_file}.corrupted" && sudo mv "$emergency_db" "$db_file"; then
            sudo chown plex:plex "$db_file"
            sudo chmod 644 "$db_file"
            log_success "Emergency database installed"
            log_warning "You will need to re-add library sections and re-scan media"

            # Clean up
            rm -f "/tmp/emergency_schema.sql"

            return 0
        else
            log_error "Failed to install emergency database"
        fi
    else
        log_error "Failed to create emergency database"
    fi

    # Clean up on failure
    rm -f "/tmp/emergency_schema.sql" "$emergency_db"
    return 1
}

# Recovery Method 4: Restore from backup
recovery_method_backup_restore() {
    local db_file="$1"
    local backup_dir="/mnt/share/media/backups/plex"

    log_info "Recovery Method 4: Restore from most recent backup"

    if [ "$DRY_RUN" = true ]; then
        log_info "DRY RUN: Would attempt backup restoration"
        return 0
    fi

    # Find most recent backup
    local latest_backup=$(find "$backup_dir" -maxdepth 1 -name "plex-backup-*.tar.gz" -type f 2>/dev/null | sort -r | head -1)

    if [ -z "$latest_backup" ]; then
        log_error "No backup files found in $backup_dir"
        return 1
    fi

    log_info "Found latest backup: $(basename "$latest_backup")"

    if [ "$AUTO_MODE" = false ]; then
        read -p "Restore from backup $(basename "$latest_backup")? [y/N]: " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Backup restoration cancelled by user"
            return 1
        fi
    fi

    # Extract and restore database from backup
    local temp_extract="/tmp/plex-recovery-extract-$(date +%Y%m%d_%H%M%S)"
    mkdir -p "$temp_extract"

    log_info "Extracting backup..."
    if tar -xzf "$latest_backup" -C "$temp_extract" 2>/dev/null; then
        local backup_db_file="$temp_extract/$(basename "$db_file")"

        if [ -f "$backup_db_file" ]; then
            # Verify backup database integrity
            if sqlite3 "$backup_db_file" "PRAGMA integrity_check;" | grep -q "ok"; then
                log_success "Backup database integrity verified"

                # Replace corrupted database with backup
                if sudo mv "$db_file" "${db_file}.corrupted" && sudo cp "$backup_db_file" "$db_file"; then
                    sudo chown plex:plex "$db_file"
                    sudo chmod 644 "$db_file"
                    log_success "Database restored from backup"

                    # Clean up
                    rm -rf "$temp_extract"

                    return 0
                else
                    log_error "Failed to replace database with backup"
                fi
            else
                log_error "Backup database also has integrity issues"
            fi
        else
            log_error "Database file not found in backup"
        fi
    else
        log_error "Failed to extract backup"
    fi

    # Clean up on failure
    rm -rf "$temp_extract"
    return 1
}

# Main recovery function
main_recovery() {
    local db_file="$PLEX_DB_DIR/$MAIN_DB"

    log_info "Starting Plex database recovery process"
    log_info "Recovery log: $RECOVERY_LOG"

    # Check dependencies
    if ! check_dependencies; then
        exit 1
    fi

    # Stop Plex service
    if ! stop_plex_service; then
        exit 1
    fi

    # Change to database directory
    cd "$PLEX_DB_DIR" || {
        log_error "Failed to change to database directory"
        exit 1
    }

    # Check if database exists
    if [ ! -f "$MAIN_DB" ]; then
        log_error "Main database file not found: $MAIN_DB"
        exit 1
    fi

    # Create backup of current corrupted state
    log_info "Creating backup of current corrupted database..."
    if [ "$DRY_RUN" = false ]; then
        sudo cp "$MAIN_DB" "${MAIN_DB}.${BACKUP_SUFFIX}"
        log_success "Corrupted database backed up as: ${MAIN_DB}.${BACKUP_SUFFIX}"
    fi

    # Check current integrity
    log_info "Verifying database corruption..."
    if check_database_integrity "$MAIN_DB"; then
        log_success "Database integrity check passed - no recovery needed!"
        start_plex_service
        exit 0
    fi

    log_warning "Database corruption confirmed, attempting recovery..."

    # Try recovery methods in order
    local recovery_methods=(
        "recovery_method_sqlite_recover"
        "recovery_method_partial_extraction"
        "recovery_method_emergency_extraction"
        "recovery_method_backup_restore"
    )

    for method in "${recovery_methods[@]}"; do
        log_info "Attempting: $method"

        if $method "$MAIN_DB"; then
            log_success "Recovery successful using: $method"

            # Verify the recovered database
            if check_database_integrity "$MAIN_DB"; then
                log_success "Recovered database integrity verified"
                start_plex_service
                log_success "Database recovery completed successfully!"
                log_info "Please check your Plex server and verify your libraries"
                exit 0
            else
                log_error "Recovered database still has integrity issues"
                # Restore backup for next attempt
                if [ "$DRY_RUN" = false ]; then
                    sudo cp "${MAIN_DB}.${BACKUP_SUFFIX}" "$MAIN_DB"
                fi
            fi
        else
            log_warning "Recovery method failed: $method"
        fi
    done

    log_error "All recovery methods failed"
    log_error "Manual intervention required"

    # Restore original corrupted database
    if [ "$DRY_RUN" = false ]; then
        sudo cp "${MAIN_DB}.${BACKUP_SUFFIX}" "$MAIN_DB"
    fi

    start_plex_service
    exit 1
}

# Trap to ensure Plex service is restarted
trap 'start_plex_service' EXIT

# Run main recovery
main_recovery "$@"