#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
SCRIPT_NAME=$(basename "$0")

# Default images directory — override with PVE_CLOUDIMG_DIR env var or -d flag
DEFAULT_IMAGES_DIR="${PVE_CLOUDIMG_DIR:-$SCRIPT_DIR/cloud-images}"
IMAGES_DIR="$DEFAULT_IMAGES_DIR"

# Auto-sudo if not root
if [[ "$(id -u)" -ne 0 ]]; then
    echo "[*] Elevating to root via sudo..." >&2
    exec sudo "$0" "$@"
fi

# Default settings
PVE_STORAGE="main-storage"
PVE_BRIDGE="vmbr0"
PVE_VLAN=""
PVE_BIOS="seabios"
VM_CORES=2
VM_MEMORY=2048
VM_DISK_SIZE=10
VM_ID=""
CLOUD_USER=""
SSH_KEY_FILE=""
INSTALL_AGENT=1
EXTRA_PACKAGES=""
FORCE_RELOAD=0
FORCE_YES=0

# Distro definitions (distro:version => URL, filename, user, format)
# Format: url|filename|default_user|image_format
declare -A DISTROS
declare -A DISTRO_PACKAGES

# Each distro gets: most current version + one back.
# Each DISTROS entry can be followed by a DISTRO_PACKAGES entry for extra packages (space-separated).

# Ubuntu LTS — current 26.04 (resolute), previous 24.04 (noble)
DISTROS["ubuntu:26.04"]="https://cloud-images.ubuntu.com/resolute/current/resolute-server-cloudimg-amd64.img|resolute-server-cloudimg-amd64.img|ubuntu|img"
DISTRO_PACKAGES["ubuntu:26.04"]=""
# ---
DISTROS["ubuntu:24.04"]="https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img|noble-server-cloudimg-amd64.img|ubuntu|img"
DISTRO_PACKAGES["ubuntu:24.04"]=""

# Debian — current stable 13 (trixie), previous 12 (bookworm)
DISTROS["debian:13"]="https://cloud.debian.org/images/cloud/trixie/latest/debian-13-genericcloud-amd64.qcow2|debian-13-genericcloud-amd64.qcow2|debian|qcow2"
DISTRO_PACKAGES["debian:13"]=""
# ---
DISTROS["debian:12"]="https://cloud.debian.org/images/cloud/bookworm/latest/debian-12-genericcloud-amd64.qcow2|debian-12-genericcloud-amd64.qcow2|debian|qcow2"
DISTRO_PACKAGES["debian:12"]=""

# Arch Linux — rolling release
DISTROS["arch:rolling"]="https://geo.mirror.pkgbuild.com/images/latest/Arch-Linux-x86_64-cloudimg.qcow2|Arch-Linux-x86_64-cloudimg.qcow2|arch|qcow2"
DISTRO_PACKAGES["arch:rolling"]=""

declare -A DISTRO_DEFAULT_VERSION=(
    [ubuntu]="26.04"
    [debian]="13"
    [arch]="rolling"
)

# Proxmox template name prefix
TEMPLATE_PREFIX="cloud"

log() { echo "[*] $*" >&2; }
warn() { echo "[!] $*" >&2; }
die() { echo "[FATAL] $*" >&2; exit 1; }

usage() {
    cat <<EOF
Usage: $SCRIPT_NAME [options] <command> [args]

Commands:
  list [distro]           List available distros, cached images, and VM templates
  download <distro> [ver] Download image into cache only (no modifications)
                           Use "all" to download every supported distro
  create <distro> [ver]   Download cloud image (if needed) and create VM template
  update <distro> [ver]   Re-download and re-create a template (remove + create)
  remove <distro> [ver]   Remove cached image and VM template for a distro

Options:
  -d, --dir DIR           Image cache directory (default: script-dir/cloud-images; set \$PVE_CLOUDIMG_DIR to change default)
  --storage STORAGE       Proxmox storage for VM disks (default: $PVE_STORAGE)
  --bridge BRIDGE         Network bridge (default: $PVE_BRIDGE)
  --vlan VLAN             VLAN tag for the network interface
  --bios BIOS             Firmware: seabios (default) or ovmf (UEFI)
  --vm-id VMID            VM ID for the template (default: auto from 9000)
  --cores CORES           CPU cores (default: $VM_CORES)
  --memory MEMORY         Memory in MB (default: $VM_MEMORY)
  --disk-size SIZE        Disk size in GB (default: $VM_DISK_SIZE)
  --user USER             Cloud-init default user (default: distro-specific)
  --ssh-key FILE          SSH public key file to inject via cloud-init
  --no-agent              Skip installing qemu-guest-agent
  --packages PKGS         Extra packages to install (comma/space-separated, passed to virt-customize)
  -f, --force             Re-download image even if already cached
  -y, --yes               Skip confirmation prompts

Examples:
  $SCRIPT_NAME list
  $SCRIPT_NAME list ubuntu
  $SCRIPT_NAME create ubuntu 24.04
  $SCRIPT_NAME create ubuntu 24.04 --storage local-btrfs --disk-size 64
  $SCRIPT_NAME create debian --user admin --ssh-key ~/.ssh/id_rsa.pub
  $SCRIPT_NAME remove ubuntu 24.04
EOF
    exit 1
}

# ---- Distro helpers ----

get_distro_field() {
    local distro=$1 version=$2 field=$3
    local key="${distro}:${version}"
    local data="${DISTROS[$key]:-}"
    if [[ -z "$data" ]]; then
        return 1
    fi
    case "$field" in
        url)    echo "${data%%|*}" ;;
        file)   echo "${data#*|}" | cut -d'|' -f1 ;;
        user)   echo "${data#*|}" | cut -d'|' -f2 ;;
        format) echo "${data#*|}" | cut -d'|' -f3 ;;
        *)      return 1 ;;
    esac
}

validate_distro_version() {
    local distro=$1 version=$2
    local key="${distro}:${version}"
    if [[ -z "${DISTROS[$key]:-}" ]]; then
        return 1
    fi
    return 0
}

list_known_versions() {
    local distro=$1
    local versions=()
    for key in "${!DISTROS[@]}"; do
        if [[ "$key" == "${distro}:"* ]]; then
            versions+=("${key#*:}")
        fi
    done
    echo "${versions[@]}"
}

list_known_distros() {
    local distros=()
    for key in "${!DISTROS[@]}"; do
        local d="${key%%:*}"
        local seen=0
        for seen_d in "${distros[@]}"; do
            if [[ "$seen_d" == "$d" ]]; then seen=1; break; fi
        done
        if [[ $seen -eq 0 ]]; then
            distros+=("$d")
        fi
    done
    echo "${distros[@]}"
}

get_default_version() {
    local distro=$1
    echo "${DISTRO_DEFAULT_VERSION[$distro]:-}"
}

image_cache_dir() {
    local distro=$1 version=$2
    echo "$IMAGES_DIR/${distro}-${version}"
}

template_name() {
    local distro=$1 version=$2
    # Replace dots with dashes for VM name
    local ver_safe="${version//./-}"
    echo "${TEMPLATE_PREFIX}-${distro}-${ver_safe}"
}

# ---- Proxmox helpers ----

check_proxmox() {
    if ! command -v qm &>/dev/null; then
        die "Proxmox VE not found — this script must be run on a PVE host"
    fi
}

check_storage() {
    if ! pvesm status 2>/dev/null | awk -v s="$PVE_STORAGE" 'NR>1 && $1==s {f=1} END{exit !f}'; then
        die "Storage '$PVE_STORAGE' not found — available: $(pvesm status 2>/dev/null | awk 'NR>1{print $1}' | tr '\n' ' ')"
    fi
    log "Storage '$PVE_STORAGE' is ready"
}

check_virt_customize() {
    if ! command -v virt-customize &>/dev/null; then
        warn "virt-customize not found — install libguestfs-tools to enable guest agent injection"
        warn "  apt-get update && apt-get install -y libguestfs-tools"
        return 1
    fi
    return 0
}

next_available_vmid() {
    local start=${1:-9000}
    while qm status "$start" &>/dev/null; do
        start=$((start + 1))
    done
    echo "$start"
}

template_exists() {
    local name=$1
    qm list 2>/dev/null | awk -v name="$name" '$2 == name && $3 == "template" {print $1}' | head -1
}

# ---- Commands ----

list_templates() {
    local prefix="${1:-}"
    for conf in /etc/pve/nodes/*/qemu-server/*.conf; do
        [[ -f "$conf" ]] || continue
        grep -q "^template: 1" "$conf" 2>/dev/null || continue
        local vmid name
        vmid=$(basename "$conf" .conf)
        name=$(grep "^name:" "$conf" 2>/dev/null | sed 's/^name: //')
        if [[ -z "$prefix" ]] || [[ "$name" == "${prefix}"* ]]; then
            echo "  VM $vmid  $name"
        fi
    done
}

cmd_list() {
    local filter="${1:-}"

    if [[ -n "$filter" ]]; then
        if [[ -z "${DISTRO_DEFAULT_VERSION[$filter]:-}" ]]; then
            warn "Unknown distro: $filter"
            return 1
        fi
        local versions
        versions=$(list_known_versions "$filter")
        echo "Distribution: $filter"
        echo "  Known versions: $versions"
        echo "  Default version: $(get_default_version "$filter")"
        echo ""
        echo "  Cached images:"
        if [[ -d "$IMAGES_DIR" ]]; then
            local found=0
            for dir in "$IMAGES_DIR/${filter}-"*/; do
                if [[ -d "$dir" ]] && (compgen -G "$dir"*.img &>/dev/null || compgen -G "$dir"*.qcow2 &>/dev/null || compgen -G "$dir"*.raw &>/dev/null); then
                    local dirname size
                    dirname=$(basename "$dir")
                    size=$(du -sh "$dir" 2>/dev/null | cut -f1)
                    echo "    $dirname ($size)"
                    found=1
                fi
            done
            [[ $found -eq 0 ]] && echo "    (none)"
        else
            echo "    (none)"
        fi
        echo ""
        echo "  Templates:"
        local tmpl
        tmpl=$(list_templates "${TEMPLATE_PREFIX}-${filter}")
        echo "${tmpl:-    (none)}"
        return
    fi

    echo "Available distributions:"
    local distros
    distros=$(list_known_distros)
    for d in $distros; do
        local versions
        versions=$(list_known_versions "$d")
        local default
        default=$(get_default_version "$d")
        echo "  $d (default: $default, versions: $versions)"
    done
    echo ""

    echo "Cached images:"
    if [[ -d "$IMAGES_DIR" ]]; then
        local found=0
        for dir in "$IMAGES_DIR"/*/; do
            local dirname
            dirname=$(basename "$dir")
            local has_files=0
            for ext in img qcow2 raw; do
                if compgen -G "$dir"*."$ext" &>/dev/null; then
                    has_files=1; break
                fi
            done
            if [[ $has_files -eq 1 ]]; then
                local size
                size=$(du -sh "$dir" 2>/dev/null | cut -f1)
                echo "  $dirname ($size)"
                found=1
            fi
        done
        [[ $found -eq 0 ]] && echo "  (none)"
    else
        echo "  (none)"
    fi
    echo ""

    echo "Proxmox templates:"
    local tmpl
    tmpl=$(list_templates)
    echo "${tmpl:-  (none)}"
}

# ---- Shared image preparation ----

# Resolve distro metadata, download to cache, optionally customize.
# Returns: <cache_path>|<ci_user>|<tname>
prepare_image() {
    local distro=$1 version=${2:-}

    if [[ -z "$version" ]]; then
        version=$(get_default_version "$distro")
        if [[ -z "$version" ]]; then
            die "Unknown distro '$distro'. Available: $(list_known_distros)"
        fi
        log "Using default version: $distro $version"
    fi

    if ! validate_distro_version "$distro" "$version"; then
        die "Unknown distro/version: $distro $version"
    fi

    local image_file
    image_file=$(get_distro_field "$distro" "$version" "file")
    local image_url
    image_url=$(get_distro_field "$distro" "$version" "url")
    local default_user
    default_user=$(get_distro_field "$distro" "$version" "user")

    local cache_dir
    cache_dir=$(image_cache_dir "$distro" "$version")
    local cache_path="$cache_dir/$image_file"
    local tname
    tname=$(template_name "$distro" "$version")
    local ci_user="${CLOUD_USER:-$default_user}"

    log "Distribution: $distro $version"
    log "Cloud-init user: $ci_user"

    # Download image if not cached (or force re-download)
    if [[ -f "$cache_path" ]]; then
        if [[ "$FORCE_RELOAD" -eq 1 ]]; then
            log "Forcing re-download: removing cached image..."
            rm -f "$cache_path"
        else
            log "Image already cached: $cache_path"
        fi
    fi

    if [[ ! -f "$cache_path" ]]; then
        log "Downloading $distro $version..."
        mkdir -p "$cache_dir"
        log "  URL: $image_url"
        curl -fL -o "$cache_path" "$image_url" || {
            rm -f "$cache_path"
            die "Download failed"
        }
        log "Downloaded: $(du -h "$cache_path" | cut -f1)"
    fi

    echo "${cache_path}|${ci_user}|${tname}|${image_file}"
}

cmd_download() {
    local distro=$1 version=${2:-}

    # Special case: download --all or download all
    if [[ "$distro" == "all" ]] || [[ "$distro" == "--all" ]]; then
        log "Downloading all supported distro versions..."
        local all_keys=()
        for key in "${!DISTROS[@]}"; do
            all_keys+=("$key")
        done
        # Sort for consistent order
        IFS=$'\n' all_keys=($(sort <<<"${all_keys[*]}")); unset IFS
        for key in "${all_keys[@]}"; do
            local d="${key%%:*}" v="${key#*:}"
            echo ""
            log "=== $d $v ==="
            cmd_download "$d" "$v"
        done
        echo ""
        log "All images downloaded."
        return
    fi

    local result
    result=$(prepare_image "$distro" "$version")
    local cache_path ci_user tname image_file
    cache_path=$(echo "$result" | cut -d'|' -f1)
    ci_user=$(echo "$result" | cut -d'|' -f2)
    tname=$(echo "$result" | cut -d'|' -f3)
    image_file=$(echo "$result" | cut -d'|' -f4)

    echo ""
    log "Image ready at: $cache_path"
    echo "  Distro:      $distro $version"
    echo "  Size:        $(du -h "$cache_path" | cut -f1)"
    echo "  Cloud-init:  user=$ci_user"
    echo ""
    echo "To create a VM template from this image:"
    echo "  $SCRIPT_NAME create $distro $version"
}

cmd_create() {
    local distro=$1 version=${2:-}

    local result
    result=$(prepare_image "$distro" "$version")
    local cache_path ci_user tname image_file
    cache_path=$(echo "$result" | cut -d'|' -f1)
    ci_user=$(echo "$result" | cut -d'|' -f2)
    tname=$(echo "$result" | cut -d'|' -f3)
    image_file=$(echo "$result" | cut -d'|' -f4)

    echo ""
    log "Template name: $tname"

    # Check for existing template
    local existing_vmid
    existing_vmid=$(template_exists "$tname")
    if [[ -n "$existing_vmid" ]]; then
        if [[ "$FORCE_YES" -ne 1 ]]; then
            warn "Template '$tname' already exists as VM $existing_vmid"
            read -r -p "Overwrite? This will destroy the existing template. [y/N] " reply
            if [[ ! "$reply" =~ ^[Yy]$ ]]; then
                die "Aborted."
            fi
        fi
        log "Removing existing template VM $existing_vmid..."
        qm stop "$existing_vmid" &>/dev/null || true
        qm destroy "$existing_vmid" --purge
    fi

    # Resolve VM ID
    local vmid="$VM_ID"
    if [[ -z "$vmid" ]]; then
        vmid=$(next_available_vmid 9000)
    fi

    # Build network config
    local netdev="virtio,bridge=${PVE_BRIDGE}"
    if [[ -n "$PVE_VLAN" ]]; then
        netdev+=",tag=${PVE_VLAN}"
    fi

    # Create VM
    local machine_args=()
    if [[ "$PVE_BIOS" == "ovmf" ]]; then
        machine_args=(--machine q35 --efidisk0 "${PVE_STORAGE}:0,format=raw,efitype=4m")
    fi

    log "Creating VM $vmid..."
    qm create "$vmid" \
        --name "$tname" \
        --memory "$VM_MEMORY" \
        --cores "$VM_CORES" \
        --net0 "$netdev" \
        --serial0 socket \
        --vga serial0 \
        --agent enabled=1 \
        "${machine_args[@]}" \
        --description "Cloud image template: $distro $version (created $(date -u +%Y-%m-%d))"

    # Build package list for first-boot install
    local PKGLIST=""
    [[ "$INSTALL_AGENT" -eq 1 ]] && PKGLIST="qemu-guest-agent"
    local distro_pkgs="${DISTRO_PACKAGES[${distro}:${version}]:-}"
    if [[ -n "$distro_pkgs" ]]; then PKGLIST="$PKGLIST $distro_pkgs"; fi
    if [[ -n "$EXTRA_PACKAGES" ]]; then PKGLIST="$PKGLIST $EXTRA_PACKAGES"; fi

    # Create a working copy: resize + firstboot installer, then import
    local import_image="$cache_path"
    local workdir=""
    if [[ "$VM_DISK_SIZE" -gt 0 ]] || [[ -n "$PKGLIST" ]]; then
        workdir=$(mktemp -d)
        import_image="$workdir/$image_file"
        cp "$cache_path" "$import_image"

        if [[ "$VM_DISK_SIZE" -gt 0 ]]; then
            log "Resizing disk to ${VM_DISK_SIZE}G..."
            qemu-img resize "$import_image" "${VM_DISK_SIZE}G" >&2
        fi

        if [[ -n "$PKGLIST" ]] && check_virt_customize; then
            log "Scheduling first-boot install:${PKGLIST}"
            virt-customize -a "$import_image" --firstboot-command "
                set -eux
                command -v apt-get && apt-get update && apt-get install -y ${PKGLIST} && { systemctl enable --now qemu-guest-agent 2>/dev/null || true; } && exit 0
                command -v pacman && pacman -S --noconfirm ${PKGLIST} && { systemctl enable --now qemu-guest-agent 2>/dev/null || true; } && exit 0
            " >&2
        fi
    fi

    # Import disk
    check_storage
    log "Importing disk to $PVE_STORAGE..."
    qm importdisk "$vmid" "$import_image" "$PVE_STORAGE"

    # Read the exact volume reference from the VM config (handles all storage types)
    local volume_ref
    volume_ref=$(qm config "$vmid" | awk '/^unused/ {print $2; exit}')
    if [[ -z "$volume_ref" ]]; then
        die "Disk import succeeded but no unused disk found in VM $vmid config"
    fi
    log "Imported volume: $volume_ref"

    # Attach disk (VirtIO SCSI — required by images lacking virtio-blk driver)
    qm set "$vmid" \
        --scsihw virtio-scsi-pci \
        --scsi0 "$volume_ref,discard=on,ssd=1" \
        --ide2 "$PVE_STORAGE:cloudinit" \
        --boot order=scsi0

    # Configure cloud-init
    log "Configuring cloud-init..."
    qm set "$vmid" --ciuser "$ci_user"

    if [[ -n "$SSH_KEY_FILE" ]]; then
        if [[ -f "$SSH_KEY_FILE" ]]; then
            qm set "$vmid" --sshkeys "$SSH_KEY_FILE"
        else
            warn "SSH key file not found: $SSH_KEY_FILE"
        fi
    fi

    qm set "$vmid" --ipconfig0 ip=dhcp

    # Convert to template
    log "Converting VM $vmid to template..."
    qm template "$vmid"

    # Clean up working copy
    if [[ -n "$workdir" ]]; then
        rm -rf "$workdir"
    fi

    echo ""
    log "Template created successfully!"
    echo "  VM ID:     $vmid"
    echo "  Name:      $tname"
    echo "  Distro:    $distro $version"
    echo "  User:      $ci_user"
    echo "  Disk:      ${VM_DISK_SIZE}G on $PVE_STORAGE"
    echo "  Memory:    ${VM_MEMORY}MB"
    echo "  Cores:     $VM_CORES"
    echo ""
    echo "To clone: qm clone $vmid <new-vmid> --name <vm-name>"
    echo "To create from CLI: qm clone $vmid 100 --name my-vm --full"
}

cmd_remove() {
    local distro=$1 version=${2:-}

    if [[ -z "$version" ]]; then
        version=$(get_default_version "$distro")
        if [[ -z "$version" ]]; then
            die "Unknown distro '$distro'"
        fi
        log "Using default version: $distro $version"
    fi

    if ! validate_distro_version "$distro" "$version"; then
        die "Unknown distro/version: $distro $version"
    fi

    local tname
    tname=$(template_name "$distro" "$version")

    # Remove VM template
    local existing_vmid
    existing_vmid=$(template_exists "$tname")
    if [[ -n "$existing_vmid" ]]; then
        if [[ "$FORCE_YES" -ne 1 ]]; then
            warn "Template '$tname' exists as VM $existing_vmid"
            read -r -p "Destroy this template? [y/N] " reply
            if [[ ! "$reply" =~ ^[Yy]$ ]]; then
                log "Skipping template removal."
            else
                log "Destroying template VM $existing_vmid..."
                qm stop "$existing_vmid" &>/dev/null || true
                qm destroy "$existing_vmid" --purge
            fi
        else
            log "Destroying template VM $existing_vmid..."
            qm stop "$existing_vmid" &>/dev/null || true
            qm destroy "$existing_vmid" --purge
        fi
    else
        log "No template found for $distro $version"
    fi

    # Remove cached image
    local cache_dir
    cache_dir=$(image_cache_dir "$distro" "$version")
    if [[ -d "$cache_dir" ]]; then
        if [[ "$FORCE_YES" -ne 1 ]]; then
            read -r -p "Remove cached image at $cache_dir? [y/N] " reply2
            if [[ "$reply2" =~ ^[Yy]$ ]]; then
                rm -rf "$cache_dir"
                log "Removed: $cache_dir"
            else
                log "Keeping cache."
            fi
        else
            rm -rf "$cache_dir"
            log "Removed: $cache_dir"
        fi
    else
        log "No cached image for $distro $version"
    fi

    log "Done."
}

# ---- Main ----

# Gather all args, parse options anywhere (before or after command)
ARGS=()
COMMAND=""
while [[ $# -gt 0 ]]; do
    case "$1" in
        -d|--dir) IMAGES_DIR="$2"; shift 2 ;;
        --storage) PVE_STORAGE="$2"; shift 2 ;;
        --bridge) PVE_BRIDGE="$2"; shift 2 ;;
        --vlan) PVE_VLAN="$2"; shift 2 ;;
        --bios) PVE_BIOS="$2"; shift 2 ;;
        --vm-id) VM_ID="$2"; shift 2 ;;
        --cores) VM_CORES="$2"; shift 2 ;;
        --memory) VM_MEMORY="$2"; shift 2 ;;
        --disk-size) VM_DISK_SIZE="$2"; shift 2 ;;
        --user) CLOUD_USER="$2"; shift 2 ;;
        --ssh-key) SSH_KEY_FILE="$2"; shift 2 ;;
        --no-agent) INSTALL_AGENT=0; shift ;;
        --packages) EXTRA_PACKAGES="$2"; shift 2 ;;
        -f|--force) FORCE_RELOAD=1; shift ;;
        -y|--yes) FORCE_YES=1; shift ;;
        --) shift; break ;;
        -*) usage ;;
        *)
            if [[ -z "$COMMAND" ]]; then
                COMMAND="$1"
            else
                ARGS+=("$1")
            fi
            shift
            ;;
    esac
done
# Append any remaining args after --
ARGS+=("$@")

[[ -z "$COMMAND" ]] && usage
set -- "${ARGS[@]}"

case "$COMMAND" in
    list)
        cmd_list "$@"
        ;;
    download)
        [[ $# -lt 1 ]] && usage
        cmd_download "$@"
        ;;
    create|update)
        check_proxmox
        [[ $# -lt 1 ]] && usage
        if [[ "$COMMAND" == "update" ]]; then
            cmd_remove "$@" || true
        fi
        cmd_create "$@"
        ;;
    remove)
        check_proxmox
        [[ $# -lt 1 ]] && usage
        cmd_remove "$@"
        ;;
    *)
        usage
        ;;
esac