#!/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, DistroList

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):
    l = DistroList(imgdir)
    image_id = l.mkid(f'{name}-{release}', 'cloud')
    image = l.find_filename(image_id)
    if image is None:
        print(f"ERROR: no {name}-{release} cloud image found")
        print("HINT: try virt-cloud-fetch")
        return None
    if os.path.exists(image):
        return image
    print(f"ERROR: file {image} does not exist")
    return None

# pylint: disable=too-many-statements
# pylint: disable=too-many-return-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('--fedora', dest = 'fedora', type = str,
                        help = 'boot fedora VER', metavar = 'VER')
    pgroup.add_argument('--fedora-uki', dest = 'fedora_uki', type = str,
                        help = 'boot fedora VER (uki variant)', metavar = 'VER')
    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 fedora/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.add_argument('--copy-file', dest = 'copyfiles',
                        action = 'append', type = str, default = [],
                        help = 'copy FILE into the guest', metavar = 'FILE')
    pgroup.add_argument('--runcmd', dest = 'runcmd',
                        action = 'append', type = str, default = [],
                        help = 'run CMD in guest', metavar = 'CMD')

    pgroup = parser.add_argument_group('firmware setup')
    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('--cc-mode', dest = 'cc_mode',
                        choices=['none', 'sev', 'snp', 'tdx'], default = 'none',
                        help = 'set confidential computing mode (default: none)')
    pgroup.add_argument('--sev', dest = 'cc_mode',
                        action = 'store_const', const = 'sev',
                        help = 'enable sev cc mode (feature "amd-sev")')
    pgroup.add_argument('--snp', dest = 'cc_mode',
                        action = 'store_const', const = 'snp',
                        help = 'enable snp cc mode (feature "amd-sev-snp")')
    pgroup.add_argument('--tdx', dest = 'cc_mode',
                        action = 'store_const', const = 'tdx',
                        help = 'enable tdx cc mode (feature "intel-tdx")')
    pgroup.add_argument('--log', dest = 'log',
                        help = 'write firmware log to FILE', metavar = 'FILE')

    pgroup = parser.add_argument_group('hardware setup')
    pgroup.add_argument('--memory', '--mem', dest = 'memory', default = '4G',
                        help = 'assign guest memory (default: 4G)', metavar = 'MEM')
    pgroup.add_argument('--smp', dest = 'smp', default = '1',
                        help = 'assign guest vcpus (default: 1)', metavar = 'CPUS')
    pgroup.add_argument('--scsi', dest = 'use_scsi',
                        action = argparse.BooleanOptionalAction, default = True,
                        help = 'use virtio scsi (default: yes)')
    pgroup.add_argument('--sata', dest = 'use_sata',
                        action = argparse.BooleanOptionalAction, default = False,
                        help = 'use sata (q35 only, default: no)')
    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', options.memory ]
    cmdline += [ '-smp', options.smp ]
    cmdline += [ '-cpu', 'host' ]
    cmdline += [ '-enable-kvm', ]
    cmdline += [ '-nographic', ]

    # qemucfg
    use_sata = False
    if machine.startswith('q35') and options.use_sata:
        use_sata = True
    qemucfg = QemuCfg(scsi = options.use_scsi,
                      nvme = options.use_nvme,
                      sata = use_sata)

    # firmware
    fw = find_firmware(sb = options.sb,
                       smm = options.smm,
                       qemuvars = options.qemuvars,
                       sev = options.cc_mode == 'sev',
                       snp = options.cc_mode == 'snp',
                       tdx = options.cc_mode == '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()
    qemucfg.add_confidential(options.cc_mode)

    # hardware
    qemucfg.add_virtio_rng()
    qemucfg.add_network()
    if options.log and os.uname().machine == 'x86_64':
        qemucfg.add_x64_firmware_log(options.log)

    # 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
    elif options.fedora:
        hostname = f'fedora-{options.fedora}'
        title = f'Fedora {options.fedora}'
        image = find_cloud_image(options.imgdir, 'fedora', options.fedora)
        if image is None:
            return 1
    elif options.fedora_uki:
        hostname = f'fedora-uki-{options.fedora_uki}'
        title = f'Fedora {options.fedora} (UKI)'
        image = find_cloud_image(options.imgdir, 'fedora-uki', options.fedora_uki)
        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.replace('.','-'))
    cloud.add_user(options.userid, passwd = options.passwd, sudo = True)
    for pkg in options.package:
        cloud.add_package(pkg)
    for f in options.copyfiles:
        cloud.add_file_copy(f)
    for cmd in options.runcmd:
        cloud.add_runcmd(cmd)
    iso = cloud.create_iso()
    qemucfg.add_cdrom_iso(iso.name)

    # go!
    cmdline += qemucfg.cmdline()
    if options.dryrun:
        print("=== cloud image ===")
        print(image)
        print("=== cloud init ===")
        cloud.dump_user()
        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())
