#!/bin/bash

# Exit immediately if a command exits with a non-zero status.
set -e

# Function to check for required utilities
function check_utilities() {
    local utilities=("qm" "wget" "xz" "sha256sum" "ssh-keygen")
    for util in "${utilities[@]}"; do
        command -v "$util" >/dev/null 2>&1 || { echo "$util not found. Please install it."; exit 1; }
    done
}

# Function to set up SSH keys
function setup_ssh_keys() {
    local default_ssh_key_dir="${HOME}/ssh-keys"
    local ssh_key_dir="$default_ssh_key_dir"

    while true; do
        # Ensure the SSH key directory exists
        if [[ ! -d "${ssh_key_dir}" ]]; then
            echo "SSH key directory not found at ${ssh_key_dir}. Creating directory..."
            mkdir -p "${ssh_key_dir}"
            chmod 700 "${ssh_key_dir}"
        fi

        # Determine the SSH public key file
        ssh_keyfile=$(ls "${ssh_key_dir}"/*.pub 2>/dev/null | head -n 1)

        if [[ -z "${ssh_keyfile}" ]]; then
            echo "No SSH public key found in ${ssh_key_dir}."
            read -p "Do you want to generate a new SSH key pair in this directory? (yes/no): " generate_key
            if [[ "${generate_key}" == "yes" ]]; then
                ssh-keygen -t rsa -b 4096 -f "${ssh_key_dir}/id_rsa" -N ""
                ssh_keyfile="${ssh_key_dir}/id_rsa.pub"
                chmod 600 "${ssh_key_dir}/id_rsa"
                chmod 644 "${ssh_key_dir}/id_rsa.pub"
                echo "SSH key pair generated at ${ssh_key_dir}/id_rsa and ${ssh_keyfile}."
                break
            else
                read -p "Do you want to specify a different directory for your SSH keys? (yes/no): " change_dir
                if [[ "${change_dir}" == "yes" ]]; then
                    read -p "Please enter the full path to your SSH key directory: " ssh_key_dir
                else
                    echo "Cannot proceed without an SSH public key. Exiting."
                    exit 1
                fi
            fi
        else
            echo "Using existing SSH public key: ${ssh_keyfile}"
            break
        fi
    done

    chmod 700 "${ssh_key_dir}"
    chmod 600 "${ssh_key_dir}"/id_* 2>/dev/null || true
    chmod 644 "${ssh_key_dir}"/*.pub 2>/dev/null || true
}

# Function to generate cloud-init user data with CloudPanel installation
function generate_cloud_init_config() {
    local vm_id="$1"
    cat > "vm-${vm_id}-cloud-init.yaml" << 'EOF'
#cloud-config
package_update: true
package_upgrade: true
packages:
  - curl
  - wget
  - sudo
  - apt-transport-https
  - ca-certificates

runcmd:
  - curl -sS https://installer.cloudpanel.io/ce/v2/install.sh -o /root/install.sh
  - echo "a3ba69a8102345127b4ae0e28cfe89daca675cbc63cd39225133cdd2fa02ad36 /root/install.sh" | sha256sum -c
  - DB_ENGINE=MARIADB_11.4 bash /root/install.sh

power_state:
  mode: reboot
  timeout: 1800
  condition: True
EOF
    
    # Import the cloud-init config
    qm set "${vm_id}" --cicustom "user=local:snippets/vm-${vm_id}-cloud-init.yaml"
}

# Function to create or update a template
function create_template() {
    local vm_id="$1"
    local vm_name="$2"
    local image_file="$3"

    echo "Processing template ${vm_name} (${vm_id})"

    # Compute the checksum of the image file
    local image_checksum
    image_checksum=$(sha256sum "${image_file}" | awk '{print $1}')
    local checksum_file="checksums/${vm_id}.sha256"

    # Check if template exists and handle accordingly
    if qm status "${vm_id}" &>/dev/null; then
        echo "Template with VM ID ${vm_id} already exists. Removing..."
        qm destroy "${vm_id}" --destroy-unreferenced-disks yes
    fi

    echo "Creating template ${vm_name} (${vm_id})"

    # Create new VM
    qm create "${vm_id}" --name "${vm_name}" --ostype l26

    # Set networking to default bridge
    qm set "${vm_id}" --net0 virtio,bridge=vmbr0

    # Set display to serial
    qm set "${vm_id}" --serial0 socket --vga serial0

    # Set memory and CPU (increased for CloudPanel requirements)
    qm set "${vm_id}" --memory 2048 --cores 2 --cpu host

    # Import the disk
    qm set "${vm_id}" --scsi0 "${storage}:0,import-from=${PWD}/${image_file},discard=on"

    # Set SCSI hardware as default boot disk
    qm set "${vm_id}" --boot order=scsi0 --scsihw virtio-scsi-single

    # Enable QEMU guest agent
    qm set "${vm_id}" --agent enabled=1,fstrim_cloned_disks=1

    # Add cloud-init device
    qm set "${vm_id}" --ide2 "${storage}:cloudinit"

    # Set cloud-init network configuration
    qm set "${vm_id}" --ipconfig0 "ip=dhcp,ip6=auto"

    # Import the SSH keyfile
    if [[ -f "${ssh_keyfile}" ]]; then
        qm set "${vm_id}" --sshkeys "${ssh_keyfile}"
    else
        echo "SSH key file not found at ${ssh_keyfile}"
        exit 1
    fi

    # Add the user
    qm set "${vm_id}" --ciuser "${username}"

    # Generate and apply cloud-init config with CloudPanel installation
    generate_cloud_init_config "${vm_id}"

    # Resize the disk to 25G for CloudPanel requirements
    qm disk resize "${vm_id}" scsi0 25G

    # Convert the VM into a template
    qm template "${vm_id}"

    # Save the checksum
    mkdir -p checksums
    echo "${image_checksum}" > "${checksum_file}"

    # Remove the image file when done
    rm -f "${image_file}"
}

# Check for required utilities
check_utilities

# User-configurable variables
export username="hackiri"        # Replace with your desired username
export storage="ceph_local"      # Replace with your Proxmox storage name

# Validate variables
if [[ -z "${username}" || "${username}" == "your_username_here" ]]; then
    echo "Please set a valid username in the script."
    exit 1
fi

if ! pvesm status | grep -q "^${storage}\s"; then
    echo "Storage '${storage}' not found. Please check your Proxmox storage configuration."
    exit 1
fi

# Set up SSH keys
setup_ssh_keys

# CloudPanel works best with Ubuntu 24.04 (Lunar Lobster)
declare -a images=(
    # Ubuntu 24.04 LTS (Lunar Lobster) with CloudPanel
    "912|ubuntu-24.04-template|https://cloud-images.ubuntu.com/releases/24.04/release/ubuntu-24.04-server-cloudimg-amd64.img"
)

# Loop through the images array
for entry in "${images[@]}"; do
    IFS='|' read -r vm_id vm_name image_url <<< "${entry}"

    # Extract the filename from the URL
    image_file="${image_url##*/}"

    # Download the image with timestamping
    echo "Downloading ${image_file}..."
    if ! wget -N "${image_url}"; then
        echo "Failed to download ${image_url}"
        exit 1
    fi

    # Create or update the template
    create_template "${vm_id}" "${vm_name}" "${image_file}"
done