#include "gpio.h"
#include "ssi_reg.h"
#include "system.h"

void __attribute__((section(".text.boot2_pre"))) boot2_copy_self(void)
{
    #define BOOT2_FLASHCMD_READ_DATA (0x03)
    #define BOOT2_FLASH_OFFSET       (256)
    extern uint32_t _boot2_copy_self_start_addr;
    extern uint32_t _boot2_copy_self_end_addr;

    uint32_t cmd;
    uint32_t length, tx_count, rx_count, rx_skip;
    uint8_t *copy_to_addr;
    uint8_t rxbyte;

    /* force QSPI_CS low level as selected */
    ioqspi_hw->io[1].ctrl = GPIO_OVER_OUT_LOW;

    /* send command and length */
    cmd = (BOOT2_FLASHCMD_READ_DATA << 24) | BOOT2_FLASH_OFFSET;
    for (int i = 0; i < 4; ++i) {
        ssi_hw->dr0 = cmd >> 24;
        cmd <<= 8;
    }
    /* init parameter */
    length = (uint32_t)&_boot2_copy_self_end_addr - (uint32_t)&_boot2_copy_self_start_addr;
    tx_count = length;
    rx_count = length;
    rx_skip = 4;
    copy_to_addr = (uint8_t *)&_boot2_copy_self_start_addr;
    /* copy data from flash to ram */
    while (tx_count || rx_skip || rx_count) {
        uint32_t tx_level = ssi_hw->txflr;
        uint32_t rx_level = ssi_hw->rxflr;
        if (tx_count && tx_level + rx_level < 14) {
            ssi_hw->dr0 = 0;
            tx_count--;
        }
        if (rx_level) {
            rxbyte = ssi_hw->dr0;
            if (rx_skip) {
                rx_skip--;
            } else {
                *copy_to_addr++ = rxbyte;
                rx_count--;
            }
        }
    }

    /* force QSPI_CS high level as disselected */
    ioqspi_hw->io[1].ctrl = GPIO_OVER_OUT_HIGH;
}

#include "common.h"
#include "resets.h"
#include "clock.h"
#include "uart.h"
#include "flash.h"
#include "watchdog.h"
#include "timer.h"
#include "jump.h"

#define UART_ID uart0_hw

uint8_t tx_buffer[512] __attribute__((aligned(4)));
uint8_t rx_buffer[512] __attribute__((aligned(4)));

void process_return_ok(void)
{
    uart_write_block(UART_ID, (uint8_t *)"OKAY", 4);
}

void process_return_fail(void)
{
    uart_write_block(UART_ID, (uint8_t *)"FAIL", 4);
}

void process_flash_erase(uint32_t addr, uint32_t length)
{
    addr = addr & (~(FLASH_ERASE_SIZE - 1));
    length = (length + FLASH_ERASE_SIZE - 1) / FLASH_ERASE_SIZE;
    for (int i = 0; i < length; i++) {
        flash_erase(addr + i * FLASH_ERASE_SIZE);
    }
}

void process_flash_write(uint32_t addr, uint8_t *data, uint32_t length)
{
    uint32_t offset;
    uint32_t chunk_size;

    offset = addr & (FLASH_WRITE_SIZE - 1);
    if (offset) {
        chunk_size = FLASH_WRITE_SIZE - offset;
        if (chunk_size > length) {
            chunk_size = length;
        }
        flash_write(addr, data, chunk_size);
        addr += chunk_size;
        data += chunk_size;
        length -= chunk_size;
    }
    while (length >= FLASH_WRITE_SIZE) {
        flash_write(addr, data, FLASH_WRITE_SIZE);
        addr += FLASH_WRITE_SIZE;
        data += FLASH_WRITE_SIZE;
        length -= FLASH_WRITE_SIZE;
    }
    if (length > 0) {
        flash_write(addr, data, length);
    }
}

void process_flash_read(uint32_t addr, uint8_t *data, uint32_t length)
{
    while (length >= FLASH_READ_SIZE) {
        flash_read(addr, data, FLASH_READ_SIZE);
        uart_write_block(UART_ID, data, FLASH_READ_SIZE);
        addr += FLASH_READ_SIZE;
        length -= FLASH_READ_SIZE;
    }
    if (length > 0) {
        flash_read(addr, data, length);
        uart_write_block(UART_ID, data, length);
    }
}

void uart_process(uint16_t code, uint16_t length)
{
    uint8_t *p;
    uint32_t checksum_read, checksum_calc;
    uint16_t i;
    uint32_t v1, v2;
    
    p = rx_buffer + 8;
    if ((length > 0) && (length <= 4)) {
        return;
    }
    if (length > 4) {
        checksum_read = ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | ((uint32_t)p[3] << 0);
        p += 4;
        length -= 4;
        checksum_calc = 0;
        for (i = 0; i < length; i++) {
            checksum_calc += p[i];
        }
        if (checksum_read != checksum_calc) {
            return;
        }
    }

    v1 = ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | ((uint32_t)p[3] << 0);
    v2 = ((uint32_t)p[4] << 24) | ((uint32_t)p[5] << 16) | ((uint32_t)p[6] << 8) | ((uint32_t)p[7] << 0);
    if (code == 0x0001) {
        process_flash_erase(v1, v2);
        process_return_ok();
    } else if (code == 0x0002) {
        process_flash_write(v1, p + 8, v2);
        process_return_ok();
    } else if (code == 0x0003) {
        process_return_ok();
        process_flash_read(v1, tx_buffer, v2);
    }
}

void uart_state_machine(uart_hw_t *uart)
{
    static uint64_t time_fifo_empty = 0;
    static uint16_t uart_rx_length = 0;
    uint16_t code, code_inv, length, length_inv;

    if (uart_get_flags(uart) & UART_FLAG_RXFE) {
        if (timer_count_read() - time_fifo_empty < 100) {
            return;
        }
        time_fifo_empty = timer_count_read();
        if (uart_rx_length < 8) {
            uart_rx_length = 0;
            return;
        }
        code = ((uint16_t)rx_buffer[0] << 8) | (uint16_t)rx_buffer[1];
        code_inv = ((uint16_t)rx_buffer[2] << 8) | (uint16_t)rx_buffer[3];
        length = ((uint16_t)rx_buffer[4] << 8) | (uint16_t)rx_buffer[5];
        length_inv = ((uint16_t)rx_buffer[6] << 8) | (uint16_t)rx_buffer[7];
        if ((code != (uint16_t)~code_inv) || (length != (uint16_t)~length_inv)) {
            uart_rx_length = 0;
            return;
        }
        uart_process(code, length);
        uart_rx_length = 0;
    } else {
        rx_buffer[uart_rx_length++] = uart0_hw->dr & 0xFF;
        time_fifo_empty = timer_count_read();
        if (uart_rx_length >= sizeof(rx_buffer)) {
            uart_rx_length = 0;
            return;
        }
    }
}

int main(void)
{
    uint32_t boot_pin_low, boot_pin_high;
    struct fw_header_s *header;

    clock_ref_set(CLOCK_REF_SRC_XOSC_GLITCHLESS, 1);
    clock_sys_set(CLOCK_SYS_SRC_REF_GLITCHLESS, 1);
    clock_peri_set(ENABLE, CLOCK_PERI_SRC_XOSC);

    reset_unreset_blocks_wait(RESETS_BLOCK_IO_BANK0 | RESETS_BLOCK_PADS_BANK0 | RESETS_BLOCK_UART0 | RESETS_BLOCK_TIMER);
    gpio_init(0, GPIO_FUNC_UART | GPIO_PULL_UP | GPIO_DRIVE_4MA); /* UART_TX pin */
    gpio_init(1, GPIO_FUNC_UART | GPIO_PULL_UP | GPIO_SCHMITT | GPIO_PAD_IE | GPIO_PAD_OD); /* UART_RX pin */
    uart_init(uart0_hw, 750 * 1000, UART_DATABITS_8 | UART_PARITY_NONE | UART_STOPBITS_1);
    watchdog_start_tick(12); /* 12MHz / 12 = 1MHz, as 1us */
    timer_start();

    gpio_init(2, GPIO_FUNC_NULL | GPIO_PULL_DOWN | GPIO_SCHMITT | GPIO_PAD_IE | GPIO_PAD_OD); /* boot pin */
    boot_pin_low = 0;
    boot_pin_high = 0;
    for (uint32_t i = 0; i < 10; i++) {
        if (gpio_read(2)) {
            boot_pin_high++;
        } else {
            boot_pin_low++;
        }
    }
    if (boot_pin_high > boot_pin_low) {
        for (uint32_t i = 0; i < 4; i++) {
            tx_buffer[i] = 0x55;
        }
        uart_write_block(UART_ID, tx_buffer, 4);
        while (1) {
            uart_state_machine(UART_ID);
        }
    } else {
        flash_read(FW_ADDRESS, rx_buffer, FLASH_READ_SIZE);
        header = (struct fw_header_s *)rx_buffer;
        if (jump_check_is_failed(header)) {
            flash_enter_quad_xip(6);
            jump_to_address((uint32_t)(uintptr_t *)(FW_ADDRESS | XIP_BASE | 1));
        } else {
            jump_pre_config(header);
            jump_to_address(((struct fw_header_s *)header)->jump_address | 1);
        }
    }

    return 0;
}