#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[FD_SETSIZE]; // 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
    int recalc_max; 
    int max_fd; //keeps track of the max fd for select()
};

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


static void iomux_select_recalc_max(struct IOMuxSelect *mux)
{
    size_t i;
    size_t n;
    mux->max_fd = -1;

    for (i = 0, n = 0; i < FD_SETSIZE && n < mux->num_descriptors; i++) {
        if (mux->descriptors[i] != NULL) {
            mux->max_fd = UC_MAX(mux->descriptors[i]->fd, mux->max_fd);
            n++;
        }
    }
}

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

    if (fd->fd < 0 || fd->fd >= FD_SETSIZE) {
        return EINVAL;
    }

    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;
    int rc;

    if (fd->callback == NULL)
        return EINVAL;

    rc = iomux_select_update_events(mux_, fd);
    if (rc != 0) {
        return ENFILE;
    }

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


    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;

    if (fd->fd < 0 || fd->fd >= FD_SETSIZE) {
        return ENFILE;
    }

    if (mux->descriptors[fd->fd] != NULL) {
        unsigned int what = fd->what; //make sure we don't alter 'what' as seen from the user

        fd->what = 0;
        iomux_select_update_events(mux_, fd);
        mux->descriptors[fd->fd] = NULL; 
        mux->recalc_max = 1;
        fd->what = what;

        rc = 0;
        mux->num_descriptors--;
    }

    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->recalc_max) {
        iomux_select_recalc_max(mux);
        mux->recalc_max = 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 < FD_SETSIZE; i++) {
        int had_event = 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);
            had_event = 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);
            had_event = 1;
        }
        
        event_cnt += had_event;
        
    }

    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);
        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;
}