/**
 * \file
 * \brief bidirectional communication through shared memory
 */

#ifndef _INIT_UMP_CHAN_H_
#define _INIT_UMP_CHAN_H_

#include <aos/aos.h>

/**
 * @brief Number of data bytes in a single buffer of the ring buffer
 */
#define UMP_RING_BUF_DATA_SIZE (CACHE_LINE_SIZE - 1)

struct ump_send_queue_entry;

typedef void (*ump_recv_header_callback_fn_t)(void *arg, size_t header_size, void *header, size_t payload_size);
typedef void (*ump_recv_payload_callback_fn_t)(void *arg, size_t payload_size, void *payload);
typedef void (*ump_send_callback_fn_t)(void *arg, struct ump_send_queue_entry *entry);

/**
 * @brief Specifies which role an endpoint would like to take in the UMP protocol
 *
 * This is just used to agree on which part of the shared frame should be used
 * to send data and which to receive data
 */
enum ump_role {
    UMP_ROLE_SERVER,
    UMP_ROLE_CLIENT,
    UMP_ROLE_COUNT, // How many roles exist
};

/**
 * @brief Stores metadata about a ring buffer
 */
struct ump_ring_state {
    /**
     * @brief first entry of the ring buffer
     */
    struct ump_ring_buf_entry *ring_start;
    /**
     * @brief Number of entries in the ring buffer
     */
    size_t ring_entry_count;
    /**
     * @brief Which entry should be read or written next
     */
    size_t ring_entry_next;
};

/**
 * @brief State of a send queue entry
 *
 * Used to keep track of which part of the message we are
 * currently sending
 */
enum ump_send_entry_state {
    UMP_SEND_ENTRY_STATE_START_MESSAGE,
    UMP_SEND_ENTRY_STATE_HEADER,
    UMP_SEND_ENTRY_STATE_PAYLOAD,
    UMP_SEND_ENTRY_STATE_COUNT, // How many states exist
};

/**
 * @brief An entry in the sending queue of a channel.
 *
 * Contains all metadata information needed to send this message
 */
struct ump_send_queue_entry {
    enum ump_send_entry_state state;

    size_t header_size;
    const void *header;

    size_t payload_size;
    const void *payload;

    ump_send_callback_fn_t callback;
    void *callback_arg;

    /**
     * @brief first this will refer to the header and when done with that
     * it will be updated to refer to the payload
     */
    size_t send_bytes_left;
    const void *send_buf_position;

    struct ump_send_queue_entry *next;
};

/**
 * @brief object containing metadata required for sending a message on a channel
 */
struct ump_send_chan {
    struct ump_ring_state ring_state;

    // required to hold this lock for any changes to the send_queue
    struct thread_mutex send_queue_lock;
    struct ump_send_queue_entry *send_queue_head;
    struct ump_send_queue_entry *send_queue_tail;
};

/**
 * @brief Used to keep track of which part of a message was received or is being received
 *
 */
enum recv_state {
    UMP_RECV_STATE_IDLE,
    UMP_RECV_STATE_START_MESSAGE,
    UMP_RECV_STATE_HEADER,
    UMP_RECV_STATE_PAYLOAD_IDLE,
    UMP_RECV_STATE_PAYLOAD,
    UMP_RECV_STATE_COUNT, // How many states exist
};

/**
 * @brief Metadata for the receiving side of the channel
 *
 */
struct ump_recv_chan {
    struct ump_ring_state ring_state;

    /**
     * @brief This must be held when editing this object
     * from a thread that is not the worker thread
     */
    struct thread_mutex recv_register_lock;

    enum recv_state state;

    /**
     * @brief Current read state.
     * This will first refer to the header and afterwards to the payload
     */
    size_t read_bytes_left;
    void *read_buf_position;

    size_t header_size;
    void *header;
    size_t next_payload_size;
    void *payload;

    ump_recv_header_callback_fn_t recv_header_callback;
    ump_recv_payload_callback_fn_t recv_payload_callback;
    void *callback_arg;
};

/**
 * @brief A single entry in a ring buffer
 */
struct ump_ring_buf_entry {
    /**
     * @brief data carried by this ring buffer entry
     */
    uint8_t data[UMP_RING_BUF_DATA_SIZE];
    /**
     * @brief specifies who currently owns this buffer entry.
     * 0 -> sender
     * 1 -> receiver
     */
    volatile uint8_t owner;
};
// for now we want each entry to be cache line sized
STATIC_ASSERT_SIZEOF(struct ump_ring_buf_entry, CACHE_LINE_SIZE);

/**
 * @brief The first data block for a new message.
 * This is only needed for the payload size right now
 */
struct ump_start_message {
    size_t payload_size;
    // NOTE rueegges: we could remove this field since the reciever also has
    // to specify it but for now it is here to ensure there are no missmatches
    // in sizes on the two endpoints
    size_t header_size;
};
STATIC_ASSERT(sizeof(struct ump_start_message) <= UMP_RING_BUF_DATA_SIZE, "struct ump_start_message is too big");

/**
 * @brief Initialize the send and receive objects for a channel.
 *
 * @param role Specifies if this side should be the server or client in the connection
 * @param send_chan send channel object to be allocated and initialized
 * @param recv_chan receive channel object to be allocated and initialized
 * @param recv_header_size size of the headers that will be received
 * @param shared_frame the frame capability shared between the client and server
 * @param run_on_waitset the waitset on which the async tasks should run. if it is NULL a thread will be started to handle the tasks
 * @return errval_t
 */
errval_t ump_chan_init(
    enum ump_role role,
    struct ump_send_chan **send_chan,
    struct ump_recv_chan **recv_chan,
    size_t recv_header_size,
    struct capref shared_frame,
    struct waitset *run_on_waitset
);

/**
 * @brief Send a message on the channel
 *
 * @param chan channel to send on
 * @param entry an allocated object to carry the send information
 * @param header_size size of the header to be sent
 * @param header buffer containing the header to be sent
 * @param payload_size size of the payload to be sent
 * @param payload buffer containing the payload to be sent
 * @param callback function to call after sending
 * @param callback_arg argument to pass to the callback function
 */
void ump_send (
    struct ump_send_chan *chan,
    struct ump_send_queue_entry *entry,
    size_t header_size,
    const void *header,
    size_t payload_size,
    const void *payload,
    ump_send_callback_fn_t callback,
    void *callback_arg
);

/**
 * @brief Receive the header of the next message. Must only be called on a fresh channel or after the receive payload completed
 *
 * @param chan channel to listen on
 * @param callback function to call with the header and the size of the payload
 * @param callback_arg argument to pass to the callback function
 */
void ump_recv_header (
    struct ump_recv_chan *chan,
    ump_recv_header_callback_fn_t callback,
    void *callback_arg
);

/**
 * @brief Receive the payload of the next message. Must only be called after the receive header completed
 *
 * @param chan channel to listen on
 * @param payload buffer to store the payload in. must be at least as large as the payload_size argument in the corresponding callback from receive header. if this is NULL then the payload will be dropped
 * @param callback function to call with payload
 * @param callback_arg argument to pass to the callback function
 */
void ump_recv_payload (
    struct ump_recv_chan *chan,
    void *payload,
    ump_recv_payload_callback_fn_t callback,
    void *callback_arg
);

#endif /* _INIT_UMP_CHAN_H_ */
