import serial
import struct
from google.protobuf.message import DecodeError
from serial.tools import list_ports

from firmware_pb2 import RequestDeviceList, Device, ResponseDeviceList, UsbPackageType

def make_header(typeid: UsbPackageType, length: int) -> bytearray: 
    struct_format = '<HHB'  # '<' for little-endian, 'H' for uint16_t, 'B' for uint8_t
    # Calculate the check byte as the sum of length and type
    typeidint = int(typeid) 
    check = (length & 0xFF) + ((length >> 8) & 0xFF) + (typeidint & 0xFF) + ((typeidint >> 8) & 0xFF)
    packed_data = struct.pack(struct_format, length, typeid, check)
    return packed_data

def send_package(typeid : UsbPackageType, data: bytearray, serial: serial.Serial):
    head = make_header(typeid, len(data))
    package = head + data
    serial.write(package)

if __name__ == "__main__":
    stm_port = None
    for port in list_ports.comports():
        print(port)
        if "STM32 Virtual ComPort" in port.description:
            stm_port = port.device
            break

    if stm_port is None:
        print("STM32 Virtual ComPort not found")
        exit(-1)
    else:
        # Open the serial port
        ser = serial.Serial(stm_port,baudrate=5000000)


        # Create a RequestDeviceList message
        request = RequestDeviceList()
        request.msg = 1  # or whatever value you want to set

        # Serialize the request to a bytearray
        request_data = request.SerializeToString()

        # Send the request
        send_package(UsbPackageType.REQUEST_DEVICE_LIST, request_data, ser)

        # Now wait for the response
        while True:
            # Read the header from the serial port
            response_header = ser.read(5)  # assuming the header is 5 bytes long

            # Unpack the header to get the length and type
            length, typeid, check = struct.unpack('<HHB', response_header)

            # Check if the type is RESPONSE_DEVICE_LIST
            if typeid == UsbPackageType.RESPONSE_DEVICE_LIST:
                # Read the response data from the serial port
                response_data = ser.read(length)

                # Try to parse the data as a ResponseDeviceList message
                try:
                    response = ResponseDeviceList.FromString(response_data)
                    # If we get here, we successfully parsed the response. Break the loop.
                    break
                except DecodeError:
                    # If we get a DecodeError, it means the data we read is not a valid
                    # ResponseDeviceList message. Ignore it and continue reading.
                    pass

        # At this point, 'response' contains the ResponseDeviceList message
        print(response)