#!/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:]
mbtype=''

if '-pmu-args' in APU_args:
    MB_args = APU_args[APU_args.index('-pmu-args')+1]
    APU_args.remove('-pmu-args')
    APU_args.remove(MB_args)
    MB_args = MB_args.split()
    PMU_rom = MB_args[MB_args.index('-kernel')+1]
    mbtype='PMU'
elif '-plm-args' in APU_args:
    MB_args = APU_args[APU_args.index('-plm-args')+1]
    APU_args.remove('-plm-args')
    APU_args.remove(MB_args)
    MB_args = MB_args.split()
    mbtype='PLM'
else:
    error_msg = '\nMultiarch not setup properly.'
    sys.exit(error_msg)

error_msg = None
if (mbtype == 'PMU' and os.path.exists(PMU_rom)) or mbtype == 'PLM':

    # 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]

    # 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 =  binpath + '/qemu-system-microblazeel ' + ' '.join(MB_args) + ' -machine-path ' + mach_path
    apu_cmd =  binpath + '/qemu-system-aarch64 ' + ' '.join(APU_args) + ' -machine-path ' + mach_path

    # Debug prints
    print('\n%s instance cmd: %s\n' % (mbtype, mb_cmd))
    print('APU instance cmd: %s\n' % apu_cmd)


    # Invoke QEMU pmu instance
    process_pmu = subprocess.Popen(mb_cmd, shell=True, stderr=subprocess.PIPE)

    # Invoke QEMU APU instance
    process_apu = subprocess.Popen(apu_cmd, shell=True, stderr=subprocess.PIPE)
    if process_apu.wait():
        error_msg = '\nQEMU APU instance failed:\n%s' % process_apu.stderr.read().decode()

else:
    if mbtype == 'PMU':
        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'

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