#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0+
#
"""Script to build/run U-Boot with QEMU

It assumes that

- you build U-Boot in ${ubdir}/<name> where <name> is the U-Boot
  board config
- your OS images are in ${imagedir}/{distroname}/

So far the script supports only ARM and x86
"""

import argparse
import os
from pathlib import Path
import subprocess
import sys
import shlex
import time

import build_helper

OUR_PATH = os.path.dirname(os.path.realpath(__file__))
OUR1_PATH = os.path.dirname(OUR_PATH)

# Bring in the patman and dtoc libraries (but don't override the first path
# in PYTHONPATH)
sys.path.insert(2, os.path.join(OUR1_PATH, 'tools'))

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


def parse_args():
    """Parses command-line arguments"""
    parser = argparse.ArgumentParser(
        description='Build and/or run U-Boot with QEMU',
        formatter_class=argparse.RawTextHelpFormatter)
    build_helper.add_common_args(parser)
    parser.add_argument('-C', '--enable-console', action='store_true',
                        help="Enable linux console (x86 only)")
    parser.add_argument('-D', '--share-dir', metavar='DIR',
                        help='Directory to share into the guest via virtiofs')
    parser.add_argument('-e', '--sct-run', action='store_true',
                        help='Run UEFI Self-Certification Test (SCT)')
    parser.add_argument('-E', '--use-tianocore', action='store_true',
                        help='Run Tianocore (OVMF) instead of U-Boot')
    parser.add_argument('-I', '--initrd',
                        help='Initial ramdisk to run using -initrd')
    parser.add_argument('-K', '--kernel',
                        help='Kernel to run using -kernel')
    parser.add_argument('-Q', '--use-qboot', action='store_true',
                        help='Run qboot instead of U-Boot')
    parser.add_argument('-x', '--xpl', action='store_true',
                        help='Use xPL image rather than U-Boot proper')
    parser.add_argument(
        '-S', '--sct-seq',
        help='SCT sequence-file to be written into the SCT image if -e')
    parser.add_argument(
        '-t', '--root',
        help='Pass the given root device to linux via root=xxx')
    parser.add_argument(
        '-U', '--uuid',
        help='Pass the given root device to linux via root=/dev/disk/by-uuid/')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Show executed commands')

    return parser.parse_args()


class BuildQemu:
    """Build and/or run U-Boot with QEMU based on command line arguments"""

    def __init__(self, args):
        """Set up arguments and configure paths"""
        self.args = args

        self.helper = build_helper.Helper()
        self.helper.read_settings()
        self.imagedir = Path(self.helper.get_setting('image_dir', '~/dev'))
        self.ubdir = Path(self.helper.get_setting('build_dir', '/tmp/b'))
        self.sctdir = Path(self.helper.get_setting('sct_dir', '~/dev/efi/sct'))
        self.tiano = Path(self.helper.get_setting('tianocore_dir',
                                                  '~/dev/tiano'))
        self.qboot = Path(self.helper.get_setting('qboot_dir', '~/dev/qboot'))
        self.mnt = Path(self.helper.get_setting('sct_mnt', '/mnt/sct'))

        self.bitness = 32 if args.word_32bit else 64
        self.qemu_extra = []
        self.mem = '512M'  # Default QEMU memory

        if args.disk:
            self.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])

        if args.sct_run:
            self.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])
            # SCT usually runs headlessly
            self.qemu_extra.extend(['-display', 'none'])
            # For potential interaction within SCT
            self.qemu_extra.extend(['-device', 'qemu-xhci'])
            self.qemu_extra.extend(['-device', 'usb-kbd'])
            sct_image_path = self.sctdir / 'sct.img'
            if not sct_image_path.exists():
                tout.fatal(f'Error: SCT image {sct_image_path} not found, '
                           'required for -e')
            self.qemu_extra.extend([
                '-drive', f'file={sct_image_path},format=raw,if=none,id=vda',
                '-device', 'virtio-blk-pci,drive=vda,bootindex=1'])
            # Basic networking for SCT, if needed
            self.qemu_extra.extend([
                '-device', 'virtio-net-pci,netdev=net0',
                '-netdev', 'user,id=net0'])
            args.serial_only = True  # SCT implies serial output

        if args.os:
            self.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])

        self.kvm_params = []
        if args.kvm:
            self.kvm_params = ['-enable-kvm', '-cpu', 'host']

        bios_override = None
        if args.use_tianocore:
            bios_override = Path(self.tiano, 'OVMF-pure-efi.x64.fd')
            if not bios_override.exists():
                tout.fatal(
                    'Error: Tianocore BIOS specified (-E) but not found at '
                    f'{bios_override}')
        elif args.use_qboot:
            bios_override = Path(self.qboot, 'bios.bin')
            if not bios_override.exists():
                tout.fatal(
                    'Error: qboot BIOS specified (-Q) but not found at '
                    f'{bios_override}')

        self.seq_fname = Path(args.sct_seq) if args.sct_seq else None
        self.img_fname = Path(args.disk) if args.disk else None

        # arch-specific setup
        if args.arch == 'arm':
            if args.xpl:
                self.board = 'qemu_arm_spl'
                default_bios = 'image.bin'
            else:
                self.board = 'qemu_arm'
                default_bios = 'u-boot.bin'
            self.qemu = 'qemu-system-arm'
            self.qemu_extra.extend(['-machine', 'virt'])
            if not args.kvm:
                self.qemu_extra.extend(['-accel', 'tcg'])
            os_arch = 'arm'
            if self.bitness == 64:
                if args.xpl:
                    self.board = 'qemu_arm64_spl'
                else:
                    self.board = 'qemu_arm64'
                self.qemu = 'qemu-system-aarch64'
                self.qemu_extra.extend(['-cpu', 'cortex-a57'])
                os_arch = 'arm64'
        elif args.arch == 'x86':
            self.board = 'qemu-x86'
            default_bios = 'u-boot.rom'
            self.qemu = 'qemu-system-i386'
            os_arch = 'i386'  # For OS image naming
            if self.bitness == 64:
                self.board = 'qemu-x86_64'
                self.qemu = 'qemu-system-x86_64'
                os_arch = 'amd64'
        else:
            raise ValueError(f"Invalid arch '{args.arch}'")

        self.os_path = None
        if args.os == 'ubuntu':
            img_name = (f'{args.os}-{args.release}-desktop-{os_arch}.iso')
            self.os_path = self.imagedir / args.os / img_name

        self.build_dir = self.ubdir / self.board
        self.bios = (bios_override if bios_override
                     else self.build_dir / default_bios)

    @staticmethod
    def execute_command(cmd_list, desc, check=True, **kwargs):
        """Execute a shell command and handle errors

        Args:
            cmd_list (list of str): The command and its arguments as a list
            desc (str): A description of the command being executed
            check (bool): Raise CalledProcessError on non-zero exit code
            kwargs: Additional arguments for subprocess.run

        Return:
            subprocess.CompletedProcess: The result of the subprocess.run call

        Raises:
            SystemExit: If the command is not found or fails and check is True
        """
        tout.info(f"Executing: {desc} -> {shlex.join(cmd_list)}")
        try:
            # Decode stdout/stderr by default if text=True
            if 'text' not in kwargs:
                kwargs['text'] = True
            return subprocess.run(cmd_list, check=check, **kwargs)
        except FileNotFoundError:
            tout.fatal(f"Error: Command '{cmd_list[0]}' not found")
        except subprocess.CalledProcessError as proc:
            tout.error(f'Error {desc}: Command failed with exit code '
                       f'{proc.returncode}')
            if proc.stdout:
                tout.error(f'Stdout:\n{proc.stdout}')
            if proc.stderr:
                tout.error(f'Stderr:\n{proc.stderr}')
            tout.fatal('Failed')

    def build_u_boot(self):
        """Build U-Boot using buildman
        """
        self.build_dir.mkdir(parents=True, exist_ok=True)
        cmd = ['buildman', '-w', '-o', str(self.build_dir), '--board',
               self.board, '-I']

        self.execute_command(
            cmd,
            f'Building U-Boot for {self.board} in {self.build_dir}')

    def update_sct_sequence(self):
        """Update the SCT image with a specified sequence file

        Requires sudo for loop device setup and mounting
        """
        if not (self.args.sct_run and self.seq_fname and
                self.seq_fname.exists()):
            if (self.args.sct_run and self.seq_fname and
                    not self.seq_fname.exists()):
                tout.warning(f'Warning: SCT sequence file {self.seq_fname}'
                             'not found')
            return

        fname = self.sctdir / 'sct.img'
        if not fname.exists():
            tout.fatal(f'Error: SCT image {fname} not found')

        loopdev = None
        try:
            # Find free loop device and attach
            loopdev = command.output_one_line(
                'sudo', 'losetup', '--show', '-f', '-P', str(fname))
            partition_path_str = f'{loopdev}p1'

            uid, gid = os.getuid(), os.getgid()
            mount_cmd = ['sudo', 'mount', partition_path_str,
                         str(self.mnt), '-o', f'uid={uid},gid={gid},rw']
            mount_cmd.extend(['-t', 'vfat'])

            self.execute_command(mount_cmd,
                                 f'Mounting {partition_path_str} to {self.mnt}')

            target_sct_path = self.mnt / self.seq_fname.name
            self.execute_command(
                ['sudo', 'cp', str(self.seq_fname), str(target_sct_path)],
                f'Copying {self.seq_fname.name} to {self.mnt}'
            )
            tout.info(f"Copied {self.seq_fname} to {target_sct_path}")

        finally:
            if Path(self.mnt).is_mount():
                self.execute_command(['sudo', 'umount', str(self.mnt)],
                                     f'Unmounting {self.mnt}', check=False)
            if loopdev:
                self.execute_command(['sudo', 'losetup', '-d', loopdev],
                                     f'Detaching loop device {loopdev}',
                                     check=False)

    def run_qemu(self):
        """Construct and run the QEMU command"""
        if not self.bios.exists():
            tout.fatal(f"Error: BIOS file '{self.bios}' not found")

        cmdline = []

        qemu_cmd = [str(self.qemu)]
        if self.bios:
            qemu_cmd.extend(['-bios', str(self.bios)])
        qemu_cmd.extend(self.kvm_params)
        qemu_cmd.extend(['-m', self.mem])

        if not self.args.sct_run and not self.qboot:
            qemu_cmd.extend(['-netdev', 'user,id=net0,hostfwd=tcp::2222-:22',
                             '-device', 'virtio-net-pci,netdev=net0'])

        # Display and Serial
        # If -e (sct_run) is used, "-display none" is in qemu_extra
        # If -s (serial_only) is used, we want no display
        has_display_option = any(
            item.startswith('-display') for item in self.qemu_extra)
        if self.args.serial_only and not has_display_option:
            qemu_cmd.extend(['-display', 'none'])
        if not any(item.startswith('-serial') for item in self.qemu_extra):
            qemu_cmd.extend(['-serial', 'mon:stdio'])

        if self.args.kernel:
            qemu_cmd.extend(['-kernel', self.args.kernel])
        if self.args.initrd:
            qemu_cmd.extend(['-initrd', self.args.initrd])

        if self.args.enable_console:
            cmdline.append('console=ttyS0,115200,8n1')
        if self.args.root:
            cmdline.append(f'root={self.args.root}')
        if self.args.uuid:
            cmdline.append(f'root=/dev/disk/by-uuid/{self.args.uuid}')

        if cmdline:
            qemu_cmd.extend(['-append'] + [' '.join(cmdline)])

        # Add other parameters gathered from options
        qemu_cmd.extend(self.qemu_extra)
        if self.os_path:
            if not self.os_path.exists():
                tout.error(f'OS image {self.os_path} specified but not found')
            qemu_cmd.extend([
                '-drive',
                f'if=virtio,file={self.os_path},format=raw,id=hd0,readonly=on'])

        if self.img_fname:
            if self.img_fname.exists():
                qemu_cmd.extend([
                    '-drive',
                    f'if=virtio,file={self.img_fname},format=raw,id=hd1'])
            else:
                tout.warning(f"Disk image '{self.img_fname}' not found")

        sock = Path('/tmp/virtiofs.sock')
        proc = None
        if self.args.share_dir:
            virtfs_dir = Path(self.args.share_dir)
            if not virtfs_dir.is_dir():
                tout.fatal(f'Error: VirtFS share directory {virtfs_dir} '
                           f'is not a valid directory')

            virtiofsd = Path('/usr/libexec/virtiofsd')
            if not virtiofsd.exists():
                tout.fatal(f'Error: virtiofsd not found at {virtiofsd}')

            # Clean up potential old socket file
            if sock.exists():
                try:
                    sock.unlink()
                    tout.info(f'Removed old socket file {sock}')
                except OSError as e:
                    tout.warning(
                        f'Warning: Could not remove old socket file {sock}: '
                        f'{e}')

            qemu_cmd.extend([
                '-chardev', f'socket,id=char0,path={sock}',
                '-device',
                'vhost-user-fs-pci,queue-size=1024,chardev=char0,tag=hostshare',
                '-object',
                f'memory-backend-file,id=mem,size={self.mem},mem-path=/dev/shm'
                    ',share=on',
                '-numa', 'node,memdev=mem'])

            virtiofsd_cmd = [
                str(virtiofsd),
                '--socket-path', str(sock),
                '--shared-dir', str(virtfs_dir),
                '--cache', 'auto']
            try:
                # Use Popen to run virtiofsd in the background
                proc = subprocess.Popen(virtiofsd_cmd, stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
                # Give virtiofsd a moment to start and create the socket
                time.sleep(0.5)
                if not sock.exists() and proc.poll() is not None:
                    stdout, stderr = proc.communicate()
                    tout.error('Error starting virtiofsd. Exit code: '
                               f'{proc.returncode}')
                    if stdout:
                        tout.error(f"virtiofsd stdout:\n{stdout.decode()}")
                    if stderr:
                        tout.error(f"virtiofsd stderr:\n{stderr.decode()}")
                    tout.fatal('Failed')

            except (subprocess.CalledProcessError, FileNotFoundError) as exc:
                tout.fatal(f'Failed to start virtiofsd: {exc}')

        tout.info(f'QEMU:\n{shlex.join(qemu_cmd)}\n')
        try:
            subprocess.run(qemu_cmd, check=True)
        except FileNotFoundError:
            tout.fatal(f"Error: QEMU executable '{self.qemu}' not found")
        except subprocess.CalledProcessError as e:
            tout.fatal(f'QEMU execution failed with exit code {e.returncode}')
        finally:
            # Clean up virtiofsd process and socket if it was started
            if proc:
                tout.info('Terminating virtiofsd')
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    tout.warning(
                        'virtiofsd did not terminate gracefully; killing')
                    proc.kill()
                if sock.exists():
                    try:
                        sock.unlink()
                    except OSError as e_os:
                        tout.warning('Warning: Could not remove virtiofs '
                                     f'socket {sock}: {e_os}')

    def start(self):
        """Build and run QEMU"""
        if not self.args.no_build and not self.args.use_tianocore:
            self.build_u_boot()

        # Update SCT sequence if -e and -S are given
        if self.args.sct_run and self.seq_fname:
            self.update_sct_sequence()

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


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

    qemu = BuildQemu(args)
    qemu.start()

if __name__ == '__main__':
    main()