#!/usr/bin/env python3
"""Runs test within Qemu and Renode simulators."""
import argparse
import os
import re
import sys
import tempfile

import io
import pexpect


parser = argparse.ArgumentParser(
    description="Run a springbok test on an simulator.")

parser.add_argument('simulator',
                    help='Select a simulator',
                    choices=['renode', 'qemu'])
parser.add_argument('elf',
                    help='Elf to execute on a simulator')
parser.add_argument('--renode-path',
                    help="Path to renode simulator")
parser.add_argument('--qemu-path',
                    help="Path to qemu simulator")
parser.add_argument('--timeout', type=int,
                    help="Timeout for test", default=1000)

class Simulation: # pylint: disable=too-few-public-methods
    """ Base class for simulation """
    def __init__(self, simulator_cmd):
        self.simulator_cmd = simulator_cmd
        self.buffer = io.StringIO()
        self.child = None
        self.termination_strings = ["main returned", "Exception occurred"]

    def run(self, timeout=1000):
        """ Run the simulation command and quit the simulation."""
        self.child = pexpect.spawn(self.simulator_cmd, encoding='utf-8')
        self.child.logfile = self.buffer
        self.child.expect(self.termination_strings, timeout=timeout)
        self.child.send("\nq\n")
        self.child.expect(pexpect.EOF, timeout=timeout)
        self.child.close()
        self.buffer.seek(0)
        return self.buffer.read()

class QemuSimulation(Simulation): # pylint: disable=too-few-public-methods
    """ Qemu simulation """
    def __init__(self, path, elf):
        self.qemu_simulator_cmd = (
            "%(sim)s -M springbok -nographic -d springbok -device loader,file=%(elf)s")
        self.sim_params = {"sim": path, "elf": elf}
        super().__init__(self.qemu_simulator_cmd % self.sim_params)


class RenodeSimulation(Simulation): # pylint: disable=too-few-public-methods
    """ Renode Simulation """
    def __init__(self, path, elf):
        # Get the ROOTDIR path if it exists
        self.rootdir = os.environ.get('ROOTDIR', default=None)
        if self.rootdir is None:
            parser.error("ROOTDIR environment variable not set.")
        renode_script = """
path set @%(rootdir)s
$bin=@%(elf)s
include @sim/config/springbok.resc
sysbus.cpu2 PerformanceInMips 2000
emulation SetGlobalQuantum "1"
start
sysbus.vec_controlblock WriteDoubleWord 0xc 0"""
        self.script_params = {
            "rootdir": self.rootdir,
            "elf":elf
        }
        self.renode_script = renode_script % self.script_params

        self.renode_args = [
            "%s" % path,
            '--disable-xwt',
            ' --console',
            '--plain',
        ]
        self.renode_simulator_cmd = " ".join(self.renode_args)
        super().__init__(self.renode_simulator_cmd)

    def run(self, timeout=120):
        file_desc, script_path = tempfile.mkstemp(suffix=".resc")
        try:
            with os.fdopen(file_desc, 'w') as tmp:
                tmp.write(self.renode_script)
                tmp.flush()
            self.simulator_cmd += " %s" % script_path
            test_output = super().run()
        finally:
            os.remove(script_path)
        return test_output


Simulators = {
    "qemu": QemuSimulation,
    "renode": RenodeSimulation
}

args = parser.parse_args()

simulators_paths = {
    "renode": args.renode_path,
    "qemu": args.qemu_path
}


def main():
    """ Run a test and check for Pass or Fail """
    simulator_path = simulators_paths[args.simulator]
    if simulator_path is None:
        parser.error(
            "Must provide path to simulator %s, use argument --%s-path" % (args.simulator,
                                                                           args.simulator))

    simulator_class = Simulators[args.simulator]
    simulator = simulator_class(simulator_path, args.elf)
    output = simulator.run(timeout=args.timeout)
    # mono API generates escape characters at the termination. Need to clean up.
    # TODO(hcindyl): Remove this when Renode fix the mono call.
    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
    output = ansi_escape.sub("", output)
    print(output)
    failure_strings = ["FAILED", "Exception occurred"]
    if any(x in output for x in failure_strings):
        sys.exit(1)
    # Grab the return code from the output string with regex
    # Syntax: "main returned: ", <code> (<hex_code>)
    return_string = re.compile(
        r"\"main returned:\s\",(?P<ret_code>\s[0-9]+\s*)")
    code = return_string.search(output)
    sys.exit(int(code.group(1)))


if __name__ == "__main__":
    main()