#include <aos/aos.h>
#include <aos/ump_chan.h>
#include <aos/paging.h>
#include <aos/waitset_chan.h>

/**
 * @brief initializes a ring buffer state
 * 
 * @param ring_state object to initialize
 * @param buf start of the ring buffer memory
 * @param buf_len size of the ring buffer memory
 */
static void ump_ring_init(struct ump_ring_state *ring_state, void *buf, size_t buf_len) {
    // the ring buffers should be cache-line aligned
    assert(((uintptr_t) buf) % CACHE_LINE_SIZE == 0);

    ring_state->ring_start = buf;
    ring_state->ring_entry_next = 0;
    ring_state->ring_entry_count = buf_len / sizeof(struct ump_ring_buf_entry);
}

/**
 * @brief Tries to write a ring buffer entry without blocking. Advances in the ring if successful
 * 
 * @param ring_state ring to write to
 * @param buf start of the buffer containing the data
 * @param bytes number of bytes from the buffer to write to the ring
 * @return true successfully wrote the entry
 * @return false failed to write the entry because the buffer is full
 */
static bool ump_ring_try_write_next(struct ump_ring_state *ring_state, const void *buf, size_t bytes) {
    assert(bytes <= UMP_RING_BUF_DATA_SIZE);
    assert(buf != NULL);

    // check if we own the entry
    if (ring_state->ring_start[ring_state->ring_entry_next].owner) {
        return false;
    }
    // memory barrier because the other endpoint might be on another core
    __asm volatile (
        "dmb sy\n"
    );

    memcpy(&ring_state->ring_start[ring_state->ring_entry_next].data, buf, bytes);
    
    // memory barrier because the other endpoint might be on another core
    __asm volatile (
        "dmb sy\n"
    );
    // mark the entry as owned by the receiver
    ring_state->ring_start[ring_state->ring_entry_next].owner = 1;
    // advance in the ring
    ring_state->ring_entry_next = (ring_state->ring_entry_next + 1) % ring_state->ring_entry_count;

    return true;
};

/**
 * @brief Tries to read an entry of the ring buffer. Advances the ring if successful
 * 
 * @param ring_state ring to read from
 * @param buf start of the buffer to write to. if NULL the data will be dropped
 * @param bytes number of bytes to read from this buffer (from the start of the buffer. the rest is dropped)
 * @return true successfully read an entry
 * @return false could not read an entry because no data was available
 */
static bool ump_ring_try_read_next(struct ump_ring_state *ring_state, void *buf, size_t bytes) {
    assert(bytes <= UMP_RING_BUF_DATA_SIZE);

    // check if we own the entry
    if(!ring_state->ring_start[ring_state->ring_entry_next].owner) {
        return false;
    }
    // memory barrier because the other endpoint might be on another core
    __asm volatile (
        "dmb sy\n"
    );

    // if we get a NULL buf then we need to drop the data
    if (buf != NULL) {
        memcpy(buf, &ring_state->ring_start[ring_state->ring_entry_next].data, bytes);
    }
    // memory barrier because the other endpoint might be on another core
    __asm volatile (
        "dmb sy\n"
    );
    // mark the entry as owned by the sender
    ring_state->ring_start[ring_state->ring_entry_next].owner = 0;
    // advance in the ring
    ring_state->ring_entry_next = (ring_state->ring_entry_next + 1) % ring_state->ring_entry_count;
    
    return true;
}

/**
 * @brief Try to read the next message. The message might also be read only partially but it will continue in the next call to this function
 * 
 * @param recv_chan channel to read the message on
 */
static void ump_try_read(struct ump_recv_chan *recv_chan) {
    // make sure we own the read metadata, otherwise we have nothing to do anyways
    thread_mutex_lock(&recv_chan->recv_register_lock);
    if(
        recv_chan->state != UMP_RECV_STATE_START_MESSAGE && 
        recv_chan->state != UMP_RECV_STATE_HEADER && 
        recv_chan->state != UMP_RECV_STATE_PAYLOAD
    ) {
        thread_mutex_unlock(&recv_chan->recv_register_lock);
        return;
    }
    thread_mutex_unlock(&recv_chan->recv_register_lock);

    // first make sure we have read the start message
    if (recv_chan->state == UMP_RECV_STATE_START_MESSAGE) {
        struct ump_start_message start_message;
        if(!ump_ring_try_read_next(&recv_chan->ring_state, &start_message, sizeof(struct ump_start_message))) {
            return;
        }
        // the sender and receiver MUST agree on the header size
        assert(recv_chan->header_size == start_message.header_size);

        // update the receive state to start reading the header
        recv_chan->read_bytes_left   = recv_chan->header_size;
        recv_chan->read_buf_position = recv_chan->header;
        recv_chan->next_payload_size = start_message.payload_size;
        recv_chan->state = UMP_RECV_STATE_HEADER;
    }

    // read as much as popssible of the header or payload
    while(recv_chan->read_bytes_left > 0) {
        size_t bytes_to_read = MIN(recv_chan->read_bytes_left, UMP_RING_BUF_DATA_SIZE);
        if(!ump_ring_try_read_next(&recv_chan->ring_state, recv_chan->read_buf_position, bytes_to_read)) {
            // buffer is empty, we need to try again later
            return;
        }
        if(recv_chan->read_buf_position != NULL) {
            recv_chan->read_buf_position += bytes_to_read;
        }
        recv_chan->read_bytes_left -= bytes_to_read;
    }

    // we managed to read the current stage, advance the state and call the callbacks
    if (recv_chan->state == UMP_RECV_STATE_HEADER) {
        recv_chan->state = UMP_RECV_STATE_PAYLOAD_IDLE;
        recv_chan->recv_header_callback(
            recv_chan->callback_arg,
            recv_chan->header_size,
            recv_chan->header,
            recv_chan->next_payload_size
        );
    } else if (recv_chan->state == UMP_RECV_STATE_PAYLOAD) {
        recv_chan->state = UMP_RECV_STATE_IDLE;
        recv_chan->recv_payload_callback(
            recv_chan->callback_arg,
            recv_chan->next_payload_size,
            recv_chan->payload
        );
    } else {
        // cannot happen but just to be sure...
        assert(false);
    }
}

/**
 * @brief Try to send the next message. The message might also be sent only partially but it will continue in the next call to this function
 * 
 * @param send_chan channel to send the message on
 */
static void ump_try_send(struct ump_send_chan *send_chan) {
    // check if there is anything to do.
    thread_mutex_lock(&send_chan->send_queue_lock);
    if (send_chan->send_queue_head == NULL) {
        thread_mutex_unlock(&send_chan->send_queue_lock);
        return;
    }
    thread_mutex_unlock(&send_chan->send_queue_lock);

    // get the entry to handle next
    struct ump_send_queue_entry *entry = send_chan->send_queue_head;

    // first we need to send the start message so the receiver knows how much payload to expect
    if (entry->state == UMP_SEND_ENTRY_STATE_START_MESSAGE) {
        struct ump_start_message start_message = {
            .header_size = entry->header_size,
            .payload_size = entry->payload_size,
        };
        if(!ump_ring_try_write_next(&send_chan->ring_state, &start_message, sizeof(struct ump_start_message))) {
            // ring is full, we need to retry later
            return;
        }
        entry->state = UMP_SEND_ENTRY_STATE_HEADER;
        entry->send_buf_position = entry->header;
        entry->send_bytes_left = entry->header_size;
    }

    // send as much in the current stage as possible
    while(entry->send_bytes_left > 0) {
        size_t bytes_to_send = MIN(entry->send_bytes_left, UMP_RING_BUF_DATA_SIZE);
        if(!ump_ring_try_write_next(&send_chan->ring_state, entry->send_buf_position, bytes_to_send)) {
            return;
        }
        entry->send_buf_position += bytes_to_send;
        entry->send_bytes_left -= bytes_to_send;
    }

    // we managed to send the current stage, advance the state!
    if (entry->state == UMP_SEND_ENTRY_STATE_HEADER) {
        entry->send_buf_position = entry->payload;
        entry->send_bytes_left = entry->payload_size;
        entry->state = UMP_SEND_ENTRY_STATE_PAYLOAD;
    } else if (entry->state == UMP_SEND_ENTRY_STATE_PAYLOAD) {
        // we are done sending, remove the entry from the queue
        thread_mutex_lock(&send_chan->send_queue_lock);
        // we take the current send queue state to get "next" because it might have changed since
        // we read the entry (NULL -> something)
        if (send_chan->send_queue_head->next == NULL) {
            send_chan->send_queue_tail = NULL;
        }
        send_chan->send_queue_head = send_chan->send_queue_head->next;
        thread_mutex_unlock(&send_chan->send_queue_lock);

        // send the callback
        assert(entry->callback != NULL);
        entry->callback(entry->callback_arg, entry);
    }
}

/**
 * @brief Contains references to information required by the ump channel worker
 */
struct ump_worker_context {
    struct waitset *run_on_waitset;
    struct waitset_chanstate waitset_chan;
    struct ump_recv_chan *recv_chan;
    struct ump_send_chan *send_chan;
};

/**
 * @brief Do the async work on the channel
 * 
 * @param arg 
 */
static void ump_work(void *arg) {
    struct ump_worker_context *worker_context = arg;

    // check the recv queue
    ump_try_read(worker_context->recv_chan);

    // check the send queue
    ump_try_send(worker_context->send_chan);
}

/**
 * @brief used to handle the async work on a separate thread
 * 
 * @param arg the worker context
 * @return int not used
 */
static int ump_thread_worker(void *arg) {

    while (true) {

        ump_work(arg);
        
        thread_yield();
    }

    return 0;
}

static void ump_waitset_worker(void *arg);
/**
 * @brief Schedule the worker on its waitset.
 * 
 * @param worker_context 
 */
static void ump_worker_schedule(struct ump_worker_context *worker_context) {
    // we want to be automatically re-registered upon completion of an event
    // worker_context->waitset_chan.persistent = true;
    errval_t err = waitset_chan_trigger_closure(
        worker_context->run_on_waitset, 
        &worker_context->waitset_chan,
        MKCLOSURE(ump_waitset_worker, worker_context)
    );
    // this is supposed to be called single threaded and only upon event trigger
    // hence this cannot be a reregister and thus not fail
    assert(err_is_ok(err));
}

/**
 * @brief used to handle the async work on a waitset
 * 
 * @param arg the worker context
 */
static void ump_waitset_worker(void *arg) {
    ump_work(arg);

    // put back in the queue
    ump_worker_schedule(arg);
}

errval_t ump_chan_init(
    enum ump_role role,
    struct ump_send_chan **send_chan_ret,
    struct ump_recv_chan **recv_chan_ret,
    size_t recv_header_size,
    struct capref shared_frame,
    struct waitset *run_on_waitset
) {
    errval_t err;

    // check that we have a valid role for the channel
    if (role >= UMP_ROLE_COUNT) return ERR_INVALID_ARGS;

    // identify the shared frame to determine its size
    struct frame_identity shared_frame_id;
    err = frame_identify(shared_frame, &shared_frame_id);
    if (err_is_fail(err)) return LIB_ERR_FRAME_IDENTIFY;
    
    // map the full frame
    void *shared_mem;
    err = paging_map_frame(get_current_paging_state(), &shared_mem, shared_frame_id.bytes, shared_frame);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);

    // get the size of the rings for each direction
    size_t ring_size = shared_frame_id.bytes / 2;
    debug_printf("[ump_chan_init] Each channel is %lu bytes\n", ring_size);

    // allocate the channel structures
    struct ump_send_chan *send_chan = malloc(sizeof(struct ump_send_chan));
    if(send_chan == NULL) {
        paging_unmap(get_current_paging_state(), shared_mem);
        return LIB_ERR_MALLOC_FAIL;
    }
    struct ump_recv_chan *recv_chan = malloc(sizeof(struct ump_recv_chan));
    if(recv_chan == NULL) {
        paging_unmap(get_current_paging_state(), shared_mem);
        free(send_chan);
        return LIB_ERR_MALLOC_FAIL;
    }

    // decide on which part of the buffer is used for sending and receiving
    // give the server receiving side the first half of the frame
    void *receive_buf = role == UMP_ROLE_SERVER ? shared_mem : shared_mem + ring_size;
    // give the server sending side the second half of the frame
    void *send_buf = role == UMP_ROLE_SERVER ? shared_mem + ring_size : shared_mem;


    // initialize recv state
    recv_chan->header = malloc(recv_header_size);
    if (recv_chan->header == NULL) {
        paging_unmap(get_current_paging_state(), shared_mem);
        free(send_chan);
        free(recv_chan);
        return LIB_ERR_MALLOC_FAIL;
    }
    recv_chan->header_size = recv_header_size;
    recv_chan->state = UMP_RECV_STATE_IDLE;
    ump_ring_init(&recv_chan->ring_state, receive_buf, ring_size);
    thread_mutex_init(&recv_chan->recv_register_lock);

    // initialize send state
    send_chan->send_queue_head = NULL;
    send_chan->send_queue_tail = NULL;
    ump_ring_init(&send_chan->ring_state, send_buf, ring_size);
    thread_mutex_init(&send_chan->send_queue_lock);

    // allocate and initialize the worker state for handling the async work
    struct ump_worker_context *worker_context = malloc(sizeof(struct ump_worker_context));
    if(worker_context == NULL) {
        paging_unmap(get_current_paging_state(), shared_mem);
        free(send_chan);
        free(recv_chan);
        free(recv_chan->header);
        return LIB_ERR_MALLOC_FAIL;
    }
    worker_context->run_on_waitset = run_on_waitset;
    worker_context->recv_chan = recv_chan;
    worker_context->send_chan = send_chan;

    
    if (run_on_waitset != NULL) {
        // we need to run on a waitset so initialize it and schedule the worker on it
        waitset_chanstate_init(&worker_context->waitset_chan, CHANTYPE_OTHER);
        ump_worker_schedule(worker_context);
    } else {
        // we need to run on a dedicated thread so create it
        struct thread *t = thread_create(ump_thread_worker, worker_context);
        if (t == NULL) {
            paging_unmap(get_current_paging_state(), shared_mem);
            free(send_chan);
            free(recv_chan);
            free(recv_chan->header);
            free(worker_context);
            return LIB_ERR_THREAD_CREATE;
        }
    }

    // return the channel metadata
    *send_chan_ret = send_chan;
    *recv_chan_ret = recv_chan;

    return SYS_ERR_OK;
}

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
) {
    assert(chan != NULL);
    
    // initialize the queue entry
    entry->state = UMP_SEND_ENTRY_STATE_START_MESSAGE;
    entry->header_size = header_size;
    entry->header = header;
    entry->payload_size = payload_size;
    entry->payload = payload;
    entry->callback = callback;
    entry->callback_arg = callback_arg;
    entry->next = NULL;

    // add the new entry to the end of the send queue
    thread_mutex_lock(&chan->send_queue_lock);
    if(chan->send_queue_tail != NULL) {
        chan->send_queue_tail->next = entry;
    } else {
        chan->send_queue_head = entry;
    }
    chan->send_queue_tail = entry;
    thread_mutex_unlock(&chan->send_queue_lock);
}

void ump_recv_header (
    struct ump_recv_chan *chan,
    ump_recv_header_callback_fn_t callback,
    void *callback_arg
) {
    assert(chan != NULL);

    thread_mutex_lock(&chan->recv_register_lock);
    // NOTE rueegges: caller has to ensure it only registers at most once?! Or should we just make it a noop?
    assert(chan->state == UMP_RECV_STATE_IDLE);

    chan->recv_header_callback = callback;
    chan->callback_arg = callback_arg;

    // this must be last so we have written everything when the worker thread sees this state
    chan->state = UMP_RECV_STATE_START_MESSAGE;

    thread_mutex_unlock(&chan->recv_register_lock);
}

void ump_recv_payload (
    struct ump_recv_chan *chan,
    void *payload,
    ump_recv_payload_callback_fn_t callback,
    void *callback_arg
) {
    assert(chan != NULL);
    
    thread_mutex_lock(&chan->recv_register_lock);
    // NOTE rueegges: caller has to ensure it only registers at most once?! Or should we just make it a noop?
    assert(chan->state == UMP_RECV_STATE_PAYLOAD_IDLE);

    chan->recv_payload_callback = callback;
    chan->callback_arg = callback_arg;

    chan->read_bytes_left   = chan->next_payload_size;
    chan->read_buf_position = payload;
    chan->payload = payload;

    // this must be last so we have written everything when the worker thread sees this state
    chan->state = UMP_RECV_STATE_PAYLOAD;

    thread_mutex_unlock(&chan->recv_register_lock);
}