#include "uart.h"

#include "async.h"
#include "gpio.h"
#include "lock.h"
#include "pol0.h"
#include "ring_buffer.h"
#include "trace.h"
#include "uart_async.h"
#include "xuartlite.h"

namespace {
using async::AwaitableType;

constexpr uintptr_t kUart0BaseAddress = UART0_BASE;
XUartLite uart0_inst;
XUartLite_Config uart0_config = {
    .DeviceId = 0,
    .RegBaseAddr = kUart0BaseAddress,
    .BaudRate = 115200,
    .UseParity = false,
    .DataBits = 8,
};

constexpr size_t kUartRxBufferSize = 256;
std::array<std::byte, kUartRxBufferSize> rx_buffer = {};
RingBuffer rx_ring_buffer{.buffer = rx_buffer};

constexpr size_t kUartTxBufferSize = 256;
std::array<std::byte, kUartTxBufferSize> tx_buffer = {};
RingBuffer tx_ring_buffer{.buffer = tx_buffer};

XUartLite* uart0 = &uart0_inst;

volatile bool sending;

void StartReceiving() {
    while (1) {
        if (rx_ring_buffer.FreeSpace() < 1) {
            // woops, full. discard some data
            // TODO: keep track of overrun stats
            rx_ring_buffer.Pop(1);
        }
        if (XUartLite_Recv(uart0, rx_ring_buffer.RawWritePointer(), 1) < 1) {
            break;
        }
        rx_ring_buffer.Push(1);
    }
}

void StartSending() {
    if (sending) {
        return;
    }
    size_t tosend = tx_ring_buffer.ContiguousAvailableData();
    if (tosend < 1) {
        return;
    }
    sending = true;
    XUartLite_Send(uart0, tx_ring_buffer.RawReadPointer(), tosend);
}
}  // namespace

void InitUarts() {
    XUartLite_CfgInitialize(uart0, &uart0_config, uart0_config.RegBaseAddr);

    XUartLite_SetSendHandler(uart0, HandleUartTxFromIsr, nullptr);
    XUartLite_SetRecvHandler(uart0, HandleUartRxFromIsr, nullptr);
    StartReceiving();
    XUartLite_EnableInterrupt(uart0);

    sending = false;
}

void UartWriteCrash(std::span<const std::byte> data) {
    XUartLite_DisableInterrupt(uart0);
    while (data.size() > 0) {
        while (XUartLite_IsSending(uart0)) {
        }
        auto* dat =
            reinterpret_cast<uint8_t*>(const_cast<std::byte*>(data.data()));
        uint8_t sent = XUartLite_Send(uart0, dat, data.size());
        data = data.subspan(sent);
    }
    while (XUartLite_IsSending(uart0)) {
    }
    XUartLite_Send(uart0, nullptr, 0);  // reset buffer before enabling interrupts
    XUartLite_EnableInterrupt(uart0);
}

async::task<> UartWrite(std::span<const std::byte> data) {
    while (!tx_ring_buffer.Store(data)) {
        co_await async::await(AwaitableType::kUartTx);
    }

    {
        InterruptLock lock;
        StartSending();
    }
}

async::task<> UartWriteLoop(
    async::gimme<std::span<const std::byte>>& data_gen) {
    while (1) {
        auto data = co_await data_gen;
        while (!tx_ring_buffer.Store(data)) {
            co_await async::await(AwaitableType::kUartTx);
        }

        {
            InterruptLock lock;
            StartSending();
        }
    }
}

// TODO: use chunks to allow receiving more than 256 bytes at once
void UartReadBlocking(std::span<std::byte> data) {
    while (!rx_ring_buffer.Load(data)) {
    }
}

void UartWriteBlocking(std::span<const std::byte> data) {
    while (!tx_ring_buffer.Store(data)) {
    }

    {
        InterruptLock lock;
        StartSending();
    }
}

async::task<std::byte> UartReadLoop() {
    std::byte c;
    while (1) {
        while (!rx_ring_buffer.Load(std::span{&c, 1})) {
            co_await async::await(AwaitableType::kUartRx);
        }

        co_yield c;
    }
}

async::task<buffer> UartRead(int size) {
    auto buff = buffer::make(size);
    while (!rx_ring_buffer.Load(buff.data)) {
        co_await async::await(AwaitableType::kUartRx);
    }
    co_return buff;
}

void HandleUartTxFromIsr(void*, unsigned int transmitted) {
    sending = false;
    tx_ring_buffer.Pop(transmitted);
    StartSending();
    async::resume(AwaitableType::kUartTx);
}

void HandleUartRxFromIsr(void*, unsigned int transmitted) {
    rx_ring_buffer.Push(transmitted);
    StartReceiving();
    async::resume(AwaitableType::kUartRx);
}

void HandleUartIsr() { XUartLite_InterruptHandler(uart0); }