#include <fe_global.h>
#include <fe_clocks.h>
#include <bootloader.h>
#include <cobs.h>

/* User configuration */
#include <fe_config_backend.h>

struct usart_config {
    uint32_t rx_rcc_ahbenr_flags;
    GPIO_TypeDef *rx_gpio;
    int rx_pin;
    int rx_alt;

    uint32_t tx_rcc_ahbenr_flags;
    GPIO_TypeDef *tx_gpio;
    int tx_pin;
    int tx_alt;

    uint32_t apb1enr_mask;
    uint32_t apb2enr_mask;
    USART_TypeDef *usart;
};

struct usart_config usart_gpios[FE_CONFIG_USART_COUNT] = {
    [FE_CONFIG_USART1_PB7] = {FE_CONFIG_GPIOB, 7, 7, FE_CONFIG_GPIOB, 6, 7, 0, RCC_APB2ENR_USART1EN, USART1},
};

void run_bootloader(void);
int bootloader_handle_cmd(USART_TypeDef *us, char *buf, size_t len);
void usart_init(struct usart_config *uc);
void usart_putc(USART_TypeDef *us, char c);
void usart_puts(USART_TypeDef *us, const char *data);
int usart_rx(USART_TypeDef *us);
int usart_tx_framed(USART_TypeDef * us, char *buf, size_t len);

void __libc_init_array(void) { /* we don't need this. */ }

void __assert_func (unused_a const char *file, unused_a int line, unused_a const char *function, unused_a const char *expr) {
    asm volatile ("bkpt #69");
    while(1) {}
}

void gpio_config(GPIO_TypeDef *gpio, int pin, int mode, int speed, int pullups, int alt) {
    gpio->MODER &= ~(3<<(2*pin));
    gpio->MODER |= mode<<(2*pin);
    gpio->OSPEEDR &= ~(3<<(2*pin));
    gpio->OSPEEDR |= speed<<(2*pin);
    gpio->PUPDR &= ~(3<<(2*pin));
    gpio->PUPDR |= pullups<<(2*pin);
    if (pin < 8) {
        gpio->AFR[0] &= ~(0xf << (4*pin));
        gpio->AFR[0] |= alt << (4*pin);
    } else {
        pin -= 8;
        gpio->AFR[1] &= ~(0xf << (4*pin));
        gpio->AFR[1] |= alt << (4*pin);
    }
}

int main(void)
{
    const struct fe_config_def *c = &fe_config;

    uint32_t old_moder = c->bootloader_enable_pin.gpio->MODER;
    uint32_t old_ospeedr = c->bootloader_enable_pin.gpio->OSPEEDR;
    uint32_t old_pupdr = c->bootloader_enable_pin.gpio->PUPDR;
    RCC->AHBENR |= c->bootloader_enable_pin.rcc_ahbenr_flags;
    int pullups = c->bootloader_enable_level ? 2 : 1;
    gpio_config(c->bootloader_enable_pin.gpio, c->bootloader_enable_pin.pin_number, 0, 0, pullups, 0);

    delay_ms(50);
    int enable_pin_state = ((c->bootloader_enable_pin.gpio->IDR >> c->bootloader_enable_pin.pin_number) & 1);
    if (enable_pin_state == c->bootloader_enable_level || !fe_check_img_valid()) {
        fe_config_clocks();
        run_bootloader();
        fe_system_reset();

    } else {
        c->bootloader_enable_pin.gpio->MODER = old_moder;
        c->bootloader_enable_pin.gpio->OSPEEDR = old_ospeedr;
        c->bootloader_enable_pin.gpio->PUPDR = old_pupdr;
        fe_jump_to_application();
    }

    /* Should never be reached. */
    assert(0);
    return 0;
}

#define MAX_RX_SIZE 512
char rx_buf[MAX_RX_SIZE];

void run_bootloader() {
    assert(fe_config.usart <= FE_CONFIG_USART_COUNT);
    struct usart_config *uc = &usart_gpios[fe_config.usart];

    usart_init(uc);
    usart_puts(uc->usart, fe_config.welcome_string);
    usart_putc(uc->usart, 0x00);

    struct cobs_decode_state cobs_st;
    cobs_decode_incremental_initialize(&cobs_st);
    while (42) {
        int c = usart_rx(uc->usart);
        if (c == -2) {
            /* ignore errors for now */
            continue;
        } else if (c == -1) {
            /* We received nothing */
            continue;
        }

        int rc = cobs_decode_incremental(&cobs_st, rx_buf, sizeof(rx_buf), c);
        if (rc == -1) {
            continue;

        } else if (rc < 0) {
            /* Ignore errors for now */
            continue;
        }

        (void)bootloader_handle_cmd(uc->usart, rx_buf, rc); /* ignore errors for now */
    }
}

enum bootloader_cmd {
    /* Generic commands */
    FE_CMD_REPLY = 0,
    FE_CMD_PING = 1,
    FE_CMD_IDENTIFY = 2,

    /* Bootloader commands */
    FE_CMD_ERASE = 16 + 0,
    FE_CMD_WRITE_BLOCK = 16 + 1,
    FE_CMD_REBOOT = 16 + 2,

    /* Reserved commands */
    FE_CMD_RESERVED0 = 'e', /* 0x65 / 101 Reserved for welcome string */
};

enum error_codes {
    FE_SUCCESS = 0,
    FE_ECMD = 1,
    FE_ESIZE = 2,
    FE_EINVAL = 3,
    FE_ESYS = 4,
};

#define FE_MAX_PACKET_SIZE 254
struct cmd_header {
    uint8_t cmd; /* Command code, see enum bootloader_cmd */
    uint32_t tag; /* Tag for request/response identification. Echoed back to requestor in response. */
} __attribute__((packed));

#define FE_MAX_PAYLOAD_SIZE (FE_MAX_PACKET_SIZE - sizeof(struct cmd_header))

struct id_response {
    uint8_t major_version;
    uint8_t minor_version;
    uint32_t idcode;
    uint8_t serial[6];
    char id_string[FE_MAX_PAYLOAD_SIZE - 2];
};

struct simple_response {
    int32_t return_code;
};

struct cmd_packet {
    struct cmd_header hdr;
    union {
        uint8_t data[FE_MAX_PAYLOAD_SIZE];
        struct id_response id_response;
        struct simple_response simple_response;
    };
};

static const uint8_t major_version = 1;
static const uint8_t minor_version = 0;
static const char *id_string = "Fenris Bootloader";

static int err_simple(USART_TypeDef *us, struct cmd_packet *pkt, int rc);
static int err_simple(USART_TypeDef *us, struct cmd_packet *pkt, int rc) {
    struct simple_response *res = &pkt->simple_response;
    pkt->hdr.cmd = FE_CMD_REPLY;
    /* leave tag unaffected */
    res->return_code = rc;
    return usart_tx_framed(us, (char *)pkt, sizeof(struct cmd_header) + sizeof(struct simple_response));
}

static size_t flash_write_addr = 0;

int bootloader_handle_cmd(USART_TypeDef *us, char *buf, size_t len) {
    struct cmd_packet *pkt = (struct cmd_packet *)buf;
    int rc = 0;
    if (len < sizeof(struct cmd_header)) {
        memset(buf, 0, sizeof(struct cmd_header));
        return err_simple(us, pkt, -FE_ESIZE);
    }
    int payload_len = len-sizeof(struct cmd_header);

    pkt->hdr.cmd = FE_CMD_REPLY;

    switch (pkt->hdr.cmd) {
        case FE_CMD_PING:
            return usart_tx_framed(us, buf, len);

        case FE_CMD_IDENTIFY: {
            if (payload_len != 0)
                return err_simple(us, pkt, -FE_ESIZE);

            struct id_response *res = &pkt->id_response;
            memcpy(res->serial, (void*)UID_BASE, 12);
            res->idcode = DBGMCU->IDCODE;
            res->major_version = major_version;
            res->minor_version = minor_version;
            strncpy(res->id_string, id_string, sizeof(res->id_string));
            len = sizeof(struct cmd_header) + 2 + strlen(id_string) + 1;
            if (len < FE_MAX_PACKET_SIZE)
                return usart_tx_framed(us, buf, len);
            return -10;
        }

        case FE_CMD_ERASE: {
            if (payload_len != 4)
                return err_simple(us, pkt, -FE_ESIZE);

            if (*((uint32_t*)pkt->data) != 0x54454c44)
                return err_simple(us, pkt, -FE_EINVAL);
            
            flash_unlock();

            if (erase_user_flash()) {
                flash_lock();
                return err_simple(us, pkt, -FE_ESYS);
            }

            flash_lock();
            flash_write_addr = flash_base;

            return err_simple(us, pkt, 0);
        }

        case FE_CMD_WRITE_BLOCK: {
            if (payload_len == 0)
                return err_simple(us, pkt, -FE_ESIZE);

            if ((payload_len&1) != 0)
                return err_simple(us, pkt, -FE_ESIZE);
            
            if (flash_write_addr + payload_len == flash_size) {
                /* FIXME sig check */
                if (0) { /* invalid signature */
                    return err_simple(us, pkt, -FE_EINVAL);
                }

            } else if (flash_write_addr + payload_len > flash_size) {
                return err_simple(us, pkt, -FE_EINVAL);
            }

            if (!flash_write(flash_write_addr, pkt->data, payload_len))
                return err_simple(us, pkt, -FE_ESYS);
            flash_write_addr += payload_len;
        }
        case FE_CMD_REBOOT: {
            fe_system_reset();
        }

        default: {
            return err_simple(us, pkt, -FE_ECMD);
        }
    }
}

void usart_init(struct usart_config *uc) {
    RCC->AHBENR |= uc->rx_rcc_ahbenr_flags | uc->tx_rcc_ahbenr_flags;
    gpio_config(uc->rx_gpio, uc->rx_pin, 2, 3, 1, uc->rx_alt);
    gpio_config(uc->tx_gpio, uc->tx_pin, 2, 3, 0, uc->tx_alt);
    
    RCC->APB1ENR |= uc->apb1enr_mask;
    RCC->APB2ENR |= uc->apb2enr_mask;

    uc->usart->CR1 = USART_CR1_TE | USART_CR1_RE;
    
    int bus_speed = uc->apb1enr_mask ? apb1_speed : apb2_speed;
    uc->usart->BRR = bus_speed * 16 / fe_config.baudrate / 16;
    uc->usart->CR1 |= USART_CR1_UE;
}

void usart_putc(USART_TypeDef *us, char c) {
    while (!(us->ISR & USART_ISR_TXE))
        ;
    us->TDR = c;
}

void usart_puts(USART_TypeDef *us, const char *data) {
    for (const char *c = data; *c; c++)
        usart_putc(us, *c);
}

int usart_rx(USART_TypeDef *us) {
    if ((us->ISR & USART_ISR_ORE) || (us->ISR & USART_ISR_PE)) {
        /* Ignore overruns or parity errors */
        us->ICR = USART_ICR_ORECF | USART_ICR_PECF;
        return -2;
    }

    if (us->ISR & USART_ISR_RXNE)
        return us->RDR;

    return -1;
}

static int cobs_encode_usart_output(void *userdata, char c);
static int cobs_encode_usart_output(void *userdata, char c) {
    usart_putc((USART_TypeDef *)userdata, c);
    return 0;
}

int usart_tx_framed(USART_TypeDef * us, char *buf, size_t len) {
    return cobs_encode_usart(cobs_encode_usart_output, us, buf, len);
}

void *memcpy(void *restrict dest, const void *restrict src, size_t n)
{
	unsigned char *d = dest;
	const unsigned char *s = src;

	for (; n; n--) *d++ = *s++;
	return dest;
}

void *memset(void *dest, int c, size_t n)
{
	unsigned char *s = dest;
	size_t k;

	/* Fill head and tail with minimal branching. Each
	 * conditional ensures that all the subsequently used
	 * offsets are well-defined and in the dest region. */

	if (!n) return dest;
	s[0] = c;
	s[n-1] = c;
	if (n <= 2) return dest;
	s[1] = c;
	s[2] = c;
	s[n-2] = c;
	s[n-3] = c;
	if (n <= 6) return dest;
	s[3] = c;
	s[n-4] = c;
	if (n <= 8) return dest;

	/* Advance pointer to align it at a 4-byte boundary,
	 * and truncate n to a multiple of 4. The previous code
	 * already took care of any head/tail that get cut off
	 * by the alignment. */

	k = -(uintptr_t)s & 3;
	s += k;
	n -= k;
	n &= -4;

	/* Pure C fallback with no aliasing violations. */
	for (; n; n--, s++) *s = c;

	return dest;
}