#!/usr/bin/env python3

# Xilinx QEMU wrapper to launch both PMU and APU instances (multiarch)
import os
import subprocess
import sys
import tempfile
import shutil

binpath = os.path.dirname(os.path.abspath(__file__))
mach_path = tempfile.mkdtemp()


# Separate PMU and APU arguments
APU_args = sys.argv[1:]
PMU_args = []
PLM_args = []

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 PMU_args and PLM_args:
    sys.exit("\nError: -pmu-args can not be used with -plm-args\n")

if ('--help' in APU_args) or (not PMU_args and not PLM_args):
    print("AMD FPGA QEMU multiarch wrapper\n")
    print("Version 2024.1")
    print("")
    print("Usage:")
    print("  %s <APU options> [-pmu-args <pmu options>]" % (sys.argv[0]))
    print("  %s <APU options> [-plm-args <plm options>]" % (sys.argv[0]))
    print("")
    sys.exit(1)

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

    if not os.path.exists(PMU_rom):
        error_msg = '\nError: Missing PMU ROM: %s' % PMU_rom
        error_msg += '\nSee "meta-xilinx/README.qemu.md" for more information on accquiring the PMU ROM.\n'
        sys.exit(error_msg)

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

#NEED TO FIX 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 = ""
if PMU_args:
    mb_cmd =  binpath + '/qemu-system-microblazeel ' + ' '.join(PMU_args) + ' -machine-path ' + mach_path

    print("PMU instance cmd: %s\n" % mb_cmd)

if PLM_args:
    mb_cmd =  binpath + '/qemu-system-microblazeel ' + ' '.join(PLM_args) + ' -machine-path ' + mach_path

    print("PLM instance cmd: %s\n" % mb_cmd)

apu_cmd =  binpath + '/qemu-system-aarch64 ' + ' '.join(APU_args) + ' -machine-path ' + mach_path

print("APU instance cmd: %s\n" % apu_cmd)


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)

error_msg = ""
if apu_cmd and process_apu.wait():
    # We only check for failures on the MB instance if APU fails
    error_msg += '\nQEMU APU instance failed:\n%s' % process_apu.stderr.read().decode()

    if mb_cmd and process_mb.wait():
        error_msg += '\nQEMU MB instance failed:\n%s' % process_mb.stderr.read().decode()

shutil.rmtree(mach_path)
sys.exit(error_msg)