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

#ifndef MESSAGINGINTERFACE_H
#define MESSAGINGINTERFACE_H

#include <memory>
#include <unordered_map>
#include <vector>

#include "flatbuffers_generated/TopologyMessage_generated.h"
#include "CallBuilder.h"
#include "BlockingQueue.h"
#include "constants/app_comms.h"
#include "CommunicationRouter.h"

#define FN_CALL_TAG 100  // reserved tag for RPC functionality

struct SizeAndSource {
    size_t bytes_written;
    uint8_t sender;
};

class MessagingInterface;  // fwd def

class Writer {
public:
    Writer(uint8_t unique_id, MessagingInterface &messaging_interface) : m_unique_id(unique_id), m_messaging_interface(messaging_interface) {}

    // Return a value from an RPC call.
    bool write(std::vector<uint8_t>& return_value);
private:
    uint8_t m_unique_id;
    MessagingInterface &m_messaging_interface;
    Flatbuffers::CallBuilder m_builder{};  // todo: this may be expensive to keep allocating, potentially move to MI & guard with lock
    bool m_called = false;
};

class MessagingInterface {
public:
    explicit MessagingInterface()
        : m_config_manager(ConfigManager::get_instance()),
            m_mpi_rx_queue(std::make_unique<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>>(RX_QUEUE_SIZE)),
            m_rpc_call_queue(std::make_unique<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>>(RPC_CALL_QUEUE_SIZE)),
            m_router(std::make_unique<CommunicationRouter>([this](std::unique_ptr<std::vector<uint8_t>>&& buffer) { handleRecv(std::move(buffer)); })),
            m_map_semaphore(xSemaphoreCreateMutex()) {
                xTaskCreate(handleRecvCall, "rpc_rx", 3072, this, 5, &m_handle_recv_call_task);
                m_thread_pool.resize(RPC_CALL_THREAD_POOL_SIZE);
                for (int i = 0; i < RPC_CALL_THREAD_POOL_SIZE; i++) {
                     xTaskCreate(processRPCCall, "rpc_thread", 3072, this, 5, &m_thread_pool[i]);
                }
            };

    ~MessagingInterface();

    int send(uint8_t* buffer, size_t size, uint8_t destination, uint8_t tag, bool durable);
    int broadcast(uint8_t* buffer, size_t size, bool durable);
    std::optional<SizeAndSource> recv(uint8_t * buffer, size_t size, uint8_t tag);
    int sendrecv(uint8_t *send_buffer, size_t send_size, uint8_t dest, uint8_t send_tag, uint8_t *recv_buffer, size_t recv_size, uint8_t recv_tag, bool durable);
    std::pair<std::vector<uint8_t>, std::vector<Orientation>> get_physically_connected_modules() const;
    Messaging::ConnectionType get_connection_type() const;
    uint8_t get_leader() const;

    // Register a function for RPC (to be executed in a threadpool). Similar interface to GRPC with the writer.
    bool register_function(uint8_t function_tag, std::function<void(const uint8_t *, size_t, Writer&)>);

private:
    void handleRecv(std::unique_ptr<std::vector<uint8_t>>&& buffer);
    static void handleRecvCall(void *args);
    static void processRPCCall(void *args);

    void checkOrInsertTag(uint8_t tag);

    ConfigManager& m_config_manager;
    uint16_t m_sequence_number = 0;
    std::unique_ptr<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>> m_mpi_rx_queue;
    std::unique_ptr<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>> m_rpc_call_queue;
    std::unique_ptr<CommunicationRouter> m_router;
    SemaphoreHandle_t m_map_semaphore;
    std::unordered_map<uint8_t, std::unique_ptr<BlockingQueue<std::unique_ptr<std::vector<uint8_t>>>>> m_tag_to_queue;
    std::unordered_map<uint8_t, std::function<void(const uint8_t *, size_t, Writer&)>> m_tag_to_fn;

    std::atomic<bool> m_stop_thread{false};
    TaskHandle_t m_handle_recv_call_task;
    std::vector<TaskHandle_t> m_thread_pool;
};

#endif //MESSAGINGINTERFACE_H