#include <cstdio>
#include <cstring>
#include <initializer_list>
#include "pico/stdlib.h"
#include "pico/error.h"
#include "hardware/dma.h"
#include "hardware/clocks.h"
#include "wizchip_spi.h"
#include "wizchip_qspi_pio.h"
#include "wizchip_qspi_pio.pio.h"

#define PIO_PROGRAM_NAME wizchip_pio_spi_quad_write_read
#define PIO_PROGRAM_FUNC __CONCAT(PIO_PROGRAM_NAME, _program)
#define PIO_PROGRAM_GET_DEFAULT_CONFIG_FUNC __CONCAT(PIO_PROGRAM_NAME, _program_get_default_config)
#define PIO_OFFSET_WRITE_BITS __CONCAT(PIO_PROGRAM_NAME, _offset_write_bits)
#define PIO_OFFSET_WRITE_BITS_END __CONCAT(PIO_PROGRAM_NAME, _offset_write_bits_end)
#define PIO_OFFSET_READ_BITS_END __CONCAT(PIO_PROGRAM_NAME, _offset_read_bits_end)

static constexpr uint32_t PADS_DRIVE = PADS_BANK0_GPIO0_DRIVE_VALUE_12MA;
static constexpr uint32_t IRQ_DELAY_NS = 100;
static constexpr uint32_t QSPI_LOOP_CNT = 2;

static struct {
    pio_hw_t *pio;
    uint8_t pio_func_sel;
    int8_t pio_offset;
    int8_t pio_sm;
    int8_t dma_out;
    int8_t dma_in;
} state;

static uint16_t mk_cmd_buf(uint8_t *pdst, uint8_t opcode, uint16_t addr) {
    pdst[0] = ((opcode >> 7 & 0x01) << 4) | ((opcode >> 6 & 0x01) << 0);
    pdst[1] = ((opcode >> 5 & 0x01) << 4) | ((opcode >> 4 & 0x01) << 0);
    pdst[2] = ((opcode >> 3 & 0x01) << 4) | ((opcode >> 2 & 0x01) << 0);
    pdst[3] = ((opcode >> 1 & 0x01) << 4) | ((opcode >> 0 & 0x01) << 0);
    pdst[4] = (uint8_t)(addr >> 8);
    pdst[5] = (uint8_t)(addr);
    pdst[6] = 0;
    return 7;
}

static uint32_t data_pin_mask() {
    return (1u << PIO_SPI_DATA_IO0_PIN) | (1u << PIO_SPI_DATA_IO1_PIN) |
           (1u << PIO_SPI_DATA_IO2_PIN) | (1u << PIO_SPI_DATA_IO3_PIN);
}

static __noinline void ns_delay(uint32_t ns) {
    uint32_t cycles = ns * (clock_get_hz(clk_sys) >> 16u) / (1000000000u >> 16u);
    busy_wait_at_least_cycles(cycles);
}

void wizchip_pio_init() {
    for (auto pin : {PIO_SPI_DATA_IO0_PIN, PIO_SPI_DATA_IO1_PIN, PIO_SPI_DATA_IO2_PIN, PIO_SPI_DATA_IO3_PIN}) {
        gpio_init(pin);
        gpio_set_dir(pin, GPIO_OUT);
        gpio_put(pin, false);
    }
    gpio_init(PIN_CS);
    gpio_set_dir(PIN_CS, GPIO_OUT);
    gpio_put(PIN_CS, true);
    gpio_init(PIN_INT);
    gpio_set_dir(PIN_INT, GPIO_IN);
    gpio_set_pulls(PIN_INT, false, false);

    pio_hw_t *pios[2] = {pio0, pio1};
    uint pio_index = 1;
    if (!pio_can_add_program(pios[pio_index], &PIO_PROGRAM_FUNC)) {
        pio_index ^= 1;
        assert(pio_can_add_program(pios[pio_index], &PIO_PROGRAM_FUNC));
    }

    state.pio = pios[pio_index];
    state.dma_in = -1;
    state.dma_out = -1;

    static_assert(GPIO_FUNC_PIO1 == GPIO_FUNC_PIO0 + 1);
    state.pio_func_sel = GPIO_FUNC_PIO0 + pio_index;
    state.pio_sm = (int8_t)pio_claim_unused_sm(state.pio, true);
    state.pio_offset = pio_add_program(state.pio, &PIO_PROGRAM_FUNC);

    pio_sm_config sm_config = PIO_PROGRAM_GET_DEFAULT_CONFIG_FUNC(state.pio_offset);
    sm_config_set_clkdiv_int_frac(&sm_config, WIZNET_SPI_CLKDIV_MAJOR_DEFAULT, WIZNET_SPI_CLKDIV_MINOR_DEFAULT);

    hw_write_masked(&pads_bank0_hw->io[PIO_SPI_SCK_PIN],
                    (uint)PADS_DRIVE << PADS_BANK0_GPIO0_DRIVE_LSB,
                    PADS_BANK0_GPIO0_DRIVE_BITS);
    hw_write_masked(&pads_bank0_hw->io[PIO_SPI_SCK_PIN],
                    1u << PADS_BANK0_GPIO0_SLEWFAST_LSB,
                    PADS_BANK0_GPIO0_SLEWFAST_BITS);

    sm_config_set_out_pins(&sm_config, PIO_SPI_DATA_IO0_PIN, 4);
    sm_config_set_in_pins(&sm_config, PIO_SPI_DATA_IO0_PIN);
    sm_config_set_set_pins(&sm_config, PIO_SPI_DATA_IO0_PIN, 4);
    sm_config_set_sideset(&sm_config, 1, false, false);
    sm_config_set_sideset_pins(&sm_config, PIO_SPI_SCK_PIN);
    sm_config_set_in_shift(&sm_config, false, true, 8);
    sm_config_set_out_shift(&sm_config, false, true, 8);
    hw_set_bits(&state.pio->input_sync_bypass, data_pin_mask());
    pio_sm_set_config(state.pio, state.pio_sm, &sm_config);
    pio_sm_set_consecutive_pindirs(state.pio, state.pio_sm, PIO_SPI_SCK_PIN, 1, true);

    for (auto pin : {PIO_SPI_DATA_IO0_PIN, PIO_SPI_DATA_IO1_PIN, PIO_SPI_DATA_IO2_PIN, PIO_SPI_DATA_IO3_PIN}) {
        gpio_set_function(pin, (gpio_function_t)state.pio_func_sel);
        gpio_set_pulls(pin, false, true);
        gpio_set_input_hysteresis_enabled(pin, true);
    }

    pio_sm_exec(state.pio, state.pio_sm, pio_encode_set(pio_pins, 1));

    state.dma_out = (int8_t)dma_claim_unused_channel(true);
    state.dma_in = (int8_t)dma_claim_unused_channel(true);
}

void wizchip_pio_close() {
    if (state.pio_sm >= 0) {
        if (state.pio_offset != -1) {
            pio_remove_program(state.pio, &PIO_PROGRAM_FUNC, state.pio_offset);
        }
        pio_sm_unclaim(state.pio, state.pio_sm);
    }
    if (state.dma_out >= 0) dma_channel_unclaim(state.dma_out);
    if (state.dma_in >= 0) dma_channel_unclaim(state.dma_in);
}

void wizchip_pio_frame_start() {
    for (auto pin : {PIO_SPI_DATA_IO0_PIN, PIO_SPI_DATA_IO1_PIN, PIO_SPI_DATA_IO2_PIN, PIO_SPI_DATA_IO3_PIN})
        gpio_set_function(pin, (gpio_function_t)state.pio_func_sel);
    gpio_set_function(PIO_SPI_SCK_PIN, (gpio_function_t)state.pio_func_sel);
    gpio_pull_down(PIO_SPI_SCK_PIN);
    gpio_put(PIN_CS, false);
}

void wizchip_pio_frame_end() {
    gpio_put(PIN_CS, true);
    ns_delay(IRQ_DELAY_NS);
}

void wizchip_pio_read(uint8_t opcode, uint16_t addr, uint8_t* buf, uint16_t len) {
    uint8_t cmd[8] = {};
    uint16_t cmd_len = mk_cmd_buf(cmd, opcode, addr);

    pio_sm_set_enabled(state.pio, state.pio_sm, false);
    pio_sm_set_wrap(state.pio, state.pio_sm, state.pio_offset, state.pio_offset + PIO_OFFSET_READ_BITS_END - 1);
    pio_sm_clear_fifos(state.pio, state.pio_sm);
    pio_sm_set_pindirs_with_mask(state.pio, state.pio_sm, data_pin_mask(), data_pin_mask());
    pio_sm_restart(state.pio, state.pio_sm);
    pio_sm_clkdiv_restart(state.pio, state.pio_sm);

    pio_sm_put(state.pio, state.pio_sm, cmd_len * QSPI_LOOP_CNT - 1);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_out(pio_x, 32));
    pio_sm_put(state.pio, state.pio_sm, len - 1);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_out(pio_y, 32));
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_jmp(state.pio_offset));

    dma_channel_abort(state.dma_out);
    dma_channel_abort(state.dma_in);

    dma_channel_config out_cfg = dma_channel_get_default_config(state.dma_out);
    channel_config_set_transfer_data_size(&out_cfg, DMA_SIZE_8);
    channel_config_set_bswap(&out_cfg, true);
    channel_config_set_dreq(&out_cfg, pio_get_dreq(state.pio, state.pio_sm, true));
    dma_channel_configure(state.dma_out, &out_cfg, &state.pio->txf[state.pio_sm], cmd, cmd_len, true);

    dma_channel_config in_cfg = dma_channel_get_default_config(state.dma_in);
    channel_config_set_transfer_data_size(&in_cfg, DMA_SIZE_8);
    channel_config_set_bswap(&in_cfg, true);
    channel_config_set_dreq(&in_cfg, pio_get_dreq(state.pio, state.pio_sm, false));
    channel_config_set_write_increment(&in_cfg, true);
    channel_config_set_read_increment(&in_cfg, false);
    dma_channel_configure(state.dma_in, &in_cfg, buf, &state.pio->rxf[state.pio_sm], len, true);

    pio_sm_set_enabled(state.pio, state.pio_sm, true);
    __compiler_memory_barrier();
    dma_channel_wait_for_finish_blocking(state.dma_out);
    dma_channel_wait_for_finish_blocking(state.dma_in);
    __compiler_memory_barrier();
    pio_sm_set_enabled(state.pio, state.pio_sm, false);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_mov(pio_pins, pio_null));
}

void wizchip_pio_write(uint8_t opcode, uint16_t addr, uint8_t* buf, uint16_t len) {
    uint8_t cmd[8] = {};
    uint16_t cmd_len = mk_cmd_buf(cmd, opcode, addr);
    uint16_t total = len + cmd_len;

    pio_sm_set_enabled(state.pio, state.pio_sm, false);
    pio_sm_set_wrap(state.pio, state.pio_sm, state.pio_offset, state.pio_offset + PIO_OFFSET_WRITE_BITS_END - 1);
    pio_sm_clear_fifos(state.pio, state.pio_sm);
    pio_sm_set_pindirs_with_mask(state.pio, state.pio_sm, data_pin_mask(), data_pin_mask());
    pio_sm_restart(state.pio, state.pio_sm);
    pio_sm_clkdiv_restart(state.pio, state.pio_sm);

    pio_sm_put(state.pio, state.pio_sm, total * QSPI_LOOP_CNT - 1);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_out(pio_x, 32));
    pio_sm_put(state.pio, state.pio_sm, 0);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_out(pio_y, 32));
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_jmp(state.pio_offset));

    dma_channel_abort(state.dma_out);

    dma_channel_config out_cfg = dma_channel_get_default_config(state.dma_out);
    channel_config_set_transfer_data_size(&out_cfg, DMA_SIZE_8);
    channel_config_set_bswap(&out_cfg, true);
    channel_config_set_dreq(&out_cfg, pio_get_dreq(state.pio, state.pio_sm, true));

    pio_sm_set_enabled(state.pio, state.pio_sm, true);
    dma_channel_configure(state.dma_out, &out_cfg, &state.pio->txf[state.pio_sm], cmd, cmd_len, true);
    dma_channel_wait_for_finish_blocking(state.dma_out);
    dma_channel_configure(state.dma_out, &out_cfg, &state.pio->txf[state.pio_sm], buf, len, true);
    dma_channel_wait_for_finish_blocking(state.dma_out);

    const uint32_t stall = 1u << (PIO_FDEBUG_TXSTALL_LSB + state.pio_sm);
    state.pio->fdebug = stall;
    while (!(state.pio->fdebug & stall)) tight_loop_contents();

    __compiler_memory_barrier();
    pio_sm_set_consecutive_pindirs(state.pio, state.pio_sm, PIO_SPI_DATA_IO0_PIN, 4, false);
    pio_sm_exec(state.pio, state.pio_sm, pio_encode_mov(pio_pins, pio_null));
    pio_sm_set_enabled(state.pio, state.pio_sm, false);
}