#!/usr/bin/env python3

# Copyright (C) 2017-2022, Xilinx, Inc.  All rights reserved.
# Copyright (C) 2022-2025, Advanced Micro Devices, Inc.  All rights reserved.
#
#
# SPDX-License-Identifier: MIT

# AMD QEMU wrapper to launch both PMU and APU instances (multiarch)

import os
import subprocess
import sys
import tempfile
import re
import shutil

binpath = os.path.dirname(os.path.abspath(__file__))

# Separate PMU and APU arguments
APU_args = sys.argv[1:]
PMU_args = []
PLM_args = []
ASU_args = []
bootbin_arg = None
mach_path_arg = None
dtb_arg = None

if '-pmu-args' in APU_args:
    pmu_args_idx = APU_args.index('-pmu-args')
    PMU_args = APU_args[pmu_args_idx+1].split()
    del APU_args[pmu_args_idx:pmu_args_idx+2]

if '-plm-args' in APU_args:
    plm_args_idx = APU_args.index('-plm-args')
    PLM_args = APU_args[plm_args_idx+1].split()
    del APU_args[plm_args_idx:plm_args_idx+2]

if '-asu-args' in APU_args:
    asu_args_idx = APU_args.index('-asu-args')
    ASU_args = APU_args[asu_args_idx+1].split()
    del APU_args[asu_args_idx:asu_args_idx+2]

if '-bootbin' in APU_args:
    bootbin_args_idx = APU_args.index('-bootbin')
    bootbin_arg = APU_args[bootbin_args_idx+1]
    del APU_args[bootbin_args_idx:bootbin_args_idx+2]

if '-machine-path' in APU_args:
    mach_path_args_idx = APU_args.index('-machine-path')
    mach_path_arg = APU_args[mach_path_args_idx+1]
    del APU_args[mach_path_args_idx:mach_path_args_idx+2]

if '-dtb' in APU_args:
    dtb_args_idx = APU_args.index('-dtb')
    dtb_arg = APU_args[dtb_args_idx+1]
    del APU_args[dtb_args_idx:dtb_args_idx+2]

# Filter this out
if '-kernel' in APU_args:
    kernel_args_idx = APU_args.index('-kernel')
    del APU_args[kernel_args_idx:kernel_args_idx+2]

help_options = ['-h', '-help', '--help']
def help(status):
    print("AMD FPGA QEMU multiarch wrapper\nVersion 2025.1\n\nUsage:")
    print(f"  {sys.argv[0]} <APU options> [-pmu-args <pmu options>] [-machine-path <path>]")
    print(f"  {sys.argv[0]} <APU options> [-plm-args <plm options>] [-machine-path <path>]\n")
    print(f"  {sys.argv[0]} <APU options> [-plm-args <plm options>] [-asu-args <asu options>] [-machine-path <path>]\n")
    if status == 0:
        print(f"\n")
        if set(PMU_args).intersection(set(help_options)):
            print(f"PMU Options:\n")
            pmu_args_s = ' '.join(PMU_args)
            help_cmd = f'{binpath}/qemu-system-microblazeel {pmu_args_s}'
        elif set(PLM_args).intersection(set(help_options)):
            print(f"PLM Options:\n")
            plm_args_s = ' '.join(PLM_args)
            help_cmd = f'{binpath}/qemu-system-microblazeel {plm_args_s}'
        elif set(ASU_args).intersection(set(help_options)):
            print(f"ASU Options:\n")
            asu_args_s = ' '.join(ASU_args)
            help_cmd = f'{binpath}/qemu-system-riscv32 {asu_args_s}'
        else:
            if not set(APU_args).intersection(set(help_options)):
                APU_args.append('-help')
            print(f"APU Options:\n")
            print(f"    -bootbin [arch:]<boot.bin> - Use a boot.bin instead of individual firmware, device trees and bootloader.  While arch: is optonal, it is recommended.  The default arch is 'versal'.\n")
            print(f"    -dtb - Specify hardware dtb for the APU\n")
            print(f"    -kernel - option is filtered out, but is available for compatibility\n")
            apu_args_s = ' '.join(APU_args)
            help_cmd = f'{binpath}/qemu-system-aarch64 {apu_args_s}'

        print(f"    -machine-path <path> - Optional path to use for multiarch files\n     Note: if specified caller is responsible for cleaning up the directory on exit\n")

        print(f"{help_cmd}\n")

        # Without this the Popen below can end up printing before the above
        sys.stdout.flush()

        process = subprocess.Popen(help_cmd, shell=True, stderr=subprocess.PIPE)
        status = process.wait()
    sys.exit(status)

if set(APU_args).intersection(set(help_options)) or set(PMU_args).intersection(set(help_options)) or set(PLM_args).intersection(set(help_options)) or set(ASU_args).intersection(set(help_options)):
    help(0)

if not PMU_args and not PLM_args:
    help(1)

# The APU QEMU will hang if it does get get a -hw-dtb
if dtb_arg and '-hw-dtb' in APU_args:
    sys.exit(f'\nERROR: Arguments -dtb and -hw-dtb conflict, only one should be specified.')
elif not dtb_arg and not '-hw-dtb' in APU_args:
    sys.exit(f'\nERROR: You must specify at least a -dtb or -hw-dtb for the APU.')

if mach_path_arg:
    if not os.path.isdir(mach_path_arg):
        sys.exit(f'\nERROR: Missing machine path for qemu: {mach_path_arg}')
    else:
        mach_path = os.path.realpath(mach_path_arg)
else:
    mach_path = tempfile.mkdtemp()

if PMU_args and (PLM_args or ASU_args):
    sys.exit("\nError: -pmu-args can not be used with -plm-args or -asu-args\n")

if ASU_args and not PLM_args:
    sys.exit("\nError: -asu-args required -plm-args\n")

if PMU_args:
    try:
        PMU_rom = PMU_args[PMU_args.index('-kernel')+1]
    except:
        PMU_rom = ""

    if not os.path.exists(PMU_rom):
        sys.exit(f'\nERROR: Missing PMU ROM: {PMU_rom}'
            '\nSee "meta-xilinx/README.qemu.md" for more information on accquiring the PMU ROM.\n')

if bootbin_arg:
    # List of valid bootgen arches and settings
    archs = [ 'versal', 'versal_2ve_2vm' ]
    boot_header_addr = {'versal' : '0xf201e000', 'versal_2ve_2vm' : '0xf201eec0'}
    hash_block_addr = {'versal' : None, 'versal_2ve_2vm' : '0xf201ecc0'}

    # The bootbin is separated into "arch:file"
    # Default to versal for legacy reasons
    arch = 'versal'
    bootbin_fn = bootbin_arg

    if ':' in bootbin_arg:
        (arch, bootbin_fn) = bootbin_arg.split(':', 1)

    if not os.path.isfile(bootbin_fn):
        print(f"\nERROR: bootbin file not found at {bootbin_fn}\n")
        sys.exit(1)

    if arch not in archs:
        print(f"\nERROR: bootbin arch {arch} not valid ({archs})\n")
        sys.exit(1)

    shutil.copyfile(bootbin_fn, f'{mach_path}/boot.bin')

    bootgen_command = [f'{binpath}/bootgen', '-arch', arch, '-dump', 'boot.bin']
    subprocess.run(bootgen_command + ['boot_files'], check=True, cwd=mach_path, stdout=subprocess.DEVNULL)

    bootgen_command = f"{binpath}/bootgen -arch {arch} -read {bootbin_fn}"
    result = subprocess.check_output(bootgen_command.split())
    bootgen_output = result.decode().splitlines()

    for i, l in enumerate(bootgen_output):
        if 'PARTITION HEADER TABLE (pmc_subsys.0.0)' in l:
            plm_line = bootgen_output[i+4]
        if 'BOOT HEADER' in l:
            pmc_line = bootgen_output[i+6]

    plm_load_addr = re.search(r"exec_addr_lo \(0x10\) : (0x\w*)\s*", plm_line).group(1)
    pmc_load_addr = re.search(r"pmccdo_load_addr \(0x20\) : (0x\w*)", pmc_line).group(1)
    print(f"qemu-system-aarch64-multiarch: INFO: Using QSPI/OSPI bootbin file")

    boot_bh_addr = boot_header_addr[arch]
    if boot_bh_addr and os.path.exists(f'{mach_path}/boot_bh.bin'):
        print(f"qemu-system-aarch64-multiarch: INFO: boot header load addr: {boot_bh_addr}")
        PLM_args.append(f"-device loader,file={mach_path}/boot_bh.bin,addr={boot_bh_addr},force-raw=on")
    else:
        print(f"\nERROR: boot header ({mach_path}/boot_bh.bin) or boot header ({boot_bh_addr}) address is missing\n")
        sys.exit(1)

    hash_addr = hash_block_addr[arch]
    if hash_addr:
        if os.path.exists(f'{mach_path}/HashBlock0.bin'):
            print(f"qemu-system-aarch64-multiarch: INFO: boot HashBlock0 load addr: {hash_addr}")
            PLM_args.append(f"-device loader,file={mach_path}/HashBlock0.bin,addr={hash_addr}")
        else:
            print(f"\nERROR: Hash Block 0 ({mach_path}/HashBlock0.bin) is missing\n")
            sys.exit(1)

    print(f"qemu-system-aarch64-multiarch: INFO: bootbin file dump PMC load addr: {pmc_load_addr}")
    PLM_args.append(f"-device loader,file={mach_path}/pmc_cdo.bin,addr={pmc_load_addr},force-raw=on")

    print(f"qemu-system-aarch64-multiarch: INFO: bootbin file dump PLM load addr: {plm_load_addr}")
    PLM_args.append(f"-device loader,file={mach_path}/plm.bin,addr={plm_load_addr},force-raw=on")
    PLM_args.append(f"-device loader,addr={plm_load_addr},cpu-num=1")

    # Add blank to make the above easier to read
    print()

# We use the normal -dtb argument to specify the hw-dtb for the multiarch runner, as -dtb is not applicable.
if dtb_arg:
    APU_args.append(f"-hw-dtb {dtb_arg}")

# We need to switch tcp serial arguments (if they exist, e.g. qemurunner) to get the output
tcp_serial_ports = [i for i, s in enumerate(APU_args) if 'tcp:127.0.0.1:' in s]

#FIXME for next yocto release (dont need to switch ports anymore, they will be provided correctly upstream
# We can only switch these if there are exactly two, otherwise we can't assume what is being executed so we leave it as is
if len(tcp_serial_ports) == 2:
    APU_args[tcp_serial_ports[0]],APU_args[tcp_serial_ports[1]] = APU_args[tcp_serial_ports[1]],APU_args[tcp_serial_ports[0]]

mb_cmd = ""
riscv_cmd = ""
if PMU_args:
    pmu_args_s = ' '.join(PMU_args)
    mb_cmd = f'{binpath}/qemu-system-microblazeel {pmu_args_s} -machine-path {mach_path}'
    print(f"PMU instance cmd: {mb_cmd}\n")

if PLM_args:
    plm_args_s = ' '.join(PLM_args)
    mb_cmd = f'{binpath}/qemu-system-microblazeel {plm_args_s} -machine-path {mach_path}'
    print(f"PLM instance cmd: {mb_cmd}\n")

if ASU_args:
    asu_args_s = ' '.join(ASU_args)
    riscv_cmd = f'{binpath}/qemu-system-riscv32 {asu_args_s} -machine-path {mach_path}'
    print(f"ASU instance cmd: {riscv_cmd}\n")

apu_args_s = ' '.join(APU_args)
apu_cmd = f'{binpath}/qemu-system-aarch64 {apu_args_s} -machine-path {mach_path}'
print(f"APU instance cmd: {apu_cmd}\n")

if riscv_cmd:
    process_riscv = subprocess.Popen(riscv_cmd, shell=True, stderr=subprocess.PIPE)

if mb_cmd:
    process_mb = subprocess.Popen(mb_cmd, shell=True, stderr=subprocess.PIPE)

if apu_cmd:
    process_apu = subprocess.Popen(apu_cmd, shell=True, stderr=subprocess.PIPE)

rc = 0
error_msg = ""
if apu_cmd:
    apu_rc = process_apu.wait()
    if apu_rc:
        rc = apu_rc
        error_msg += '\nQEMU APU instance failed (%s):\n%s' % (apu_rc, process_apu.stderr.read().decode())

# We only report errors for MB and Risc-V if the APU failed (or APU wasn't started)
# otherwise we assume that the shutdown of the APU is the cause of the error
try:
    if mb_cmd:
        mb_rc = process_mb.wait(timeout=5)
        if mb_rc and (apu_cmd and apu_rc):
            rc += mb_rc
            error_msg += '\nQEMU MB instance failed (%s):\n%s' % (mb_rc, process_mb.stderr.read().decode())
except subprocess.TimeoutExpired:
    process_mb.kill()

try:
    if riscv_cmd:
        riscv_rc = process_riscv.wait(timeout=5)
        if riscv_rc and (apu_cmd and apu_rc):
            rc += riscv_rc
            error_msg += '\nQEMU Risc-V instance failed (%s):\n%s' % (riscv_rc, process_riscv.stderr.read().decode())
except subprocess.TimeoutExpired:
    process_riscv.kill()

if not mach_path_arg:
    shutil.rmtree(mach_path)

if apu_rc != 0:
    print(error_msg)

sys.exit(apu_rc)
