import serial
import struct
import time
import sys
from google.protobuf.message import DecodeError
from serial.tools import list_ports
from crccheck.crc import Crc

from firmware_pb2 import FirmwareStart, FirmwarePackage, FirmwarePackageAck, FirmwareDone, UsbPackageType


Crc32 = Crc(32,0x04C11DB7,0xffffffff) 

def load_firmware(filename):
    # Open the file in binary mode
    with open(filename, 'rb') as f:
        data = f.read()
    # Calculate the CRC32 checksum
    crc = Crc32.calc(data)
    # Get the file size
    size = len(data)
    # Split the data into 256-byte chunks
    packages = [data[i:i+256] for i in range(0, len(data), 256)]
    return crc, size, packages


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)


def receive_ack(serial):
    # Read the header
    header_data = serial.read(5)  # header size is 5 bytes
    length, typeid, check = struct.unpack('<HHB', header_data)

    # Verify the check byte
    check_calculated = (length & 0xFF) + ((length >> 8) & 0xFF) + (typeid & 0xFF) + ((typeid >> 8) & 0xFF)
    if check != check_calculated:
        print('Header check byte mismatch')
        return None

    # Read the message
    message_data = serial.read(length)
    if len(message_data) != length:
        print('Incomplete message')
        return None

    if typeid != UsbPackageType.FIRMWAREPACKAGEACK:
        return None

    # Parse the message
    ack = FirmwarePackageAck()
    try:
        ack.ParseFromString(message_data)
    except DecodeError:
        print('Failed to parse FirmwarePackageAck')
        return None

    return ack

FILENAME = 'firmware.bin'
ID = 0
if len(sys.argv) > 1:
    FILENAME = sys.argv[1]
if len(sys.argv) > 2:
    ID = int(sys.argv[2])

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)

    crc, size, packages = load_firmware(FILENAME)

    # Create a FirmwareStart message
    start = FirmwareStart()
    start.name = FILENAME
    start.size = size
    start.packages = len(packages)
    start.device_id = ID
    start.crc_fw = crc

    # Send the FirmwareStart message
    print(start)
    send_package(UsbPackageType.FIRMWARESTART, start.SerializeToString(), ser)
    #time.sleep(1)  # wait for the device to process the message

    # Send the firmware packages
    for (i,pack_data) in enumerate(packages):
        package = FirmwarePackage()
        package.counter = i
        package.crc_pac = Crc32.calc(pack_data)
        package.device_id = start.device_id
        package.data = pack_data

        # Send the FirmwarePackage message
        print(package)
        print(hex(package.crc_pac))
        send_package(UsbPackageType.FIRMWAREPACKAGE, package.SerializeToString(), ser)


        # Wait for the FirmwarePackageAck message
        ack = receive_ack(ser)
        print(ack)
        if not ack.ack:
            print(f'Package {i} not acknowledged')
            exit(-1)


    # Send the FirmwareDone message
    done = FirmwareDone()
    done.size = start.size
    done.crc_fw = start.crc_fw
    done.device_id = start.device_id
    print(done)
    send_package(UsbPackageType.FIRMWAREDONE, done.SerializeToString(), ser)