#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <assert.h>

#include "common.h"
#include "render_struct.h"
#include "render_multi.h"
#include "render_remote.h"
#include "remote_node.h"
#include "remote_pool.h"
#include "render_slices_struct.h"
#include "render_slices.h"

/*
// the states we can go through
enum render_multi_state {
    STATE_INIT,                 // the render_multi is in the process of being initialized
    STATE_SENDING,              // we're still waiting for some of the requests to be sent
    STATE_RENDER,               // we're handling data now!
    STATE_DATA_DONE,            // we're finished with all the data
    STATE_PNG_DONE,             // PNG data's done
    STATE_FAILED,               // we failed
};

enum render_multi_node_state {
    STATE_INIT,                 // the render_multi_node is in the process of being initialized
    STATE_SENDING,              // we're waiting for the requests to be sent
    STATE_FILL_ROW,             // we're filling the row with data
    STATE_ROW_FULL,             // our row is full
    STATE_DATA_DONE,            // we're finished with all the data
    STATE_PNG_DONE,             // PNG data's done
    STATE_FAILED,               // we failed

};
*/

struct render_multi {
    // these are used as arguments to render_remote
    struct render_multi_node {
        // the slice info
        struct render_slice_info *info;

        // the render_remote_ctx
        struct render_remote *render_remote;
        
        // _render_multi_sent called for this?
        int render_sent;

        // how wide our slice is, static
        size_t slice_width;

        // how many bytes we have already written into the current row
        size_t col;

        // passed in the last segment of data?
        int render_done;
        
        // a pointer to ourself
        struct render_multi *self;

    } nodes[RENDER_SLICES_MAX];
    
    // how many nodes we have
    int node_count;

    // is this still alive?
    int alive;

    // how many remote renders have been succesfully cb_sent?
    int renders_sent;

    // how many of the renders are done?
    int renders_done;

    // have we called cb_sent?
    int have_sent;

    // the render_slices thing
    struct render_slices slices;

    // has render_slices_done returned?
    int slices_done;

    // buffer render_png output in this
    struct evbuffer *out_buf;

    // our own callbacks that we call
    void (*cb_sent)(void *arg);
    void (*cb_data)(struct evbuffer *buf, void *arg);
    void (*cb_done)(void *arg);
    void (*cb_fail)(void *arg);

    void *cb_arg;
};

#define FAIL_PARTIAL 0x01
#define FAIL_SILENT 0x02

// prototypes
static void _render_multi_do_free (struct render_multi *ctx);
static void _render_multi_do_sent (struct render_multi *ctx);
static void _render_multi_do_png_data (struct render_multi *ctx);
static void _render_multi_do_png_done (struct render_multi *ctx);
static void _render_multi_do_done (struct render_multi *ctx);
static void _render_multi_do_fail (struct render_multi *ctx, int flags);

/*
 * Actually free the request. Should be de-initialized (either _render_multi_error, or _render_done when this is called
 */
static void _render_multi_do_free (struct render_multi *ctx) {
    assert(ctx && ctx->alive == 0);
    
    render_slices_deinit(&ctx->slices);

    if (ctx->out_buf) {
        evbuffer_free(ctx->out_buf);
        ctx->out_buf = NULL;
    }

    free(ctx);
}

// the request has been sent
static void _render_multi_do_sent (struct render_multi *ctx) {
    int i;

    // check that all the nodes have indeed been sent
    for (i = 0; i < ctx->node_count; i++) {
        assert(ctx->nodes[i].render_sent);
    }

    // assert the callbacks are still valid
    assert(ctx->cb_sent && ctx->cb_fail && ctx->cb_done);

    // call cb_sent and then invalidate it
    ctx->cb_sent(ctx->cb_arg);
    ctx->cb_sent = NULL;
    ctx->have_sent = 1;
    
    // we're going to always have the PNG header data buffered at this point, so give that to the user right away
    assert(evbuffer_get_length(ctx->out_buf) > 0);
    _render_multi_do_png_data(ctx);
}

// possibly call cb_data, and if the renders are all done and the buffer is empty, cb_done
static void _render_multi_do_png_data (struct render_multi *ctx) {
    // at first we have to wait until we've called cb_sent
    if (!ctx->have_sent)
        return;

    // got any PNG data in there?
    if (evbuffer_get_length(ctx->out_buf))
        ctx->cb_data(ctx->out_buf, ctx->cb_arg);
    
    // was that the last piece of PNG data?
    if (ctx->slices_done && evbuffer_get_length(ctx->out_buf) == 0) {
        // PNG data done!
        _render_multi_do_png_done(ctx);
    }
}

// the PNG rendering completed succesfully
static void _render_multi_do_png_done (struct render_multi *ctx) {
    // check that ctx is still valid
    assert(ctx->alive);

    // mark as not alive
    ctx->alive = 0;

    // call cb_done
    ctx->cb_done(ctx->cb_arg);

    // don't free ourself, our user does that (probably did already)
    return;
}

// the request completed normally, flush the png data and return
// _render_multi_png_done takes care of calling cb_done, not us!
static void _render_multi_do_done (struct render_multi *ctx) {
    assert(ctx->alive);

    int i;
    
    // check that all the remote_renders are indeed complete
    for (i = 0; i < ctx->node_count; i++) {
        assert(ctx->nodes[i].render_remote == NULL);
        assert(ctx->nodes[i].col == 0);
    }
    
    // finish off the render_slices
    if (render_slices_done(&ctx->slices))
        ERROR("render_slices_done");
    
    // mark this as complete, all data is now in the out buffer
    ctx->slices_done = 1;

    // if that all the data handled now, we're done
    _render_multi_do_png_data(ctx);

    // don't free ourself, our user does that (probably already did, via render_png_done)
    return;

error:    
    /* render_slices_done -> render_png_done failed, probably because we didn't have enough data */
    _render_multi_do_fail(ctx, FAIL_PARTIAL);
}

// the request completed abnormally. Flags:
//  FAIL_SILENT     - don't call cb_fail
//  FAIL_PARTIAL    - assume png_info may be NULL
static void _render_multi_do_fail (struct render_multi *ctx, int flags) {
    int i;

    // check that ctx is still valid
    assert(ctx->alive || flags & FAIL_PARTIAL);

    // mark as not alive
    ctx->alive = 0;

    // cancel any in-progress remote renders
    for (i = 0; i < ctx->node_count; i++)
        if (ctx->nodes[i].render_remote) {
            render_remote_cancel(ctx->nodes[i].render_remote);
            ctx->nodes[i].render_remote = NULL;
        }
    
    // XXX: does render_slices need an abort?
    
    // check that both callbacks are still valid
    assert(ctx->cb_fail && ctx->cb_done);
    
    if (!(flags & FAIL_SILENT)) {
        // call cb_fail and then invalidate it
        ctx->cb_fail(ctx->cb_arg);
    }

    ctx->cb_fail = NULL;

    // don't free ourself, our user does that
}

/*
 * One of the remote render commands has succesfully been sent.
 *
 * Once all of these commands have been sent, invoke our cb_sent.
 */
static void _render_multi_sent (void *arg) {
    struct render_multi_node *ctx = arg;
    
    // mark these as sent
    ctx->render_sent = 1;
    ctx->self->renders_sent++;

    // have all render_sub_ctxs been sent?
    if (ctx->self->renders_sent == ctx->self->node_count) {
        // tell our user
        _render_multi_do_sent(ctx->self);
    }
}

/*
 * One render node failed, abort the whole thing
 */
static void _render_multi_fail (void *arg) {
    struct render_multi_node *ctx = arg;

    // free this ctx's remote render
    render_remote_free(ctx->render_remote);
    ctx->render_remote = NULL;
    
    // cancel the rest + cb_fail
    _render_multi_do_fail(ctx->self, 0);
}

/*
 * Got new data for some remote render
 */
static void _render_multi_data_raw (int fd, short event, void *arg) {
    struct render_multi_node *ctx = arg;
    int ret;

    assert(ctx->col <= ctx->slice_width);   // check it isn't out of range

    // if our slice is full, we don't want to receive any more data
    if (ctx->col == ctx->slice_width)
        return;

    // read new data into our slice
    ret = read(fd, 
        ctx->info->render_buf + ctx->col,               // our segment buffer + partial segment offset
        ctx->slice_width - ctx->col                     // how many bytes left in the segment
    );
    
    // errors/EOF?
    if (ret == -1) {
        if (errno == EAGAIN) {
            // false alarm
            goto reschedule;

        } else
            ERROR("read");

    } else if (ret == 0) {
        // this ctx's remote render is done
        render_remote_done(ctx->render_remote);
        ctx->render_remote = NULL;

        // count how many are done
        ctx->self->renders_done++;

        // are all of them done?
        if (ctx->self->renders_done == ctx->self->node_count) {
            // finish it off
            _render_multi_do_done(ctx->self);

        } // else, just wait for the rest to complete
        
        // do *NOT* reschedule ourself, ctx->render_remote is invalid anyways (as is ctx!)
        return;
    }

    // ok, we received some data normally
    ctx->col += ret;

    // is this segment full now?
    if (ctx->col == ctx->slice_width) {
        int status;
        
        // pass the segment in to render_slices
        if ((status = render_slices_segment_done(&ctx->self->slices, ctx->info->index)) == -1)
            ERROR("render_slices_segment_done");

        // reset the col marker
        ctx->col = 0;

        // row done?
        if (status & SLICE_PROCESS_ROW) {
            // process the row via render_slices
            status = render_slices_process_row(&ctx->self->slices);
            
        }

        // do we need to continue?
        if (status & SLICE_CONTINUE) {
            // reschedule the reads in case they were ~SLICE_CONTINUE
            int i;

            for (i = 0; i < ctx->self->node_count; i++) {
                // just don't reschedule those that are already EOF...
                if (ctx->self->nodes[i].render_remote)
                    render_remote_reschedule(ctx->self->nodes[i].render_remote);
            }

        } else {
            // don't read any more data yet, we'll be rescheduled by someone else
            return;

        }
    }

    // ok, reschedule ourselves

reschedule:
    // reschedule a new call once we get more data
    render_remote_reschedule(ctx->render_remote);

    return;

error:
    _render_multi_do_fail(ctx->self, 0);
}

/*
 * We fed a row of pixels into render_png, and this PNG data came out.
 *
 * We need to pass it back to our caller
 */
static int _render_multi_png_data (const unsigned char *data, size_t length, void *arg) {
    struct render_multi *ctx = arg;

    // XXX: need a better user-API to avoid these data copies
    if (evbuffer_add(ctx->out_buf, data, length))
        ERROR("evbuffer_add");
    
    // handle cb_data/cb_done
    _render_multi_do_png_data(ctx);
    
    // ok
    return 0;

error:
    // don't do any failing here, this will return control to a _render_* function that will handle it
    return -1;
}

struct render_multi *render_multi (
        struct render *render,               // what to render
        struct remote_pool *pool_info,    // what render pool to use
        void (*cb_sent)(void *arg),
        void (*cb_data)(struct evbuffer *buf, void *arg),
        void (*cb_done)(void *arg),
        void (*cb_fail)(void *arg),
        void *cb_arg
) {
    struct render_multi *ctx = NULL;
    
    // alloc the render_multi
    ctx = calloc(1, sizeof(struct render_multi));

    if (!ctx)
        ERROR("calloc");

    // store the provided callback functions
    ctx->cb_sent = cb_sent;
    ctx->cb_data = cb_data;
    ctx->cb_done = cb_done;
    ctx->cb_fail = cb_fail;
    ctx->cb_arg = cb_arg;

    // store our render_png callbacks, must be before png_info
    if (render_io_custom(render, &_render_multi_png_data, NULL, ctx))
        ERROR("render_io_custom");

    // evbuffer, must be before png_info
    if (!(ctx->out_buf = evbuffer_new()))
        ERROR("evbuffer_new");

    // then initialize the render_slices
    if (render_slices_init(&ctx->slices, render))
        ERROR("render_slices_init");
    
    // how many nodes?
    ctx->node_count = render_slices_get_count(&ctx->slices);

    // load them
    int i;

    for (i = 0; i < ctx->node_count; i++) {
        // store the info struct
        ctx->nodes[i].info = render_slices_get_slice_info(&ctx->slices, i);

        // some simple attributes
        ctx->nodes[i].slice_width = ctx->nodes[i].info->render_info->img_w;
        ctx->nodes[i].self = ctx;

       
        // the render_remote
        if (!(ctx->nodes[i].render_remote = render_remote_rawio(ctx->nodes[i].info->render_info, pool_info, 
                &_render_multi_sent, &_render_multi_fail, &_render_multi_data_raw, &ctx->nodes[i]))
        )
            ERROR("render_remote_rawio");
    }
    
    // we are now alive
    ctx->alive = 1;
    
    // I guess that's a succesfull start now
    return ctx;

error:
    _render_multi_do_fail(ctx, FAIL_SILENT | FAIL_PARTIAL);

    return NULL;
}


void render_multi_set_recv (struct render_multi *ctx, size_t recv_threshold, size_t unread_buffer) {

}

int render_multi_flush (struct render_multi *ctx) {
    _render_multi_do_png_data(ctx);

    return 0;
}

void render_multi_cancel (struct render_multi *ctx) {
    _render_multi_do_fail(ctx, FAIL_SILENT);
}

void render_multi_free (struct render_multi *ctx) {
    _render_multi_do_free(ctx);
}