#include <sys/select.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <assert.h>
#include "ucore/utils.h"
#include "iomux_impl.h"

#define IOMUX_GROW_CHUNK (32)

struct IOMuxSelect {
    fd_set master_read_set;
    fd_set master_write_set;

    struct IOMuxFD **descriptors; // array of pointers to the user IOMuxFD's
    size_t alloc_descriptors; //allocated descriptors in 'descriptors'
    size_t num_descriptors; //number of descriptors we're watching
    size_t deleted_cnt; //used to rebuild the descriptors array
    int max_fd; //keeps track of the max fd for select()
};

static int iomux_select_update_events(struct IOMux *mux_, struct IOMuxFD *fd);

//grows the descriptor array. We never shrink it.
static int iomux_select_grow(struct IOMuxSelect *mux)
{
    void *tmp = realloc(mux->descriptors, (mux->num_descriptors + IOMUX_GROW_CHUNK) * sizeof *mux->descriptors);
    assert(tmp != NULL);
    if (tmp == NULL)
        return ENOMEM;

    mux->alloc_descriptors += IOMUX_GROW_CHUNK;
    mux->descriptors = tmp;

    return 0;
}

static void iomux_select_remove(struct IOMuxSelect *mux, size_t idx)
{
    if (mux->num_descriptors > 1) {
        //move last element into the empty slot
        size_t last_idx = mux->num_descriptors - 1;
        mux->descriptors[idx] = mux->descriptors[last_idx];
        mux->descriptors[last_idx] = NULL;
    }

    mux->num_descriptors--;
}

static void iomux_select_rebuild(struct IOMuxSelect *mux)
{
    size_t i;

    mux->max_fd = -1;
    for (i = 0; i < mux->num_descriptors; ) {
        if (mux->descriptors[i] == NULL) {
            iomux_select_remove(mux, i);
            //a deleted fd slot needs to be rechecked, no i++ here
        } else {
            mux->max_fd = UC_MAX(mux->descriptors[i]->fd, mux->max_fd);
            i++;
        }
    }
}

//sets or clears the events.
static int iomux_select_update_events(struct IOMux *mux_, struct IOMuxFD *fd)
{
    struct IOMuxSelect *mux = mux_->instance;

    if (fd->what & MUX_EV_READ)
        FD_SET(fd->fd, &mux->master_read_set);
    else 
        FD_CLR(fd->fd, &mux->master_read_set);

    if (fd->what & MUX_EV_WRITE)
        FD_SET(fd->fd, &mux->master_write_set);
    else 
        FD_CLR(fd->fd, &mux->master_write_set);

    return 0;
}

static int iomux_select_register_fd(struct IOMux *mux_, struct IOMuxFD *fd)
{
    struct IOMuxSelect *mux = mux_->instance;
    if (fd->what == 0 || fd->callback == NULL)
        return EINVAL;

    if (mux->num_descriptors == mux->alloc_descriptors) {
        if (iomux_select_grow(mux) != 0)
            return ENOMEM;
    }

    mux->descriptors[mux->num_descriptors] = fd;
    mux->num_descriptors++;

    iomux_select_update_events(mux_, fd);

    mux->max_fd = UC_MAX(fd->fd, mux->max_fd);

    return 0;
}

static int iomux_select_unregister_fd(struct IOMux *mux_, struct IOMuxFD *fd)
{
    struct IOMuxSelect *mux = mux_->instance;
    int rc = ENOENT;

    size_t i;
    for (i = 0; i < mux->num_descriptors; i++) {
        if (fd == mux->descriptors[i]) {
            unsigned int what = fd->what; //make sure we don't alter 'what' as seen from the user
            fd->what &= ~MUX_EV_MASK;
            iomux_select_update_events(mux_, fd);
            mux->descriptors[i] = NULL; //set the slot to NULL, the event loop needs to rebuild it.
            mux->deleted_cnt++;
            fd->what = what;
            rc = 0;
            break;
        }
    }

    return rc;
}

static int iomux_select_run(struct IOMux *mux_, struct timeval *timeout)
{
    struct IOMuxSelect *mux = mux_->instance;

    int rc;
    size_t i;
    int event_cnt = 0;
    fd_set read_set;
    fd_set write_set;

    if (mux->deleted_cnt) {
        iomux_select_rebuild(mux);
        mux->deleted_cnt = 0;
    }

    if (mux->max_fd == -1 && timeout == NULL)
        return 0;

    do {

        read_set = mux->master_read_set;
        write_set = mux->master_write_set;
        rc = select(mux->max_fd + 1, &read_set, &write_set, NULL, timeout); //we can use timeout directly

    } while (rc == -1 && errno == EINTR);

    assert(rc >= 0);
    if (rc < 0)
        return -1;

    event_cnt = iomux_timers_run(mux_); //fire the timers

    if (rc == 0) //no fd's were ready, no more work to do
        return event_cnt;
    
    for (i = 0; rc > 0 && i < mux->num_descriptors; i++) {
        unsigned int events = 0;

        //deliver only one event at a time , to simplify application
        //programming
        if (mux->descriptors[i] != NULL &&
                FD_ISSET(mux->descriptors[i]->fd, &read_set)) {
            rc--;
            assert(mux->descriptors[i]->callback != NULL);
            mux->descriptors[i]->callback(mux_, mux->descriptors[i], MUX_EV_READ);
            events = 1;
        }

        //need to re-check in case the callback deleted the IOMuxFD
        if (mux->descriptors[i] != NULL &&
                FD_ISSET(mux->descriptors[i]->fd, &write_set)) {
            rc--;
            assert(mux->descriptors[i]->callback != NULL);
            mux->descriptors[i]->callback(mux_, mux->descriptors[i], MUX_EV_WRITE);
            events = 1;
        }
        
        event_cnt += events;
    }

    return event_cnt;
}

static void iomux_select_delete(struct IOMux *mux)
{
    if (mux != NULL && mux->instance != NULL) {
        struct IOMuxSelect *mux_select = mux->instance;
        free(mux_select->descriptors);
        mux_select->descriptors = NULL;
        free(mux_select);
        mux->instance = NULL;
    }
}

static const struct IOMuxOps select_ops = {
    .run_impl           = iomux_select_run,
    .delete_impl        = iomux_select_delete,
    .register_fd_impl   = iomux_select_register_fd,
    .unregister_fd_impl = iomux_select_unregister_fd,
    .update_events_impl = iomux_select_update_events,
};
    
int iomux_select_init(struct IOMux *mux)
{
    struct IOMuxSelect *mux_select = calloc(1, sizeof *mux_select);
    if (mux_select == NULL)
        return ENOMEM;

    mux_select->max_fd = -1;

    mux->instance = mux_select;
    mux->ops = select_ops;

    return 0;
}