#include "uart.h"

#include "async.h"
#include "buffer.h"
#include "ring_buffer.h"

#include <array>
#include <mutex>
#include <vector>

namespace {
using async::AwaitableType;

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

uint8_t* tx_buff = nullptr;
uint16_t tx_size = 0;
uint16_t tx_size_orig = 0;

uint8_t* rx_buff = nullptr;
uint16_t rx_size = 0;
uint16_t rx_size_orig = 0;
std::vector<std::byte> rx_overflow;

std::mutex interrupt_lock;
std::mutex fu_mutex;
}  // namespace

//// Internal API
void FakeUart_Send(uint8_t* ptr, uint16_t size) {
    std::scoped_lock lock(fu_mutex);

    tx_buff = ptr;
    tx_size = size;
    tx_size_orig = size;
}
size_t FakeUart_Recv(uint8_t* ptr, uint16_t size) {
    std::scoped_lock lock(fu_mutex);

    if (size <= rx_overflow.size()) {
        std::copy(rx_overflow.begin(), rx_overflow.begin() + size,
                  std::as_writable_bytes(std::span{ptr, size}).begin());
        std::rotate(rx_overflow.begin(), rx_overflow.begin() + size, rx_overflow.end());
        rx_overflow.resize(rx_overflow.size() - size);

        return size;
    }

    rx_buff = ptr;
    rx_size = size;

    if (rx_overflow.empty()) {
        return 0;
    }

    std::copy(rx_overflow.begin(), rx_overflow.end(),
              std::as_writable_bytes(std::span{ptr, size}).begin());
    size_t copied = rx_overflow.size();
    rx_overflow.clear();

    rx_buff += copied;
    rx_size -= copied;

    return copied;
}
uint8_t FakeUart_IsSending() {
    std::scoped_lock lock(fu_mutex);
    return tx_size > 0;
}

//// Backdoor access API
/// Those do trigger UART interrupts when tx/rx buffer state changes
buffer FakeUart_Drain(uint16_t size) {
    std::scoped_lock isr_lock(interrupt_lock);  // not really the same, but that'll do
    std::scoped_lock lock(fu_mutex);

    size_t drained = std::min(size, tx_size);

    if (drained < 1) { 
        return buffer{};
    }

    buffer buff = buffer::make(drained);
    auto txb = std::as_bytes(std::span{tx_buff, tx_size});
    std::copy(txb.begin(), txb.begin() + drained, buff.data.begin());

    tx_buff += drained;
    tx_size -= drained;

    if (tx_size < 1) {
        fu_mutex.unlock();
        HandleUartTxFromIsr(nullptr, tx_size_orig);
        fu_mutex.lock();
    }

    return buff;
}
void FakeUart_Feed(std::span<const std::byte> data) {
    std::scoped_lock lock(fu_mutex);

    if (data.empty()) {
        return;
    }

    auto rxb = std::as_writable_bytes(std::span{rx_buff, rx_size});
    size_t fed = std::min(static_cast<size_t>(rx_size), data.size());

    if (data.size() > fed) {
        rx_overflow.insert(rx_overflow.end(), data.begin() + fed, data.end());
    }

    if (fed > 0) {
        std::copy(data.begin(), data.begin() + fed, rxb.begin());

        rx_buff += fed;
        rx_size -= fed;

        if (rx_size < 1) {
            HandleUartRxFromIsr(nullptr, rx_size_orig);
        }
    }
}
void FakeUart_Reset() {
    rx_overflow.clear();

    rx_buff = nullptr;
    rx_size = 0;
    tx_buff = nullptr;
    tx_size = 0;
}

//// Public API
void InitUarts() {}
void UartWriteCrash(std::span<const std::byte>) {}
void LogStuff() {}
void UartReadBlocking(std::span<std::byte>) {}
void UartWriteBlocking(std::span<const std::byte>) {}
void HandleUartIsr() {}
async::task<buffer> UartRead(int size) {
    co_return buffer{};
}

void HandleUartTxFromIsr(void*, unsigned int transmitted) {
    tx_ring_buffer.Pop(transmitted);
    if (tx_ring_buffer.AvailableData() > 0 && !FakeUart_IsSending()) {
        FakeUart_Send(tx_ring_buffer.RawReadPointer(),
                       tx_ring_buffer.ContiguousAvailableData());
    }
    async::resume(AwaitableType::kUartTx);
}

void HandleUartRxFromIsr(void*, unsigned int) {
    async::resume(AwaitableType::kUartRx);
}

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

    {
        std::scoped_lock lock(interrupt_lock);

        if (!FakeUart_IsSending()) {
            FakeUart_Send(tx_ring_buffer.RawReadPointer(),
                          tx_ring_buffer.ContiguousAvailableData());
        }
    }
}

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);
        }

        {
            std::scoped_lock lock(interrupt_lock);

            if (tx_ring_buffer.AvailableData() && !FakeUart_IsSending()) {
                FakeUart_Send(tx_ring_buffer.RawReadPointer(),
                              tx_ring_buffer.ContiguousAvailableData());
            }
        }
    }
}

async::task<uint8_t> UartReadLoop() {
    uint8_t c;
    while (1) {
        size_t received = FakeUart_Recv(&c, 1);
        // some data may already be in the fifo, but if not, wait:
        if (received < 1) {
            co_await async::await(AwaitableType::kUartRx);
        }

        co_yield c;
    }
}