//
// Created by Johnathon Slightham on 2025-05-25.
//

#include <chrono>
#include <cstddef>
#include <ranges>

#include "MessagingInterface.h"
#include "AngleControlMessageBuilder.h"
#include "CallBuilder.h"
#include "ConfigManager.h"
#include "constants/module.h"
#include "flatbuffers_generated/SendCall_generated.h"
#include "freertos/FreeRTOS.h"
#include "freertos/idf_additions.h"
#include "freertos/queue.h"
#include "freertos/semphr.h"
#include "MPIMessageBuilder.h"

#define TAG "MessagingInterface"
#define MAX_RX_WAIT_TIME_MS 3000
#define MAX_TX_WAIT_TIME_MS 200
#define THREAD_POOL_SLEEP_AFTER_FAIL_MS 10

bool Writer::write(std::vector<uint8_t>& return_value) {
    if (m_called) { // Only allow returning once
        return false;
    }
    m_called = true;

    auto [send_data, send_size] = m_builder.build_return_call(m_unique_id, return_value);
    m_messaging_interface.send((uint8_t*)send_data, send_size, PC_ADDR, FN_CALL_TAG, true);

    return true;
}

MessagingInterface::~MessagingInterface() {
    vSemaphoreDelete(m_map_semaphore);
}

int MessagingInterface::send(uint8_t* buffer, const size_t size, const uint8_t destination, const uint8_t tag, const bool durable) {
    Flatbuffers::MPIMessageBuilder builder;
    const auto [mpi_buffer, mpi_size] = builder.build_mpi_message(Messaging::MessageType_PTP, m_config_manager.get_module_id(), destination, 0, durable, tag, std::vector<uint8_t>(buffer, buffer + size));
    // Intentionally set sequence_number = 0 so that the messages can get cached

    m_router->send_msg((uint8_t *)mpi_buffer, mpi_size);
    return 0;
}

int MessagingInterface::broadcast(uint8_t* buffer, size_t size, bool durable) {
    // todo: impl
    return 0;
}

std::optional<SizeAndSource> MessagingInterface::recv(uint8_t* buffer, size_t size, const uint8_t tag) {
    checkOrInsertTag(tag);

    // The message should already be validated if added to message queue
    auto maybe_message = m_tag_to_queue.at(tag)->dequeue(std::chrono::milliseconds(MAX_RX_WAIT_TIME_MS));
    if (!maybe_message) { return std::nullopt; }
    const auto message = std::move(*maybe_message);
    const auto mpi_message = Flatbuffers::MPIMessageBuilder::parse_mpi_message(message->data());

    const auto bytes_written = std::min(size, (size_t)mpi_message->length());
    std::memcpy(buffer, mpi_message->payload()->data(), bytes_written);

    return SizeAndSource{(size_t)bytes_written, (uint8_t)mpi_message->sender()};
}

int MessagingInterface::sendrecv(uint8_t* send_buffer, const size_t send_size, const uint8_t dest, const uint8_t send_tag, uint8_t* recv_buffer, const size_t recv_size, const uint8_t recv_tag, const bool durable) {
    send(send_buffer, send_size, dest, send_tag, durable);
    recv(recv_buffer, recv_size, recv_tag);

    return 0;
}

void MessagingInterface::handleRecv(std::unique_ptr<std::vector<uint8_t>>&& buffer) {
    const auto mpi_message = Flatbuffers::MPIMessageBuilder::parse_mpi_message(buffer->data());

    checkOrInsertTag(mpi_message->tag());

    m_tag_to_queue.at(mpi_message->tag())->enqueue(std::move(buffer), std::chrono::milliseconds(MAX_TX_WAIT_TIME_MS));
}

void MessagingInterface::checkOrInsertTag(const uint8_t tag) {
    xSemaphoreTake(m_map_semaphore, portMAX_DELAY);
    if (!m_tag_to_queue.contains(tag)) {
        m_tag_to_queue[tag] = std::make_unique<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>>(MPI_QUEUE_SIZE);
    }
    xSemaphoreGive(m_map_semaphore);
}

std::pair<std::vector<uint8_t>, std::vector<Orientation>> MessagingInterface::get_physically_connected_modules() const {
    return m_router->get_physically_connected_modules();
}

Messaging::ConnectionType MessagingInterface::get_connection_type() const {
    if (this->m_router->get_leader() == m_config_manager.get_module_id()) {
        return Messaging::ConnectionType_DIRECT;
    }
    return Messaging::ConnectionType_HOP;
}

uint8_t MessagingInterface::get_leader() const {
    return this->m_router->get_leader();
}

void MessagingInterface::handleRecvCall(void *args) {
    const auto that = (MessagingInterface *)args;

    while (!that->m_stop_thread) {
        auto buffer = std::make_unique<std::vector<uint8_t>>();
        buffer->resize(RPC_CALL_PARAMETERS_MAX_SIZE);
        if (const auto rx_info = that->recv(buffer->data(), RPC_CALL_PARAMETERS_MAX_SIZE, FN_CALL_TAG)) {
            buffer->resize((*rx_info).bytes_written);

            flatbuffers::Verifier verifier(buffer->data(), buffer->size());
            if (!Messaging::VerifySendCallBuffer(verifier)) {
                continue;
                ESP_LOGW(TAG, "Got invalid rpc call flatbuffer");
            }

            if (!that->m_rpc_call_queue->enqueue(std::move(buffer), std::chrono::milliseconds(RPC_CALL_ENQUEUE_TIMEOUT_MS))) {
                ESP_LOGW(TAG, "Dropping RPC calls since queue is filled, and passed timeout");
            }
        }
    }
}

void MessagingInterface::processRPCCall(void *args) {
    const auto that = (MessagingInterface *)args;

    Flatbuffers::CallBuilder builder{};

    while (!that->m_stop_thread) {
        if(auto maybe_rpc_call = that->m_rpc_call_queue->dequeue(std::chrono::milliseconds(RPC_CALL_DEQUEUE_TIMEOUT_MS))) {
            auto rpc_call = *std::move(maybe_rpc_call);
            auto parsed_rpc_call = builder.parse_send_call(rpc_call->data());

            if (!that->m_tag_to_fn[parsed_rpc_call->tag()]) {
                ESP_LOGW(TAG, "RPC call failed, tried to call unregistered tag %d", parsed_rpc_call->tag());
                vTaskDelay(pdMS_TO_TICKS(THREAD_POOL_SLEEP_AFTER_FAIL_MS));
                continue;
            }

            Writer w{parsed_rpc_call->unique_id(), *that};
            that->m_tag_to_fn[parsed_rpc_call->tag()](parsed_rpc_call->parameters()->data(), (size_t)parsed_rpc_call->length(), w);
        }
    }
}

bool MessagingInterface::register_function(uint8_t function_tag, std::function<void(const uint8_t *, size_t, Writer&)> f) {
    if (m_tag_to_fn[function_tag]) {
        ESP_LOGW(TAG, "Attempt to register a function that was already registered");
        return false;
    }

    m_tag_to_fn[function_tag] = f;
    return true;
}