#!/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 LiteralStr, 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_distro(imgdir, name, release):
    l = DistroList(imgdir)
    image_id = l.mkid(f'{name}-{release}', 'cloud')
    distro = l.find_entry(image_id)
    if distro is None:
        print(f"ERROR: no {name}-{release} cloud image found")
        print("HINT: use virt-cloud-fetch to download images")
        print("HINT: use 'virt-cloud-boot --list' to show images")
        return None
    return distro

def list_cloud_images(imgdir):
    l = DistroList(imgdir)
    print(f'{"Name":16} Version')
    print('---------------- -------')
    for item in filter(lambda e: e.kind == 'cloud', l.list_entries()):
        print(f'{item.name:16} {item.version}')

# 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('--ubuntu', dest = 'ubuntu', type = str,
                        help = 'boot ubuntu 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.add_argument('--list', dest = 'imglist',
                        action = 'store_true', default = False,
                        help = f'list cloud images in {defdir}')

    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('--no-uefi', dest = 'uefi',
                        action = 'store_false', default = True,
                        help = 'skip uefi firmware setup (boot with qemu default firmware)')
    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)')
    pgroup.add_argument('--vblk', dest = 'use_vblk',
                        action = argparse.BooleanOptionalAction, default = False,
                        help = 'use virtio block (default: no)')

    options = parser.parse_args()

    if options.imglist:
        list_cloud_images(options.imgdir)
        return 0

    # 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,
                      vblk = options.use_vblk)

    # firmware
    varsfile = None
    if options.uefi:
        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:
            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:
        # filename
        hostname = 'cloud-image'
        distro = None
        title = options.image
        image = options.image
    else:
        # distro
        if options.rhel:
            title = f'RHEL {options.rhel}'
            distro = find_cloud_distro(options.imgdir, 'rhel', options.rhel)
        elif options.centos:
            title = f'CentOS Stream {options.centos}'
            distro = find_cloud_distro(options.imgdir, 'centos-stream', options.centos)
        elif options.fedora:
            title = f'Fedora {options.fedora}'
            distro = find_cloud_distro(options.imgdir, 'fedora', options.fedora)
        elif options.fedora_uki:
            title = f'Fedora {options.fedora_uki} (UKI)'
            distro = find_cloud_distro(options.imgdir, 'fedora-uki', options.fedora_uki)
        elif options.ubuntu:
            title = f'ubuntu {options.ubuntu}'
            distro = find_cloud_distro(options.imgdir, 'ubuntu', options.ubuntu)
        else:
            print("ERROR: no cloud image specified")
            return 1
        if distro is None:
            # find_cloud_distro prints error message
            return 1
        image = distro.filepath
        if not os.path.exists(image):
            print(f"ERROR: file {image} does not exist")
            return 1
        hostname = f'{distro.name}-{distro.version}'.replace('.','-')
    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)
    for f in options.copyfiles:
        cloud.add_file_copy(f)
    for cmd in options.runcmd:
        cloud.add_runcmd(cmd)
    if distro is not None:
        repoconfig = distro.repoconfig()
        if repoconfig:
            cloud.add_file_content('/etc/yum.repos.d/cloudinit.repo',
                                   LiteralStr(repoconfig))
    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())
