#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0+
"""
Script to build an EFI thing suitable for booting with QEMU, possibly running
it also.

UEFI binaries for QEMU used for testing this script:

OVMF-pure-efi.i386.fd at
https://drive.google.com/file/d/1jWzOAZfQqMmS2_dAK2G518GhIgj9r2RY/view?usp=sharing

OVMF-pure-efi.x64.fd at
https://drive.google.com/file/d/1c39YI9QtpByGQ4V0UNNQtGqttEzS-eFV/view?usp=sharing

Use ~/.build-efi to configure the various paths used by this script.

When --bootcmd is specified, a uboot.env file is created on the EFI partition
containing the boot command. U-Boot needs to be configured to import this file
on startup, for example by adding to CONFIG_PREBOOT or the default bootcmd:

  load ${devtype} ${devnum}:${distro_bootpart} ${loadaddr} uboot.env; \
  env import -t ${loadaddr} ${filesize}
"""

from argparse import ArgumentParser
import os
from pathlib import Path
import shutil
import sys
import tempfile

import build_helper

# pylint: disable=C0413
from u_boot_pylib import command
from u_boot_pylib import tools
from u_boot_pylib import tout


def parse_args():
    """Parse the program arguments

    Return:
        Namespace object
    """
    parser = ArgumentParser(
        epilog='Script for running U-Boot as an EFI app/payload')
    build_helper.add_common_args(parser)
    parser.add_argument('-g', '--debug', action='store_true',
                        help="Run QEMU with gdb")
    parser.add_argument('--write-kernel', action='store_true',
                        help='Add a kernel to the disk image')
    parser.add_argument('-O', '--old', action='store_true',
                        help='Use old EFI app build (before 32/64 split)')
    parser.add_argument('-p', '--payload', action='store_true',
                        help='Package up the payload instead of the app')
    parser.add_argument('-P', '--partition', action='store_true',
                        help='Create a partition table')
    parser.add_argument('--spice', action='store_true',
                        help='Enable SPICE for clipboard sharing')
    parser.add_argument('-N', '--net', action='store_true',
                        help='Enable networking (with SSH forwarding on port 2222)')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Show executed commands')
    parser.add_argument('--include-dir',
                        help='Directory containing additional files to include in the image')

    args = parser.parse_args()

    return args


class BuildEfi:
    """Class to collect together the various bits of state while running"""
    def __init__(self, args):
        self.helper = build_helper.Helper(args)
        self.helper.read_settings()
        self.img_fname = self.helper.get_setting('efi_image_file', 'efi.img')
        self.img = None
        self.build_topdir = self.helper.get_setting("build_dir", '/tmp')
        self.build_dir = None
        self.args = args
        self.imagedir = Path(self.helper.get_setting('image_dir', '~/dev'))

    def run_qemu(self, serial_only):
        """Run QEMU

        Args:
            serial_only (bool): True to run without a display
        """
        extra = []
        efi_dir = self.helper.get_setting('efi_dir')
        if self.args.arch == 'arm':
            qemu_arch = 'aarch64'
            extra += ['--machine', 'virt', '-cpu', 'max']
            bios = os.path.join(efi_dir, 'OVMF-pure-efi.aarch64.fd.64m')
            var_store = os.path.join(efi_dir, 'varstore.img')
            extra += [
                '-drive', f'if=pflash,format=raw,file={bios},readonly=on',
                '-drive', f'if=pflash,format=raw,file={var_store}'
                ]
            extra += ['-drive',
                      f'if=virtio,file={self.img},format=raw,id=hd0']
        else:  # x86
            if self.helper.bitness == 64:
                qemu_arch = 'x86_64'
                bios = 'OVMF-release-x64.fd'
            else:
                qemu_arch = 'i386'
                bios = 'OVMF-pure-efi.i386.fd'
            bios = os.path.join(efi_dir, bios)
            var_store = os.path.join(efi_dir, 'OVMF_VARS_4M.fd')
            extra += [
                '-drive', f'if=pflash,format=raw,file={bios},readonly=on',
                '-drive', f'if=pflash,format=raw,file={var_store}'
                ]
            extra += ['-drive', f'id=disk,file={self.img},if=none,format=raw']
            extra += ['-device', 'ahci,id=ahci']
            extra += ['-device', 'ide-hd,drive=disk,bus=ahci.0']
        qemu = f'qemu-system-{qemu_arch}'
        if serial_only:
            extra += ['-display', 'none', '-serial', 'mon:stdio']
            serial_msg = ' (Ctrl-a x to quit)'
        else:
            if self.args.arch == 'arm':
                extra += ['-device', 'virtio-gpu-pci']
                extra += ['-device', 'qemu-xhci', '-device', 'usb-kbd',
                          '-device', 'usb-tablet']
                extra += ['-display', 'default,show-cursor=on']
            else:  # x86
                extra += ['-device', 'qemu-xhci', '-device', 'usb-kbd',
                          '-device', 'usb-mouse']

            # This uses QEMU's GTK clipboard integration with SPICE vdagent
            if self.args.spice:
                extra += ['-device', 'virtio-serial-pci']
                extra += ['-chardev', 'qemu-vdagent,id=spicechannel0,name=vdagent,clipboard=on']
                extra += ['-device', 'virtserialport,chardev=spicechannel0,name=com.redhat.spice.0']
            extra += ['-serial', 'mon:stdio']
            serial_msg = ''
        if self.args.kvm:
            extra.extend(['-enable-kvm', '-cpu', 'host'])

        print(f'Running {qemu}{serial_msg}')

        # Use 512MB since U-Boot EFI likes to have 256MB to play with
        if self.args.os or self.args.disk:
            mem = '4G'
            extra.extend(['-smp', '4'])
        else:
            mem = '1G'

        if self.args.debug:
            extra.extend(['-s', '-S'])

        cmd = [qemu]
        cmd += '-m', mem
        if self.args.net:
            cmd += '-netdev', 'user,id=net0,hostfwd=tcp::2222-:22'
            cmd += '-device', 'virtio-net-pci,netdev=net0'
        else:
            cmd += '-nic', 'none'
        cmd += extra
        self.helper.add_qemu_args(self.args, cmd, base_hd=1)
        tout.info(' '.join(cmd))
        sys.stdout.flush()
        command.run(*cmd)

    def setup_files(self, build, build_type, dst):
        """Set up files in the staging area

        Args:
            build (str): Name of build being packaged, e.g. 'efi-x86_app32'
            build_type (str): Build type ('app' or 'payload')
            dst (str): Destination directory
        """
        print(f'Packaging {build}')
        if self.args.custom:
            dirname, fname = os.path.split(self.args.custom)
            if not dirname:
                dirname = '.'
        else:
            fname = f'u-boot-{build_type}.efi'
            dirname = f'{self.build_dir}/'
        tools.write_file(f'{dst}/startup.nsh', f'fs0:{fname}', binary=False)
        shutil.copy(f'{dirname}/{fname}', dst)

        # Copy additional files from include directory if specified
        if self.args.include_dir:
            include_path = Path(self.args.include_dir)
            if include_path.exists() and include_path.is_dir():
                print(f'Including files from {include_path}')
                for item in include_path.iterdir():
                    if item.is_file():
                        print(f'  Copying {item.name}')
                        shutil.copy(str(item), dst)
                    elif item.is_dir():
                        dest_dir = Path(dst) / item.name
                        print(f'  Copying directory {item.name}')
                        shutil.copytree(str(item), str(dest_dir))
            else:
                print(f'Warning: Include directory {include_path} does not exist or is not a directory')

        # Write U-Boot environment file if bootcmd is specified
        if self.args.bootcmd:
            # Check if mkenvimage is available (local build or system-wide)
            mkenvimage = 'tools/mkenvimage'
            if not os.path.exists(mkenvimage):
                mkenvimage = 'mkenvimage'
                if not shutil.which(mkenvimage):
                    tout.error('Please install u-boot-tools package:')
                    tout.error('  sudo apt install u-boot-tools')
                    raise FileNotFoundError('mkenvimage not found')

            # Create text environment file
            env_content = f'bootcmd={self.args.bootcmd}\n'
            with tempfile.NamedTemporaryFile(mode='w', delete=False,
                                             suffix='.txt') as outf:
                outf.write(env_content)
                env_fname = outf.name

            try:
                # Convert to binary format with CRC using mkenvimage
                command.run(mkenvimage, '-s', '0x1000',
                           '-o', f'{dst}/uboot.env', env_fname)
                print(f'Created uboot.env with bootcmd: {self.args.bootcmd}')
            finally:
                os.unlink(env_fname)

    def do_build(self, build):
        """Build U-Boot for the selected board"""
        extra = ['-a', '~CONSOLE_PAGER'] if self.args.no_pager else []
        res = command.run_one('buildman', '-w', '-o', self.build_dir, *extra,
                              '--board', build, '-I', raise_on_error=False)
        if res.return_code and res.return_code != 101:  # Allow warnings
            raise ValueError(
                f'buildman exited with {res.return_code}: {res.combined}')

    def start(self):
        """This does all the work"""
        args = self.args
        arch = 'arm' if self.args.arch == 'arm' else 'x86'
        build_type = 'payload' if args.payload else 'app'
        build = f'efi-{arch}_{build_type}{self.helper.bitness}'

        if args.build_dir:
            self.build_dir = args.build_dir
            self.img = f'{self.build_dir}/{self.img_fname}'
        else:
            self.build_dir = f'{self.build_topdir}/{build}'
            self.img = self.img_fname
        if not args.no_build:
            self.do_build(build)

        if args.old and self.helper.bitness == 32:
            build = f'efi-{arch}_{build_type}'

        with self.helper.make_disk(self.img, fs_type='vfat',
                                   use_part=args.partition) as dirpath:
            self.setup_files(build, build_type, dirpath)
            if self.args.write_kernel:
                bzimage = self.helper.get_setting('bzimage_file', 'bzImage')
                command.run('cp', bzimage, f'{dirpath}/vmlinuz')

        if args.run:
            self.run_qemu(args.serial_only)


def main():
    """Parse arguments and start the program"""
    args = parse_args()
    tout.init(tout.INFO if args.verbose else tout.WARNING)

    qemu = BuildEfi(args)
    qemu.start()

if __name__ == '__main__':
    main()