#!/usr/bin/env python3

import os
import sys
import time
import subprocess
import serial.tools.list_ports
from test_common import (
    DeviceUnderTestAeadUARTP,
    eprint,
    FileMutex,
    run_nist_lws_aead_test,
)


def get_serial():
    ports = serial.tools.list_ports.comports()
    devices = [
        p.device
        for p in ports
        if (p.vid == 4292 and p.pid == 60000)
    ]
    devices.sort()
    return devices[0]


class ESP32(DeviceUnderTestAeadUARTP):

    def __init__(self, build_dir):
        DeviceUnderTestAeadUARTP.__init__(self)

        self.uart_device = get_serial()
        devname = os.path.basename(self.uart_device)
        self.lock = FileMutex('/var/lock/lwc-compare.%s.lock' % devname)
        self.build_dir = build_dir

    def reset(self):
        self.ser.setDTR(False)  # IO0=HIGH
        self.ser.setRTS(True)   # EN=LOW, chip in reset
        time.sleep(0.1)
        self.ser.setDTR(False)  # IO0=HIGH
        self.ser.setRTS(False)  # EN=HIGH, chip out of reset
        time.sleep(1)

    def prepare(self):
        self.ser = serial.Serial(
            self.uart_device,
            baudrate=500000,
            timeout=5)
        self.reset()
        DeviceUnderTestAeadUARTP.prepare(self)

    def flash(self):
        pipe = subprocess.PIPE
        previous_dir = os.path.abspath(os.curdir)
        os.chdir(self.build_dir)
        cmd = ['platformio', 'run', '-e', 'esp32dev', '--target', 'upload']
        cmd.extend(['--upload-port', self.uart_device])
        p = subprocess.Popen(
            cmd, stdout=sys.stderr, stdin=pipe)
        stdout, stderr = p.communicate("")
        assert p.returncode == 0
        eprint("Firmware flashed.")
        os.chdir(previous_dir)

    def dump_ram(self):
        return None


def main(argv):
    if len(argv) != 3:
        print("Usage: test LWC_AEAD_KAT.txt build_dir")
        return 1

    run_nist_lws_aead_test(argv[1], argv[2], 0x0020, ESP32)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))