#!/usr/bin/env bash
# ==============================================================================
# VM Bench Stack Provisioning & Lifecycle Script
# Target: Proxmox VE node (any version)
# Usage:
#   ./install.sh --deploy   [Interactive Setup & Deployment]
#   ./install.sh --update   [Pull Updates & Restart Services]
#   ./install.sh --remove   [Teardown Everything]
# ==============================================================================

set -euo pipefail
# Uncomment next line for debugging:
# set -x

# ------------------------------------------------------------------------------
# Formatting & Constants
# ------------------------------------------------------------------------------
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
NC='\033[0m'

REPO_URL="https://git.lohmar.co.uk/cclohmar/vm-bench.git"
SYSTEMD_DIR="/etc/systemd/system"
SERVICE_NAME="vm-bench-backend"

CT_CORES=1
CT_MEMORY=512
CT_SWAP=512
CT_DISK_SIZE=8

# ------------------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------------------
has_cmd() {
    command -v "$1" >/dev/null 2>&1
}

get_default_bridge() {
    ip -o link show type bridge | awk -F': ' '{print $2; exit}' 2>/dev/null || echo "vmbr0"
}

guess_gateway_from_bridge() {
    local bridge="$1"
    local ip_cidr
    ip_cidr=$(ip -o -4 addr show dev "$bridge" 2>/dev/null | awk '{print $4; exit}')
    if [[ -n "$ip_cidr" ]]; then
        local ip="${ip_cidr%/*}"
        local prefix="${ip%.*}"
        echo "${prefix}.1"
    else
        echo ""
    fi
}

# ------------------------------------------------------------------------------
# Environment Discovery & User Parameter Gathering
# ------------------------------------------------------------------------------
gather_environment_info() {
    echo -e "${CYAN}==> [1/3] Environment Inspection & Parameter Collection${NC}"

    # ---- Storage Selection ----
    echo ""
    echo "Inspecting available storage pools from /etc/pve/storage.cfg..."
    declare -a storage_ids=()
    declare -a storage_contents=()
    declare -a storage_has_rootdir=()

    # Parse storage.cfg for container-capable storage types
    if [ -f /etc/pve/storage.cfg ]; then
        # Extract storage IDs for types: lvmthin, zfspool, rbd, dir
        while IFS= read -r id; do
            storage_ids+=("$id")
            # Get content for this pool (ignore errors)
            local content
            content=$(pvesm config "$id" 2>/dev/null | awk '/^content/ {$1=""; print $0}' | xargs | tr ' ' ',' || true)
            storage_contents+=("${content:-none}")
            if [[ "$content" == *"rootdir"* ]]; then
                storage_has_rootdir+=("yes")
            else
                storage_has_rootdir+=("no")
            fi
        done < <(grep -E "^(lvmthin|zfspool|rbd|dir):" /etc/pve/storage.cfg | awk '{print $2}')
    fi

    # Fallback: if no storage found in config, use active pools from pvesm
    if [ ${#storage_ids[@]} -eq 0 ]; then
        echo -e "${YELLOW}No storage definitions found in /etc/pve/storage.cfg. Falling back to active pools from 'pvesm status'.${NC}"
        while IFS= read -r pool; do
            storage_ids+=("$pool")
            local content
            content=$(pvesm config "$pool" 2>/dev/null | awk '/^content/ {$1=""; print $0}' | xargs | tr ' ' ',' || true)
            storage_contents+=("${content:-none}")
            if [[ "$content" == *"rootdir"* ]]; then
                storage_has_rootdir+=("yes")
            else
                storage_has_rootdir+=("no")
            fi
        done < <(pvesm status 2>/dev/null | awk 'NR>1 && $3=="active" {print $1}' || true)
    fi

    # If still empty, abort
    if [ ${#storage_ids[@]} -eq 0 ]; then
        echo -e "${RED}No usable storage pools found. Please create a storage pool (e.g., 'local' directory) and try again.${NC}"
        exit 1
    fi

    # Show menu
    echo "Available Storage Options Detected:"
    for i in "${!storage_ids[@]}"; do
        id="${storage_ids[$i]}"
        content="${storage_contents[$i]}"
        marker=""
        [ "${storage_has_rootdir[$i]}" = "yes" ] && marker=" ✓ (has rootdir)"
        printf "  [%d] %-18s (content: %s)%s\n" "$((i+1))" "$id" "$content" "$marker"
    done

    # Pick default: first with rootdir, else first
    default_idx=0
    for i in "${!storage_ids[@]}"; do
        if [ "${storage_has_rootdir[$i]}" = "yes" ]; then
            default_idx=$i
            break
        fi
    done
    default_num=$((default_idx+1))

    read -p "Select storage pool [1-${#storage_ids[@]}] (default: ${default_num}): " POOL_CHOICE
    POOL_CHOICE="${POOL_CHOICE:-$default_num}"
    if [[ "$POOL_CHOICE" =~ ^[0-9]+$ ]] && [ "$POOL_CHOICE" -ge 1 ] && [ "$POOL_CHOICE" -le "${#storage_ids[@]}" ]; then
        TARGET_STORAGE="${storage_ids[$((POOL_CHOICE-1))]}"
        TARGET_HAS_ROOTDIR="${storage_has_rootdir[$((POOL_CHOICE-1))]}"
    else
        echo -e "${YELLOW}Invalid choice. Defaulting to '${storage_ids[$default_idx]}'.${NC}"
        TARGET_STORAGE="${storage_ids[$default_idx]}"
        TARGET_HAS_ROOTDIR="${storage_has_rootdir[$default_idx]}"
    fi

    # Ensure rootdir capability
    if [ "$TARGET_HAS_ROOTDIR" != "yes" ]; then
        echo "Enabling 'rootdir' capability on storage pool '${TARGET_STORAGE}'..."
        current_content="${storage_contents[$((POOL_CHOICE-1))]:-}"
        if [ -n "$current_content" ]; then
            new_content="${current_content},rootdir"
        else
            new_content="rootdir"
        fi
        pvesm set "$TARGET_STORAGE" --content "$new_content" || {
            echo -e "${RED}Failed to add rootdir. Please enable it manually.${NC}"
            exit 1
        }
    fi

    # ---- Path & Container Configuration ----
    echo ""
    read -p "Base Storage Path [/mnt/pve/host-shared/converter]: " PATH_INPUT
    STORAGE_PATH="${PATH_INPUT:-/mnt/pve/host-shared/converter}"
    STORAGE_PATH="${STORAGE_PATH%/}"

    read -p "Container ID [20020]: " CTID_INPUT
    CTID="${CTID_INPUT:-20020}"

    DEFAULT_BRIDGE=$(get_default_bridge)
    read -p "Network Bridge [default: ${DEFAULT_BRIDGE}]: " BRIDGE_INPUT
    CT_BRIDGE="${BRIDGE_INPUT:-$DEFAULT_BRIDGE}"

    SUGGESTED_GW=$(guess_gateway_from_bridge "$CT_BRIDGE")
    [ -z "$SUGGESTED_GW" ] && SUGGESTED_GW="10.2.0.1"

    read -p "Container IP/CIDR (e.g., 10.2.0.20/16): " CT_IP
    while [[ -z "$CT_IP" ]]; do
        echo -e "${RED}Error: Static IP with CIDR mask is required (e.g., 10.2.0.20/16).${NC}"
        read -p "Container IP/CIDR: " CT_IP
    done

    read -p "Gateway IP [${SUGGESTED_GW}]: " GW_INPUT
    CT_GW="${GW_INPUT:-$SUGGESTED_GW}"

    read -p "MAC Address (Press Enter to auto-generate): " MAC_INPUT
    CT_MAC="${MAC_INPUT:-}"

    echo ""
    echo -e "${GREEN}✓ Environment configuration gathered successfully.${NC}"
}

# ------------------------------------------------------------------------------
# Action Functions (Deploy, Update, Remove)
# ------------------------------------------------------------------------------
do_deploy() {
    echo -e "${GREEN}==> Beginning Deployment on Host $(hostname)...${NC}"

    if ! command -v pct &>/dev/null; then
        echo -e "${RED}Error: 'pct' CLI not found. Script must be executed on a Proxmox host.${NC}"
        exit 1
    fi

    gather_environment_info

    REAL_PATH=$(readlink -f "$STORAGE_PATH" 2>/dev/null || echo "$STORAGE_PATH")

    # Clean previous state
    if [ -d "$REAL_PATH" ] || pct status "$CTID" &>/dev/null; then
        echo -e "${YELLOW}! Cleaning pre-existing installation state at ${REAL_PATH} / CTID ${CTID}...${NC}"
        pct stop "$CTID" --skiplock 2>/dev/null || true
        pct destroy "$CTID" --purge 2>/dev/null || true
        systemctl stop "$SERVICE_NAME" 2>/dev/null || true
        rm -rf "$REAL_PATH"
    fi

    # Fix APT sources for Proxmox (no subscription)
    rm -f /etc/apt/sources.list.d/pve-enterprise.sources \
          /etc/apt/sources.list.d/ceph.sources \
          /etc/apt/sources.list.d/pve-enterprise.list \
          /etc/apt/sources.list.d/pve-no-subscription.list 2>/dev/null || true
    cat <<'EOF' > /etc/apt/sources.list.d/proxmox-no-sub.sources
Types: deb
URIs: http://download.proxmox.com/debian/pve
Suites: trixie
Components: pve-no-subscription
Signed-By: /usr/share/keyrings/proxmox-archive-keyring.gpg
EOF
    apt-get update -qq

    echo -e "${CYAN}==> [2/3] Provisioning Directory Structure & Backend Service${NC}"
    mkdir -p "$REAL_PATH"
    echo "Cloning source code from ${REPO_URL}..."
    git clone "$REPO_URL" "$REAL_PATH"

    CONVERTER_ROOT="$REAL_PATH"
    SCRIPT_DIR="${CONVERTER_ROOT}/backend"
    SERVICE_FILE="${SCRIPT_DIR}/${SERVICE_NAME}.service"

    mkdir -p "${CONVERTER_ROOT}/logs" "${CONVERTER_ROOT}/tmp"
    chmod 777 "${CONVERTER_ROOT}/tmp"

    apt-get install -y -qq --no-install-recommends \
        python3 python3-pip python3-venv libguestfs-tools \
        p7zip-full unzip unrar-free tar gzip bzip2 xz-utils git

    VENV_DIR="${CONVERTER_ROOT}/venv"
    echo "Setting up Python virtual environment..."
    python3 -m venv "$VENV_DIR"
    "$VENV_DIR/bin/pip" install --upgrade pip -q
    if [ -f "${SCRIPT_DIR}/requirements.txt" ]; then
        "$VENV_DIR/bin/pip" install -r "${SCRIPT_DIR}/requirements.txt" -q
    fi

    cat <<EOF > "$SERVICE_FILE"
[Unit]
Description=VM Bench Converter Backend Service
After=network.target

[Service]
Type=simple
User=root
WorkingDirectory=${SCRIPT_DIR}
ExecStart=${VENV_DIR}/bin/python ${SCRIPT_DIR}/app.py
Restart=always
RestartSec=5
StandardOutput=append:${CONVERTER_ROOT}/logs/vm-bench-backend.log
StandardError=append:${CONVERTER_ROOT}/logs/vm-bench-backend.log

[Install]
WantedBy=multi-user.target
EOF

    cp "$SERVICE_FILE" "${SYSTEMD_DIR}/${SERVICE_NAME}.service"
    systemctl daemon-reload
    systemctl enable --now "$SERVICE_NAME"
    echo -e "${GREEN}✓ Host Backend systemd service initialized and running.${NC}"

    echo -e "${CYAN}==> [3/3] Creating and Starting Frontend LXC Container (${CTID})${NC}"

    # Get LXC template
    TEMPLATE=$(pveam available 2>/dev/null | grep -m1 "debian-12.*amd64" | awk '{print $2}' || true)
    if [ -z "$TEMPLATE" ]; then
        pveam update 2>/dev/null || true
        pveam download local debian-12-standard_12.7-1_amd64.tar.zst 2>/dev/null || true
        TEMPLATE="local:vztmpl/debian-12-standard_12.7-1_amd64.tar.zst"
    else
        TEMPLATE_NAME=$(basename "$TEMPLATE")
        [ ! -f "/var/lib/vz/template/cache/${TEMPLATE_NAME}" ] && pveam download local "$TEMPLATE_NAME" 2>/dev/null || true
        TEMPLATE="local:vztmpl/${TEMPLATE_NAME}"
    fi

    NET_CFG="name=eth0,bridge=${CT_BRIDGE},firewall=1,gw=${CT_GW},ip=${CT_IP},type=veth"
    [ -n "$CT_MAC" ] && NET_CFG="${NET_CFG},hwaddr=${CT_MAC}"

    pct create "$CTID" "$TEMPLATE" \
        --arch amd64 \
        --cores "$CT_CORES" \
        --hostname vm-bench \
        --memory "$CT_MEMORY" \
        --net0 "$NET_CFG" \
        --ostype debian \
        --storage "$TARGET_STORAGE" \
        --rootfs "${TARGET_STORAGE}:${CT_DISK_SIZE}" \
        --swap "$CT_SWAP" \
        --unprivileged 0 \
        --features nesting=1 \
        --mp0 "${REAL_PATH},mp=/mnt/converter" \
        --onboot 1 \
        --start 0

    # Extra LXC config
    CONF_FILE="/etc/pve/lxc/${CTID}.conf"
    echo "lxc.cgroup2.devices.allow: a" >> "$CONF_FILE"
    echo "lxc.mount.auto: proc:rw sys:rw cgroup:rw" >> "$CONF_FILE"
    echo "lxc.apparmor.profile: unconfined" >> "$CONF_FILE"

    pct start "$CTID" 2>/dev/null || true

    echo ""
    echo -e "${GREEN}======================================================================${NC}"
    echo -e "${GREEN} Deployment Successful on $(hostname)!${NC}"
    echo " - Host Backend API    : http://127.0.0.1:9000"
    echo " - LXC Container IP    : ${CT_IP%/*}"
    echo " - Container Storage   : ${TARGET_STORAGE}"
    echo " - Mount Directory     : ${REAL_PATH} -> /mnt/converter"
    echo -e "${GREEN}======================================================================${NC}"
}

do_update() {
    echo -e "${CYAN}==> Updating VM Bench Stack...${NC}"
    read -p "Base Storage Path [/mnt/pve/host-shared/converter]: " PATH_INPUT
    STORAGE_PATH="${PATH_INPUT:-/mnt/pve/host-shared/converter}"
    REAL_PATH=$(readlink -f "$STORAGE_PATH" 2>/dev/null || echo "$STORAGE_PATH")

    read -p "Container ID [20020]: " CTID_INPUT
    CTID="${CTID_INPUT:-20020}"

    if [ -d "$REAL_PATH/.git" ]; then
        echo "Pulling latest code from Git..."
        git -C "$REAL_PATH" pull --ff-only
    fi

    if [ -d "${REAL_PATH}/venv" ]; then
        echo "Updating Python environment..."
        "${REAL_PATH}/venv/bin/pip" install --upgrade pip -q
        [ -f "${REAL_PATH}/backend/requirements.txt" ] && "${REAL_PATH}/venv/bin/pip" install -r "${REAL_PATH}/backend/requirements.txt" -q
    fi

    echo "Restarting Services..."
    systemctl restart "$SERVICE_NAME"
    pct restart "$CTID" 2>/dev/null || true
    echo -e "${GREEN}✓ Stack update complete.${NC}"
}

do_remove() {
    echo -e "${YELLOW}==> Tearing Down VM Bench Stack...${NC}"
    read -p "Base Storage Path [/mnt/pve/host-shared/converter]: " PATH_INPUT
    STORAGE_PATH="${PATH_INPUT:-/mnt/pve/host-shared/converter}"
    REAL_PATH=$(readlink -f "$STORAGE_PATH" 2>/dev/null || echo "$STORAGE_PATH")

    read -p "Container ID to Destroy [20020]: " CTID_INPUT
    CTID="${CTID_INPUT:-20020}"

    if pct status "$CTID" &>/dev/null; then
        echo "Destroying LXC Container ${CTID}..."
        pct stop "$CTID" --skiplock 2>/dev/null || true
        pct destroy "$CTID" --purge 2>/dev/null || true
    fi

    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null || systemctl is-enabled "$SERVICE_NAME" 2>/dev/null; then
        echo "Removing Host Backend Systemd Service..."
        systemctl stop "$SERVICE_NAME" 2>/dev/null || true
        systemctl disable "$SERVICE_NAME" 2>/dev/null || true
        rm -f "${SYSTEMD_DIR}/${SERVICE_NAME}.service"
        systemctl daemon-reload
    fi

    if [ -d "$REAL_PATH" ]; then
        echo "Purging directory ${REAL_PATH}..."
        rm -rf "$REAL_PATH"
    fi

    echo -e "${GREEN}✓ Removal complete.${NC}"
}

# ------------------------------------------------------------------------------
# CLI Dispatcher
# ------------------------------------------------------------------------------
case "${1:-}" in
    --deploy)
        do_deploy
        ;;
    --update)
        do_update
        ;;
    --remove)
        do_remove
        ;;
    *)
        echo -e "${CYAN}VM Bench Stack Lifecycle Manager (${HOSTNAME})${NC}"
        echo "Usage: $0 {--deploy|--update|--remove}"
        echo "  --deploy : Interactive installation, storage selection, and deployment"
        echo "  --update : Pull latest repository changes and refresh stack services"
        echo "  --remove : Purge LXC container, host systemd backend, and storage directories"
        exit 1
        ;;
esac