#!/bin/bash

set -euo pipefail

SERVICE_NAME="promdatapanel"
DEFAULT_APP_DIR="/opt/promdata-panel"
ZIP_URL="https://git.littlediary.cn/CN-JS-HuiBai/PromdataPanel/archive/main.zip"

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

APP_DIR=""
TEMP_DIR=""
BACKUP_DIR=""
ROLLBACK_REQUIRED=false

echo -e "${BLUE}=== Starting PromdataPanel Update ===${NC}"

cleanup() {
    if [ -n "${TEMP_DIR}" ] && [ -d "${TEMP_DIR}" ]; then
        rm -rf "${TEMP_DIR}"
    fi
}

rollback() {
    if [ "$ROLLBACK_REQUIRED" != true ] || [ -z "${BACKUP_DIR}" ] || [ ! -d "${BACKUP_DIR}" ]; then
        return
    fi

    echo -e "${YELLOW}Update failed. Restoring previous application state...${NC}"
    rsync -a --delete --exclude '.env' "${BACKUP_DIR}/" "${APP_DIR}/"
}

trap 'rollback' ERR
trap cleanup EXIT

validate_app_dir() {
    local dir="$1"
    [ -n "$dir" ] || return 1
    [ -d "$dir" ] || return 1
    [ -f "$dir/package.json" ] || return 1
    [ -f "$dir/server/index.js" ] || return 1
    [ -f "$dir/public/index.html" ] || return 1
    return 0
}

detect_app_dir() {
    local service_dir=""
    if command -v systemctl >/dev/null 2>&1 && systemctl list-unit-files | grep -q "^${SERVICE_NAME}\.service"; then
        echo "Detecting application directory from systemd service..."
        service_dir=$(systemctl show -p WorkingDirectory "$SERVICE_NAME" | cut -d= -f2-)
        if validate_app_dir "$service_dir"; then
            APP_DIR="$service_dir"
            return
        fi
    fi

    local current_dir
    current_dir=$(pwd)
    if validate_app_dir "$current_dir"; then
        APP_DIR="$current_dir"
        return
    fi

    if validate_app_dir "$DEFAULT_APP_DIR"; then
        APP_DIR="$DEFAULT_APP_DIR"
        return
    fi

    echo -e "${RED}Error: Could not locate a valid PromdataPanel application directory.${NC}"
    echo -e "${YELLOW}Expected markers: package.json, server/index.js, public/index.html${NC}"
    exit 1
}

ensure_tool() {
    local cmd="$1"
    if command -v "$cmd" >/dev/null 2>&1; then
        return
    fi

    echo -e "${BLUE}${cmd} is not installed. Attempting to install it...${NC}"
    if command -v apt-get >/dev/null 2>&1; then
        sudo apt-get update
        sudo apt-get install -y "$cmd"
    elif command -v dnf >/dev/null 2>&1; then
        sudo dnf install -y "$cmd"
    elif command -v yum >/dev/null 2>&1; then
        sudo yum install -y "$cmd"
    elif command -v apk >/dev/null 2>&1; then
        sudo apk add "$cmd"
    else
        echo -e "${RED}Error: '${cmd}' is not installed and could not be auto-installed.${NC}"
        exit 1
    fi
}

update_from_git() {
    echo -e "${BLUE}Git repository detected. Pulling latest code...${NC}"
    if [ -n "$(git status --porcelain)" ]; then
        echo -e "${RED}Error: Working tree has local changes. Commit or stash them before updating.${NC}"
        exit 1
    fi
    git pull --ff-only
}

update_from_zip() {
    echo -e "${BLUE}No git repository found. Updating via ZIP archive with staging and rollback...${NC}"
    ensure_tool curl
    ensure_tool unzip
    ensure_tool rsync

    TEMP_DIR=$(mktemp -d "${TMPDIR:-/tmp}/promdatapanel-update-XXXXXX")
    BACKUP_DIR="${TEMP_DIR}/backup"
    local archive_path="${TEMP_DIR}/latest.zip"
    local extracted_folder=""
    local staging_dir=""

    echo "Downloading latest version (main branch)..."
    curl -fL "$ZIP_URL" -o "$archive_path"

    echo "Extracting archive..."
    unzip -q "$archive_path" -d "$TEMP_DIR"
    extracted_folder=$(find "$TEMP_DIR" -mindepth 1 -maxdepth 1 -type d ! -name backup | head -n 1)

    if ! validate_app_dir "$extracted_folder"; then
        echo -e "${RED}Extraction failed or archive structure is invalid.${NC}"
        exit 1
    fi

    staging_dir="${TEMP_DIR}/staging"
    mkdir -p "$staging_dir"
    rsync -a --exclude '.git' "$extracted_folder/" "$staging_dir/"

    if [ -f "${APP_DIR}/.env" ]; then
        cp "${APP_DIR}/.env" "${staging_dir}/.env"
    fi

    echo "Installing dependencies in staging directory..."
    (
        cd "$staging_dir"
        npm install --production
    )

    echo "Creating rollback backup..."
    rsync -a --delete --exclude '.env' "${APP_DIR}/" "${BACKUP_DIR}/"

    echo "Applying staged update..."
    ROLLBACK_REQUIRED=true
    rsync -a --delete --exclude '.env' "${staging_dir}/" "${APP_DIR}/"
}

restart_service() {
    if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet "$SERVICE_NAME"; then
        echo -e "${BLUE}Restarting systemd service: ${SERVICE_NAME}...${NC}"
        sudo systemctl restart "$SERVICE_NAME"
        return
    fi

    if command -v pm2 >/dev/null 2>&1 && pm2 list | grep -q "$SERVICE_NAME"; then
        echo -e "${BLUE}Restarting with PM2...${NC}"
        pm2 restart "$SERVICE_NAME"
        return
    fi

    echo -e "${YELLOW}Warning: Could not detect an active systemd service or PM2 process named '${SERVICE_NAME}'.${NC}"
    echo -e "${YELLOW}Please restart the application manually.${NC}"
}

detect_app_dir
echo -e "${BLUE}Application directory: ${APP_DIR}${NC}"
cd "$APP_DIR"

if [ -d ".git" ]; then
    update_from_git
    echo -e "${BLUE}Updating npm dependencies...${NC}"
    npm install --production
else
    update_from_zip
fi

restart_service
ROLLBACK_REQUIRED=false

echo -e "${GREEN}=== Update successfully finished ===${NC}"