#!/bin/bash
set -euo pipefail

# ANSI Color Codes
readonly COLOR_RESET='\033[0m'
readonly COLOR_BLUE='\033[34m'
readonly COLOR_GREEN='\033[32m'
readonly COLOR_YELLOW='\033[33m'
readonly COLOR_RED='\033[31m'
readonly COLOR_CYAN='\033[36m'
readonly COLOR_MAGENTA='\033[35m'

log_info() { echo -e "${COLOR_BLUE}[INFO]${COLOR_RESET} $*"; }
log_success() { echo -e "${COLOR_GREEN}[SUCCESS]${COLOR_RESET} $*"; }
log_warning() { echo -e "${COLOR_YELLOW}[WARNING]${COLOR_RESET} $*"; }
log_error() { echo -e "${COLOR_RED}[ERROR]${COLOR_RESET} $*" >&2; }
log_step() { echo -e "${COLOR_CYAN}[STEP]${COLOR_RESET} $*"; }
log_debug() {
    if [[ "${DEBUG:-0}" == "1" ]]; then
        echo -e "${COLOR_MAGENTA}[DEBUG]${COLOR_RESET} $*"
    fi
}
log_separator() { echo "================================================================"; }
log_header() {
    echo ""
    log_separator
    echo -e "       ${COLOR_CYAN}$*${COLOR_RESET}"
    log_separator
    echo ""
}

confirm() {
    local prompt="$1"
    local response

    while true; do
        echo -ne "${COLOR_MAGENTA}[PROMPT]${COLOR_RESET} $prompt (y/n): "
        read -r response
        case "$response" in
            [yY][eE][sS] | [yY])
                return 0
                ;;
            [nN][oO] | [nN])
                return 1
                ;;
            *)
                log_warning "Please answer 'y' or 'n'"
                ;;
        esac
    done
}

prompt_input() {
    local prompt="$1"
    local default="$2"
    local response

    if [[ -n "$default" ]]; then
        echo -ne "${COLOR_MAGENTA}[PROMPT]${COLOR_RESET} $prompt (default: $default): "
    else
        echo -ne "${COLOR_MAGENTA}[PROMPT]${COLOR_RESET} $prompt: "
    fi

    read -r response

    if [[ -z "$response" ]] && [[ -n "$default" ]]; then
        echo "$default"
    else
        echo "$response"
    fi
}

die() {
    log_error "$*"
    exit 1
}

INSTALL_DIR=""
TARGET_VERSION=""
ACTION="rollback"
NON_INTERACTIVE=false

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --list) ACTION="list"; shift ;;
            --install-dir) INSTALL_DIR="$2"; shift 2 ;;
            --install-dir=*) INSTALL_DIR="${1#*=}"; shift ;;
            --non-interactive|-y|--yes) NON_INTERACTIVE=true; shift ;;
            --help|-h) show_help; exit 0 ;;
            *) TARGET_VERSION="$1"; shift ;;
        esac
    done
}

show_help() {
    cat << 'EOF'
Smart Home Rollback Script

USAGE:
  ./rollback.sh [OPTIONS] [version]
  curl -fsSL <url>/rollback.sh | sudo bash -s -- [OPTIONS] [version]

OPTIONS:
  --list                List available backups
  --install-dir <path>  Installation directory (default: /home/smarthome/smarthome)
  --non-interactive, -y Run without confirmation prompts
  --help, -h            Show this help

EXAMPLES:
  List backups:         ./rollback.sh --list
  Rollback to v1.0.0:   ./rollback.sh v1.0.0
  Non-interactive:      ./rollback.sh --non-interactive 1.0.0
EOF
}

list_backups() {
    log_step "Available backups:"
    if [[ ! -d "$INSTALL_DIR/backups" ]]; then
        log_warning "No backups directory found"
        return 1
    fi

    local backups=($(ls -d "$INSTALL_DIR/backups"/v* 2>/dev/null || echo ""))
    if [[ ${#backups[@]} -eq 0 ]]; then
        log_warning "No backups available"
        return 1
    fi

    for backup in "${backups[@]}"; do
        local version=$(basename "$backup")
        echo "  $version"
    done
}

validate_backup() {
    local backup_dir="$INSTALL_DIR/backups/$TARGET_VERSION"
    if [[ ! -d "$backup_dir" ]]; then
        die "Backup not found: $TARGET_VERSION"
    fi
    if [[ ! -d "$backup_dir/api" ]]; then
        die "Invalid backup: missing backend"
    fi
}

stop_services() {
    log_step "Stopping services..."
    local user
    user=$(stat -c '%U' "$INSTALL_DIR" 2>/dev/null || echo "smarthome")
    if command -v pm2 &>/dev/null; then
        sudo -u "$user" pm2 stop smart-home-api 2>/dev/null || true
        sleep 2
    fi
}

restore_backup() {
    local backup_dir="$INSTALL_DIR/backups/$TARGET_VERSION"
    log_step "Restoring backup..."
    local user
    user=$(stat -c '%U' "$INSTALL_DIR" 2>/dev/null || echo "smarthome")

    if [[ -d "$backup_dir/api" ]]; then
        log_info "Restoring backend..."
        sudo rm -rf "$INSTALL_DIR/smart-home-api"
        sudo cp -r "$backup_dir/api" "$INSTALL_DIR/smart-home-api"
        sudo chown -R "$user:$user" "$INSTALL_DIR/smart-home-api"
    fi

    if [[ -d "$backup_dir/web" ]]; then
        log_info "Restoring frontend..."
        sudo rm -rf "$INSTALL_DIR/smart-home-web"
        sudo cp -r "$backup_dir/web" "$INSTALL_DIR/smart-home-web"
        sudo chown -R "$user:$user" "$INSTALL_DIR/smart-home-web"
    fi

    log_success "Backup restored"
}

restart_services() {
    log_step "Restarting services..."
    local user
    user=$(stat -c '%U' "$INSTALL_DIR" 2>/dev/null || echo "smarthome")
    if command -v pm2 &>/dev/null; then
        local ecosystem="$INSTALL_DIR/ecosystem.config.js"
        if [[ -f "$ecosystem" ]]; then
            sudo -u "$user" pm2 restart smart-home-api || log_warning "Failed to restart PM2"
            sleep 3
            sudo -u "$user" pm2 save 2>/dev/null || true
        fi
    fi
    log_success "Services restarted"
}

verify_rollback() {
    log_step "Verifying rollback..."
    sleep 2
    local user
    user=$(stat -c '%U' "$INSTALL_DIR" 2>/dev/null || echo "smarthome")

    if command -v pm2 &>/dev/null; then
        if sudo -u "$user" pm2 status 2>/dev/null | grep -q "smart-home-api.*online"; then
            log_success "Rollback verified"
            return 0
        fi
    fi

    log_warning "Could not verify service status"
    return 1
}

main() {
    log_header "Smart Home Rollback Utility"

    parse_arguments "$@"

    # Detect install directory
    if [[ -z "$INSTALL_DIR" ]]; then
        if [[ -d "/home/smarthome/smarthome" ]]; then
            INSTALL_DIR="/home/smarthome/smarthome"
        elif [[ -d "$HOME/smarthome" ]]; then
            INSTALL_DIR="$HOME/smarthome"
        else
            die "Could not detect installation directory"
        fi
    fi

    if [[ "$ACTION" == "list" ]]; then
        list_backups
        exit 0
    fi

    if [[ -z "$TARGET_VERSION" ]]; then
        log_error "No version specified"
        show_help
        exit 1
    fi

    # Normalize version string (e.g. 1.0.0 -> v1.0.0)
    TARGET_VERSION="v${TARGET_VERSION#v}"

    validate_backup

    if [[ "$NON_INTERACTIVE" != true ]]; then
        confirm "Rollback to $TARGET_VERSION?" || { log_warning "Rollback cancelled"; exit 0; }
    else
        log_info "Non-interactive mode: proceeding with rollback to $TARGET_VERSION..."
    fi

    stop_services
    restore_backup
    restart_services
    verify_rollback || log_warning "Verification failed, check logs"

    log_success "Rollback to $TARGET_VERSION completed successfully!"
}

[[ "${BASH_SOURCE[0]:-$0}" == "$0" ]] && main "$@"