#!/bin/bash
set -euo pipefail

readonly MODE_YUM="yum"
readonly MODE_RPM="rpm"

readonly GA_SERVICE="virtcca-guest-agent.service"

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

declare -A CFG
VC_ACTIONS=()

CFG[image_path]=""
CFG[install_mode]=""
CFG[iso_repo]=""
CFG[iso_mount]="/mnt/iso"
CFG[vc_attach]=""

CFG[rpm_files]=""
CFG[ssl_enabled]="false"
CFG[ssl_cert_src]=""
CFG[ssl_key_src]=""
CFG[ssl_ca_src]=""
CFG[ssl_vm_dir]="/etc/virtcca_deploy/cert"

CFG[temp_dir]=""
CFG[log_file]="/tmp/inject_${RANDOM}.log"

# ==============================================================================
# log (safe)
# ==============================================================================

log() {
    local level="$1"
    local msg="$2"

    if [[ "$level" == "INFO" ]]; then
        echo -e "${GREEN}[INFO] $msg${NC}"
    elif [[ "$level" == "WARN" ]]; then
        echo -e "${YELLOW}[WARN] $msg${NC}"
    else
        echo -e "[ERROR] $msg"
    fi
}

die() { log ERROR "$1"; exit 1; }

# ==============================================================================
# args
# ==============================================================================

parse_args() {

    local has_mode=0

    while [[ $# -gt 0 ]]; do
        case "$1" in

            -i)
                CFG[image_path]="$2"
                shift 2
                ;;

            --yum-install)
                if [[ $has_mode -eq 1 ]]; then
                    die "install mode already set, cannot mix with other modes"
                fi
                CFG[install_mode]=$MODE_YUM
                has_mode=1
                shift
                ;;

            --iso-repo)
                CFG[iso_repo]="$2"
                shift 2
                ;;

            --rpm-file)
                CFG[rpm_files]="${CFG[rpm_files]} $2"
                shift 2
                ;;

            --ssl)
                CFG[ssl_enabled]="true"
                shift
                ;;

            --ssl-cert-src)
                CFG[ssl_cert_src]="$2"
                shift 2
                ;;

            --ssl-key-src)
                CFG[ssl_key_src]="$2"
                shift 2
                ;;

            --ssl-ca-src)
                CFG[ssl_ca_src]="$2"
                shift 2
                ;;

            -h|--help)
                show_usage
                exit 0
                ;;

            *)
                die "unknown arg: $1"
                ;;
        esac
    done

    if [[ -z "${CFG[image_path]}" ]]; then
        die "image path (-i) is required"
    fi

    if [[ -z "${CFG[install_mode]}" ]]; then
        die "install mode not set (use --yum-install | --rpm-file)"
    fi

    if [[ -n "${CFG[iso_repo]}" ]]; then
        if [[ "${CFG[install_mode]}" != "$MODE_YUM" ]]; then
            die "Option --iso-repo is only valid in yum install mode (--yum-install)"
        fi
    fi
}

# ==============================================================================
# validate
# ==============================================================================

validate() {

    [[ -f "${CFG[image_path]}" ]] || die "image missing"

    local m=0
    [[ "${CFG[install_mode]}" == "$MODE_YUM" ]] && m=$((m+1))
    [[ "${CFG[install_mode]}" == "$MODE_RPM" ]] && m=$((m+1))

    [[ "$m" -eq 1 ]] || die "choose one mode"

    if [[ "${CFG[install_mode]}" == "$MODE_RPM" ]]; then
        [[ -n "${CFG[rpm_files]}" ]] || die "rpm files required"
    fi
}

# ==============================================================================
# temp
# ==============================================================================

prepare() {
    CFG[temp_dir]="$(mktemp -d /tmp/ga_XXXX)"

    if [[ "${CFG[install_mode]}" == "$MODE_YUM" && -n "${CFG[iso_repo]}" ]]; then
        CFG[vc_attach]="--attach ${CFG[iso_repo]}"
    fi
}

# ==============================================================================
# ssl
# ==============================================================================

ssl_prepare() {

    [[ "${CFG[ssl_enabled]}" == "true" ]] || return 0

    mkdir -p "${CFG[temp_dir]}/cert"

    cp "${CFG[ssl_cert_src]}" "${CFG[temp_dir]}/cert/agent.crt"
    cp "${CFG[ssl_key_src]}" "${CFG[temp_dir]}/cert/agent.key"
    cp "${CFG[ssl_ca_src]}" "${CFG[temp_dir]}/cert/ca.crt"
}

# ==============================================================================
# build plan
# ==============================================================================

vc_add() {
    VC_ACTIONS+=("$@")
}

ssl_add() {
    if [[ "${CFG[ssl_enabled]}" == "true" ]]; then
        vc_add --copy-in "${CFG[temp_dir]}/cert:/etc/virtcca_deploy/"
        vc_add --run-command "mkdir -p /etc/virtcca_deploy/cert"
        vc_add --run-command "chown -R root:virtcca /etc/virtcca_deploy"
        vc_add --run-command "chmod 750 /etc/virtcca_deploy"
        vc_add --run-command "chmod 750 /etc/virtcca_deploy/cert"
        vc_add --run-command "chmod 640 /etc/virtcca_deploy/cert/agent.crt || true"
        vc_add --run-command "chmod 640 /etc/virtcca_deploy/cert/ca.crt || true"
        vc_add --run-command "chmod 640 /etc/virtcca_deploy/cert/agent.key || true"
        vc_add --run-command "chown root:virtcca /etc/virtcca_deploy/cert/* || true"
    fi
}

build() {

    case "${CFG[install_mode]}" in

        "yum")
            if [[ -n "${CFG[iso_repo]}" ]]; then
                vc_add --run-command "mkdir -p ${CFG[iso_mount]}"
                vc_add ${CFG[vc_attach]}
                vc_add --run-command "mount -o ro /dev/sdb ${CFG[iso_mount]}"
                vc_add --run-command "cat > /etc/yum.repos.d/local.repo <<EOF
[offline-repo]
name=Offline OpenEuler Repository
baseurl=file://${CFG[iso_mount]}
enabled=1
gpgcheck=1
gpgkey=file://${CFG[iso_mount]}/RPM-GPG-KEY-openEuler
EOF"
            fi
            vc_add --run-command "mv /etc/yum.repos.d/openEuler.repo /etc/yum.repos.d/openEuler.repo.bak"
            vc_add --run-command "yum clean all"
            vc_add --run-command "yum makecache"
            vc_add --run-command "yum install -y virtcca-guest-agent"
            vc_add --run-command "systemctl enable ${GA_SERVICE}"
            vc_add --run-command "mv /etc/yum.repos.d/openEuler.repo.bak /etc/yum.repos.d/openEuler.repo"
            ;;

        "rpm")
            vc_add --run-command "mkdir -p /etc/virtcca_deploy/cert"
            vc_add --run-command "yum -y clean all || true"
            vc_add --run-command "yum -y makecache || true"

            vc_add --run-command "yum -y install numactl pciutils python3-flask python3-gevent python3-psutil python3-requests python3-PyYAML python3-sqlalchemy python3-Flask-SQLAlchemy python3-gunicorn python3-Flask-Cors python3-dotenv"

            for f in ${CFG[rpm_files]}; do
                vc_add --copy-in "$f:/home"
            done

            vc_add --run-command "rpm -ivh /home/virtcca-deploy-common-*.rpm"
            vc_add --run-command "rpm -ivh /home/virtcca-guest-agent-*.rpm"
            vc_add --run-command "systemctl enable ${GA_SERVICE}"
            ;;
    esac
}

# ==============================================================================
# guest boot config
# ==============================================================================

configure_guest_boot() {
    vc_add --run-command 'grub2-mkimage -d /usr/lib/grub/arm64-efi -O arm64-efi --output=/boot/efi/EFI/openEuler/grubaa64.efi --prefix="(,msdos1)/efi/EFI/openEuler" fat part_gpt part_msdos linux tpm'
    vc_add --run-command 'cp -f /boot/efi/EFI/openEuler/grubaa64.efi /boot/EFI/BOOT/BOOTAA64.EFI'

    local boot_params="ima_rot=tpm cma=64M virtcca_cvm_guest=1 cvm_guest=1 swiotlb=65536,force loglevel=8"
    vc_add --run-command "grep -q 'virtcca_cvm_guest=1' /boot/efi/EFI/openEuler/grub.cfg || sed -i '/linux.*vmlinuz-6.6.0/ s/\$/ ${boot_params}/' /boot/efi/EFI/openEuler/grub.cfg"
    vc_add --run-command "grep -q 'virtcca_cvm_guest=1' /etc/default/grub || sed -i '/^GRUB_CMDLINE_LINUX=/ s/\"$/ ${boot_params}\"/' /etc/default/grub"

    vc_add --run-command 'echo "install ubfi /bin/false" > /etc/modprobe.d/disable-ubfi.conf'
}

# ==============================================================================
# inject
# ==============================================================================

run() {
    virt-customize -a "${CFG[image_path]}" "${VC_ACTIONS[@]}"
}

# ==============================================================================
# main
# ==============================================================================

main() {
    parse_args "$@"
    validate
    prepare
    ssl_prepare
    build
    ssl_add
    configure_guest_boot
    run
    log INFO "DONE"
}

main "$@"