#!/usr/bin/env python3

# This file is part of the MicroPython project, http://micropython.org/
# The MIT License (MIT)
# Copyright (c) 2019 Damien P. George

import os
import subprocess
import sys
import argparse

sys.path.append("../tools")
import pyboard

# Paths for host executables
CPYTHON3 = os.getenv("MICROPY_CPYTHON3", "python3")
MICROPYTHON = os.getenv("MICROPY_MICROPYTHON", "../ports/unix/build-coverage/micropython")

NATMOD_EXAMPLE_DIR = "../examples/natmod/"

# Supported tests and their corresponding mpy module
TEST_MAPPINGS = {
    "btree": "btree/btree_$(ARCH).mpy",
    "deflate": "deflate/deflate_$(ARCH).mpy",
    "framebuf": "framebuf/framebuf_$(ARCH).mpy",
    "heapq": "heapq/heapq_$(ARCH).mpy",
    "random": "random/random_$(ARCH).mpy",
    "re": "re/re_$(ARCH).mpy",
}

# Code to allow a target MicroPython to import an .mpy from RAM
injected_import_hook_code = """\
import sys, io, vfs
class __File(io.IOBase):
  def __init__(self):
    self.off = 0
  def ioctl(self, request, arg):
    return 0
  def readinto(self, buf):
    buf[:] = memoryview(__buf)[self.off:self.off + len(buf)]
    self.off += len(buf)
    return len(buf)
class __FS:
  def mount(self, readonly, mkfs):
    pass
  def chdir(self, path):
    pass
  def stat(self, path):
    if path == '/__injected.mpy':
      return tuple(0 for _ in range(10))
    else:
      raise OSError(-2) # ENOENT
  def open(self, path, mode):
    return __File()
vfs.mount(__FS(), '/__remote')
sys.path.insert(0, '/__remote')
sys.modules['{}'] = __import__('__injected')
"""


class TargetSubprocess:
    def __init__(self, cmd):
        self.cmd = cmd

    def close(self):
        pass

    def run_script(self, script):
        try:
            p = subprocess.run(
                self.cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, input=script
            )
            return p.stdout, None
        except subprocess.CalledProcessError as er:
            return b"", er


class TargetPyboard:
    def __init__(self, pyb):
        self.pyb = pyb
        self.pyb.enter_raw_repl()

    def close(self):
        self.pyb.exit_raw_repl()
        self.pyb.close()

    def run_script(self, script):
        try:
            self.pyb.enter_raw_repl()
            output = self.pyb.exec_(script)
            output = output.replace(b"\r\n", b"\n")
            return output, None
        except pyboard.PyboardError as er:
            return b"", er


def run_tests(target_truth, target, args, stats):
    for test_file in args.files:
        # Find supported test
        for k, v in TEST_MAPPINGS.items():
            if test_file.find(k) != -1:
                test_module = k
                test_mpy = v.replace("$(ARCH)", args.arch)
                break
        else:
            print("----  {} - no matching mpy".format(test_file))
            continue

        # Read test script
        with open(test_file, "rb") as f:
            test_file_data = f.read()

        # Create full test with embedded .mpy
        test_script = b"import sys\nsys.path.remove('')\n\n"
        try:
            with open(NATMOD_EXAMPLE_DIR + test_mpy, "rb") as f:
                test_script += b"__buf=" + bytes(repr(f.read()), "ascii") + b"\n"
        except OSError:
            print("----  {} - mpy file not compiled".format(test_file))
            continue
        test_script += bytes(injected_import_hook_code.format(test_module), "ascii")
        test_script += test_file_data

        # Run test under MicroPython
        result_out, error = target.run_script(test_script)

        # Work out result of test
        extra = ""
        if error is None and result_out == b"SKIP\n":
            result = "SKIP"
        elif error is not None:
            result = "FAIL"
            extra = " - " + str(error)
        else:
            # Check result against truth
            try:
                with open(test_file + ".exp", "rb") as f:
                    result_exp = f.read()
                error = None
            except OSError:
                result_exp, error = target_truth.run_script(test_file_data)
            if error is not None:
                result = "TRUTH FAIL"
            elif result_out != result_exp:
                result = "FAIL"
                print(result_out)
            else:
                result = "pass"

        # Accumulate statistics
        stats["total"] += 1
        if result == "pass":
            stats["pass"] += 1
        elif result == "SKIP":
            stats["skip"] += 1
        else:
            stats["fail"] += 1

        # Print result
        print("{:4}  {}{}".format(result, test_file, extra))


def main():
    cmd_parser = argparse.ArgumentParser(
        description="Run dynamic-native-module tests under MicroPython"
    )
    cmd_parser.add_argument(
        "-p", "--pyboard", action="store_true", help="run tests via pyboard.py"
    )
    cmd_parser.add_argument(
        "-d", "--device", default="/dev/ttyACM0", help="the device for pyboard.py"
    )
    cmd_parser.add_argument(
        "-a", "--arch", default="x64", help="native architecture of the target"
    )
    cmd_parser.add_argument("files", nargs="*", help="input test files")
    args = cmd_parser.parse_args()

    target_truth = TargetSubprocess([CPYTHON3])

    if args.pyboard:
        target = TargetPyboard(pyboard.Pyboard(args.device))
    else:
        target = TargetSubprocess([MICROPYTHON])

    stats = {"total": 0, "pass": 0, "fail": 0, "skip": 0}
    run_tests(target_truth, target, args, stats)

    target.close()
    target_truth.close()

    print("{} tests performed".format(stats["total"]))
    print("{} tests passed".format(stats["pass"]))
    if stats["fail"]:
        print("{} tests failed".format(stats["fail"]))
    if stats["skip"]:
        print("{} tests skipped".format(stats["skip"]))

    if stats["fail"]:
        sys.exit(1)


if __name__ == "__main__":
    main()