#include <unistd.h>
#include <assert.h>

#include "Bootloader.h"

Bootloader::Bootloader(Uart &uart) : uart(uart)
{
    syn();
    get();
    get_chip_id();
    get_flash_memory_size();
}

void Bootloader::syn()
{
    unsigned char buffer = 0x7f;
    uart.transmit(&buffer, 1);
    int rec = uart.receive(&buffer, 1);
    if ((rec >= 1) && ((buffer == ACK) || (buffer == NAK)))
    {
        is_sync = true;
    }
}

void Bootloader::get()
{
    unsigned char buffer[0xff];
    send_command(commands.get);
    usleep(10000);
    int rec = uart.receive(buffer, sizeof(buffer));
    if ((rec == 15) && ((buffer[0] == ACK) && (buffer[14] == ACK)))
    {
        chip.protocol_version_major = (buffer[2] & 0xf0) >> 4;
        chip.protocol_version_minor = buffer[2] & 0x0f;
        commands.get = buffer[3];
        commands.get_version_and_read_protection_status = buffer[4];
        commands.get_id = buffer[5];
        commands.read_memory = buffer[6];
        commands.go = buffer[7];
        commands.write_memory = buffer[8];
        commands.erase = buffer[9];
        commands.write_protect = buffer[10];
        commands.write_unprotect = buffer[11];
        commands.readout_protect = buffer[12];
        commands.readout_unprotect = buffer[13];
    }
}

void Bootloader::get_chip_id()
{
    unsigned char buffer[0xff];
    send_command(commands.get_id);
    usleep(10000);
    int rec = uart.receive(buffer, sizeof(buffer));
    if ((rec == 5) && ((buffer[0] == ACK) && (buffer[4] == ACK)))
    {
        chip.id = (buffer[2] << 8) + buffer[3];
        if (chip.id == SUPPORTED_DEVICE_ID)
        {
            is_supported = true;
        }
    }
}

void Bootloader::get_flash_memory_size()
{
    unsigned char buffer[2];
    read_memory(FLASH_MEMORY_SIZE_DATA_REGISTER_ADDRESS, buffer, sizeof(buffer));
    chip.flash_size = ((buffer[1] << 8) + buffer[0]) * 1024;
}

bool Bootloader::read_memory(uint32_t start_address, uint8_t *mem, uint16_t size)
{
    if (size > 256)
    {
        return false;
    }
    bool ret = false;
    unsigned char buffer[0x2];
    send_command(commands.read_memory);
    if (wait_for_ack())
    {
        send_address(start_address);
        if (wait_for_ack())
        {
            buffer[0] = size - 1;
            buffer[1] = buffer[0] ^ 0xff;
            uart.transmit(buffer, 2);

            if(wait_for_ack())
            {
                unsigned char* wptr = mem;
                int rec = uart.receive(wptr, size);
                int to_receice = size - rec;
                while(to_receice > 0) {
                    wptr += rec;
                    assert(wptr <= mem + size);
                    rec = uart.receive(wptr, to_receice);
                    to_receice -= rec;
                }
                if(to_receice == 0) {
                    ret = true;
                }
            }
        }
    }
    return ret;
}

bool Bootloader::write_memory(uint32_t start_address, uint8_t *mem, uint16_t size)
{
    if (size > 256)
    {
        return false;
    }
    bool ret = false;
    unsigned char buffer[0x101];
    send_command(commands.write_memory);
    if (wait_for_ack())
    {
        send_address(start_address);
        if (wait_for_ack())
        {
            uint16_t len = size;
            if (len % 4)
            {
                len = len + (4 - len % 4);
            }
            buffer[0] = (uint8_t)(len - 1);
            memcpy(&buffer[1], mem, size);

            buffer[len + 1] = 0;
            for (uint16_t i = 0; i <= len; i++)
            {
                buffer[len + 1] ^= buffer[i];
            }
            uart.transmit(buffer, len + 2);
            if (wait_for_ack())
            {
                ret = true;
            }
        }
    }
    return ret;
}

bool Bootloader::full_flash_erase()
{
    bool ret = false;
    unsigned char buffer[0x03];
    send_command(commands.erase);
    if (wait_for_ack())
    {
        buffer[0] = 0xff; // full flash erase
        buffer[1] = 0xff; // full flash erase
        buffer[2] = buffer[0] ^ buffer[1];
        uart.transmit(buffer, 3);
        if (wait_for_ack())
        {
            ret = true;
        }
    }
    return ret;
}

bool Bootloader::go(uint32_t start_address)
{
    bool ret = false;
    unsigned char buffer[0x101];
    send_command(commands.go);
    usleep(10000);
    int rec = uart.receive(buffer, sizeof(buffer));
    if ((rec == 1) && (buffer[0] == ACK))
    {
        send_address(start_address);
        if (wait_for_ack())
        {
            ret = true;
        }
    }
    return ret;
}

void Bootloader::send_command(uint8_t command)
{
    uint8_t buffer[2];
    buffer[0] = command;
    buffer[1] = buffer[0] ^ 0xff;
    uart.transmit(buffer, 2);
}

void Bootloader::send_address(uint32_t address)
{
    uint8_t buffer[5];
    buffer[0] = (address >> 24) & 0xff;
    buffer[1] = (address >> 16) & 0xff;
    buffer[2] = (address >> 8) & 0xff;
    buffer[3] = address & 0xff;
    buffer[4] = buffer[0] ^ buffer[1] ^ buffer[2] ^ buffer[3];
    uart.transmit(buffer, 5);
}

bool Bootloader::wait_for_ack()
{
    unsigned char rx;
    for (uint32_t i = 0; i < 10000; i++)
    {
        int rec = uart.receive(&rx, 1);
        if (rec > 0)
        {
            if (rx == ACK)
            {
                return true;
            }
            else if (rx == NAK)
            {
                return false;
            }
        }
    }
    return false;
}

uint16_t Bootloader::chip_id()
{
    return chip.id;
}

uint32_t Bootloader::flash_size()
{
    return chip.flash_size;
}

bool Bootloader::chip_is_supported()
{
    return is_supported;
}