#!/usr/bin/python3
#
# SPDX-License-Identifier: GPL-2.0-only
# (c) 2026 Gerd Hoffmann
#
# pylint: disable=consider-using-with
""" boot cloud images on qemu """
import os
import sys
import shutil
import argparse
import tempfile
import subprocess

from virt.qemutest.qemucfg import QemuCfg
from virt.qemutest.autoconf import find_qemu_binary, find_firmware, qemu_default_machine
from virt.qemutest.cloudinit import CloudInitISO
from virt.qemutest.distro.core import default_imgdir, FetchDistroCore

def xterm_title(msg):
    term = os.environ.get('TERM', 'dumb')
    if term.startswith('xterm'):
        print(f'\x1b]2;{msg}\x07', flush = True, end = '')

def find_cloud_image(imgdir, name, release):
    f = FetchDistroCore()
    index = f.index_load(imgdir)
    entry = index.get(f'{name}-{release}-{os.uname().machine}-cloud')
    if entry is None:
        print(f"ERROR: no {name}-{release} cloud image found")
        print("HINT: try virt-cloud-fetch")
        return None
    filename = entry.get('filename')
    image = f'{imgdir}/{filename}'
    if os.path.exists(image):
        return image
    print(f"ERROR: file {image} does not exist")
    return None

def enable_cc(sev = False, snp = False, tdx = False):
    cmdline = []
    if sev:
        cmdline += [ '-object', 'sev-guest,id=sev0,reduced-phys-bits=1,cbitpos=51' ]
        cmdline += [ '-machine', 'confidential-guest-support=sev0' ]
    elif snp:
        cmdline += [ '-object', 'sev-snp-guest,id=snp0,reduced-phys-bits=1,cbitpos=51' ]
        cmdline += [ '-machine', 'confidential-guest-support=snp0' ]
    elif tdx:
        cmdline += [ '-object', 'tdx-guest,id=tdx0' ]
        cmdline += [ '-machine', 'confidential-guest-support=tdx0' ]
    return cmdline

# pylint: disable=too-many-statements
# pylint: disable=too-many-branches
def main():
    parser = argparse.ArgumentParser(description = "boot cloud image")
    parser.add_argument('--dry-run', '--dryrun', dest = 'dryrun',
                        action = 'store_true', default = False,
                        help = 'print generated qemu command line and quit')

    defdir = default_imgdir()
    pgroup = parser.add_argument_group('select cloud image')
    pgroup.add_argument('-i', '--image', dest = 'image', type = str,
                        help = 'boot cloud image FILE', metavar = 'FILE')
    pgroup.add_argument('--centos', dest = 'centos', type = str,
                        help = 'boot centos stream VER', metavar = 'VER')
    pgroup.add_argument('--rhel', dest = 'rhel', type = str,
                        help = 'boot rhel VER', metavar = 'VER')
    pgroup.add_argument('--imgdir', dest = 'imgdir',
                        type = str, default = defdir,
                        help = f'lookup centos/rhel images in DIR (default: {defdir})',
                        metavar = 'DIR')

    pgroup = parser.add_argument_group('setup cloud image')
    pgroup.add_argument('-u', '--username', '--user', dest = 'userid', default = 'cloud',
                        help = 'login username (default: cloud)', metavar = 'USER')
    pgroup.add_argument('--password', '--passwd', dest = 'passwd', default = 'cloud',
                        help = 'login password (default: cloud)', metavar = 'PASSWD')
    pgroup.add_argument('--package', '--rpm', dest = 'package',
                        action = 'append', type = str, default = [],
                        help = 'install package PKG', metavar = 'PKG')

    pgroup = parser.add_argument_group('firmware selection')
    pgroup.add_argument('--smm', dest = 'smm',
                        action = 'store_true', default = False,
                        help = 'use smm firmware builds (feature "require-smm")')
    pgroup.add_argument('--qemu-vars', '--qemuvars', dest = 'qemuvars',
                        action = 'store_true', default = False,
                        help = 'use qemu uefi vars (feature "host-uefi-vars")')
    pgroup.add_argument('--secure-boot', '--sb', dest = 'sb',
                        action = 'store_true', default = False,
                        help = 'enable secure boot (feature "enrolled-keys")')
    pgroup.add_argument('--sev', dest = 'sev',
                        action = 'store_true', default = False,
                        help = 'enable sev cc mode (feature "amd-sev")')
    pgroup.add_argument('--snp', dest = 'snp',
                        action = 'store_true', default = False,
                        help = 'enable snp cc mode (feature "amd-sev-snp")')
    pgroup.add_argument('--tdx', dest = 'tdx',
                        action = 'store_true', default = False,
                        help = 'enable tdx cc mode (feature "intel-tdx")')

    pgroup = parser.add_argument_group('hardware setup')
    pgroup.add_argument('--scsi', dest = 'use_scsi',
                        action = argparse.BooleanOptionalAction, default = True,
                        help = 'use virtio scsi (default: yes)')
    pgroup.add_argument('--nvme', dest = 'use_nvme',
                        action = argparse.BooleanOptionalAction, default = False,
                        help = 'use NVMe emulation (default: no)')

    options = parser.parse_args()

    # basics
    qemu = find_qemu_binary()
    if qemu is None:
        print("ERROR: no qemu binary found")
        return 1
    machine = qemu_default_machine()
    cmdline = [ qemu, ]
    cmdline += [ '-machine', machine ]
    cmdline += [ '-m', '4G' ]
    cmdline += [ '-cpu', 'host' ]
    cmdline += [ '-enable-kvm', ]
    cmdline += [ '-nographic', ]

    # firmware
    fw = find_firmware(sb = options.sb,
                       smm = options.smm,
                       qemuvars = options.qemuvars,
                       sev = options.sev,
                       snp = options.snp,
                       tdx = options.tdx)

    if fw is None:
        print("ERROR: no firmware found")
        return 1
    template = fw.vars_template()
    if template:
        varsfile = tempfile.NamedTemporaryFile()
        shutil.copyfile(template, varsfile.name)
        cmdline += fw.cmdline(varsfile = varsfile.name)
    else:
        varsfile = None
        cmdline += fw.cmdline()
    cmdline += enable_cc(sev = options.sev,
                         snp = options.snp,
                         tdx = options.tdx)

    # hardware
    have_sata = False
    if machine.startswith('q35'):
        have_sata = True
    qemucfg = QemuCfg(scsi = options.use_scsi,
                      nvme = options.use_nvme,
                      sata = have_sata)
    qemucfg.add_virtio_rng()
    qemucfg.add_network()

    # cloud image
    if options.image:
        hostname = 'cloud-image'
        title = options.image
        image = options.image
    elif options.rhel:
        hostname = f'rhel-{options.rhel}'
        title = f'RHEL {options.rhel}'
        image = find_cloud_image(options.imgdir, 'rhel', options.rhel)
        if image is None:
            return 1
    elif options.centos:
        hostname = f'centos-{options.centos}'
        title = f'CentOS Stream {options.centos}'
        image = find_cloud_image(options.imgdir, 'centos-stream', options.centos)
        if image is None:
            return 1
    else:
        print("ERROR: no cloud image specified")
        return 1
    snapshot = qemucfg.create_tempfile_qcow2(image)
    qemucfg.add_disk_qcow2(snapshot.name, bootindex = 1)

    # cloud init iso
    cloud = CloudInitISO(hostname)
    cloud.add_user(options.userid, passwd = options.passwd, sudo = True)
    for pkg in options.package:
        cloud.add_package(pkg)
    iso = cloud.create_iso()
    qemucfg.add_cdrom_iso(iso.name)

    # go!
    cmdline += qemucfg.cmdline()
    if options.dryrun:
        print("=== cloud image ===")
        print(image)
        print("=== firmware ===")
        print(fw.data)
        print("=== qemu cmdline ===")
        print(cmdline)
    else:
        xterm_title(title)
        subprocess.run(cmdline, check = True)

    # cleanup
    if varsfile:
        varsfile.close()
    snapshot.close()
    iso.close()
    return 0

if __name__ == '__main__':
    sys.exit(main())
