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

#include "common.h"
#include "render_internal.h"
#include "render_multi.h"
#include "render_remote.h"
#include "remote_node.h"
#include "render_png.h"

struct render_multi {
    // these are used as arguments to render_remote
    struct render_multi_sub {
        // our offset in the list
        int index;

        // the remote_render_ctx
        struct render_remote *remote_render;
        
        // a pointer to ourself
        struct render_multi *self;

        // if the render connection is marked as done without us having read all the data, move it from render_remote's buffers into here
        struct evbuffer *in_buf;

        // _render_sent called for this?
        int render_sent;

        // our offset into the row, static
        size_t row_offset;

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

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

        // is this render op done?
        int render_done;
    } remote_renders[RENDER_MULTI_NODES_MAX];
    
    // how many remote_renders we have
    int remote_render_count;

    // have we called cb_sent?
    int have_sent;

    // the png thing
    struct render_png *png_info;

    // our pixel data row
    unsigned char *rowbuf;

    // 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_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) {
    if (!ctx)
        return;
    
    int i;
    for (i = 0; i < ctx->remote_render_count; i++) {
        if (ctx->remote_renders[i].in_buf) {
            evbuffer_free(ctx->remote_renders[i].in_buf);
            ctx->remote_renders[i].in_buf = NULL;
        }
    }

    if (ctx->rowbuf) {
        free(ctx->rowbuf);
        ctx->rowbuf = NULL;
    }

    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 remote_renders have indeed been sent
    for (i = 0; i < ctx->remote_render_count; i++) {
        assert(ctx->remote_renders[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);

    ctx->cb_data(ctx->out_buf, ctx->cb_arg);

    // ok
    return;
}

// the request compelted normally
static void _render_multi_do_done (struct render_multi *ctx) {
    int i;

    // check that all the remote_renders are indeed complete
    for (i = 0; i < ctx->remote_render_count; i++) {
        assert(ctx->remote_renders[i].remote_render == NULL);
        assert(ctx->remote_renders[i].col == 0);
    }

    // finish the png_info
    if (render_png_done(ctx->png_info)) {
        // don't free it twice, though...
        ctx->png_info = NULL;

        ERROR("render_png_done");
    }

    ctx->png_info = NULL;

    // check that both callbacks are still valid
    assert(ctx->cb_fail && ctx->cb_done);

    // call cb_done and then invalidate it
    ctx->cb_done(ctx->cb_arg);
    ctx->cb_fail = NULL;

    // free ourself
    _render_multi_do_free(ctx);
    
    // ok
    return;

error:    
    /* O_o */
    _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;

    // cancel any in-progress remote renders
    for (i = 0; i < ctx->remote_render_count; i++)
        if (ctx->remote_renders[i].remote_render) {
            render_remote_cancel(ctx->remote_renders[i].remote_render);
            ctx->remote_renders[i].remote_render = NULL;
        }
    
    if (!(flags & FAIL_PARTIAL) || ctx->png_info) {
        // abort the render_png
        render_png_abort(ctx->png_info);
        ctx->png_info = NULL;
    }
    
    // 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;

    // free ourselves
    _render_multi_do_free(ctx);
}

/*
 * 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_sub *ctx = arg;

    ctx->render_sent = 1;

    // have all render_sub_ctxs been sent?
    int i;
    for (i = 0; i < ctx->self->remote_render_count; i++) {
        if (!ctx->self->remote_renders[i].render_sent)
            break;
    }
    
    // did we loop through all of them?
    if (i == ctx->self->remote_render_count) {
        // tell our user
        _render_multi_do_sent(ctx->self);
    }
}

/*
 * We have received data from a single render node.
 *
 * Move the data into the row buffer if there's space in it, otherwise
 * leave it inside the buffer.
 *
 * If this fills up our slice, check and see if the entire row is now full. If
 * it is, pass the row to render_png_row, clear out the col values, and call
 * render_remote_shake for each remote render
 */
static void _render_multi_data (struct evbuffer *buf, void *arg) {
    struct render_multi_sub *ctx = arg;

    assert(ctx->col <= ctx->slice_width);

    // are we being called from _render_multi_done?
    if (ctx->render_done && ctx->in_buf == NULL) {
        // move all the data into our own buffer
        if (!(ctx->in_buf = evbuffer_new()))
            ERROR("evbuffer_new");

        if (evbuffer_add_buffer(ctx->in_buf, buf))
            ERROR("evbuffer_add_buffer");
        
        // don't do anything else
        return;
    }

    // is our slice full?
    if (ctx->col == ctx->slice_width) {
        // don't care for new data
        return;
    }
    
    // read new data into our slice
//    printf("rowbuf + %4d : %d");
    ctx->col += evbuffer_remove(buf, ctx->self->rowbuf + ctx->row_offset + ctx->col, ctx->slice_width - ctx->col);

    // is our slice full now?
    if (ctx->col == ctx->slice_width) {
        // is the row complete now?
        int i;
        for (i = 0; i < ctx->self->remote_render_count; i++) {
            if (ctx->self->remote_renders[i].col < ctx->self->remote_renders[i].slice_width)
                break;
        }
        
        // are they all full?
        if (i == ctx->self->remote_render_count) {
            // pass the data to render_png, this results in calls to _render_multi_png_data
            if (render_png_row(ctx->self->png_info, ctx->self->rowbuf))
                ERROR("render_png_row");

            // clear the col values
            for (i = 0; i < ctx->self->remote_render_count; i++) {
                ctx->self->remote_renders[i].col = 0;
            }
            
            // shake the buffers
            // XXX: this will result in recursion
            for (i = 0; i < ctx->self->remote_render_count; i++) {
                if (ctx->self->remote_renders[i].remote_render) {
                    render_remote_flush(ctx->self->remote_renders[i].remote_render);
                } else {
                    // already disconnected, use in_buf instead
                    assert(ctx->self->remote_renders[i].in_buf);

                    // call it directly...
                    _render_multi_data(ctx->self->remote_renders[i].in_buf, &ctx->self->remote_renders[i]);
                }
            }
        }
    }

    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: avoid these many data copies?
    if (evbuffer_add(ctx->out_buf, data, length))
        ERROR("evbuffer_add");
    
    // only call the data cb if we've already called cb_sent
    if (ctx->have_sent)
        ctx->cb_data(ctx->out_buf, ctx->cb_arg);
    
    // ok
    return 0;

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

/*
 * The render node has rendered everything that it needs to render. Our slice
 * of the last row should be full (or empty) now. If all the remote_renders are
 * done, we can call render_png_done and then cb_done.
 */
static void _render_multi_done (void *arg) {
    struct render_multi_sub *ctx = arg;
 
    // mark it as done
    ctx->render_done = 1;

    // shake out the rest of the data as needed, as render_multi won't keep it anymore
    render_remote_flush(ctx->remote_render);
   
    // invalidate this ctx's remote render
    ctx->remote_render = NULL;

    // is the data incomplete?
    if (!(ctx->col == ctx->slice_width || ctx->col == 0))
        ERROR("incomplete data for slice %d: %zu/%zu bytes", ctx->index, ctx->col, ctx->slice_width);

    // are all of them done?
    int i;
    for (i = 0; i < ctx->self->remote_render_count; i++) {
        if (!ctx->self->remote_renders[i].render_done)
            break;
    }
    
    // are they all done?
    if (i == ctx->self->remote_render_count) {
        // finish it off
        _render_multi_do_done(ctx->self);
    }

    // ok, wait for the rest to complete
    return;
    
error:
    _render_multi_do_fail(ctx->self, 0);
}

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

    // invalidate this ctx's remote render
    ctx->remote_render = NULL;

    _render_multi_do_fail(ctx->self, 0);
}

#define ROUND_DOWN(dividend, divisor) ((dividend) / (divisor))
#define ROUND_UP(dividend, divisor) (((dividend) / (divisor)) + ((dividend) % (divisor)))

#define HALF(a, b) (( a + b) / 2)

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;
    struct render r_left, r_right;
    
    // alloc the render_multi
    ctx = calloc(1, sizeof(struct render_multi));

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

    // init the remote_render
    // for now, just split it in half into two render_ts
    ctx->remote_renders[0].index = 0;
    ctx->remote_renders[0].self = ctx;
    ctx->remote_renders[0].slice_width = render->img_w / 2;
    ctx->remote_renders[0].row_offset = 0;

    ctx->remote_renders[1].index = 1;
    ctx->remote_renders[1].self = ctx;
    ctx->remote_renders[1].slice_width = render->img_w / 2 + render->img_w % 2;
    ctx->remote_renders[1].row_offset = render->img_w / 2;

    ctx->remote_render_count = 2;

    assert(RENDER_MULTI_NODES_MAX >= 2);

    if (
            render_init(&r_left, RENDER_RAW)
         || render_init(&r_right, RENDER_RAW)
         || render_set_size(&r_left, ctx->remote_renders[0].slice_width, render->img_h)
         || render_set_size(&r_right, ctx->remote_renders[1].slice_width, render->img_h)
         || render_region_raw(&r_left, render->x1, render->y1, HALF(render->x1, render->x2), render->y2)
         || render_region_raw(&r_right, HALF(render->x1, render->x2), render->y1, render->x2, render->y2)
    )
        ERROR("render_{init,set_size,set_region_raw}");
            
    // 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;

    // our rowbuf
    if (!(ctx->rowbuf = malloc(render->img_w)))
        ERROR("malloc");

    // 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");

    // png info
    if (!(ctx->png_info = render_png_init(render)))
        ERROR("render_png_init");

    // pull two nodes from the pool
    struct remote_node *node_left, *node_right;

    if (
            !(node_left = remote_pool_get(pool_info))
         || !(node_right = remote_pool_get(pool_info))
    )
        ERROR("remote_pool_get");
    
    // the two render_remote calls
    if (
            !(ctx->remote_renders[0].remote_render = render_remote(&r_left, node_left, 
                &_render_multi_sent, &_render_multi_data, &_render_multi_done, &_render_multi_fail, &ctx->remote_renders[0]))
         || !(ctx->remote_renders[1].remote_render = render_remote(&r_right, node_right,
                &_render_multi_sent, &_render_multi_data, &_render_multi_done, &_render_multi_fail, &ctx->remote_renders[1]))
    )
        ERROR("render_remote");
    
    // 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_cancel (struct render_multi *ctx) {
    _render_multi_do_fail(ctx, FAIL_SILENT);
}

int render_multi_set_recv (struct render_multi *ctx, size_t recv_threshold, size_t unread_buffer) {
    return 0;
}

int render_multi_shake (struct render_multi *ctx) {
    // call the data cb
    ctx->cb_data(ctx->out_buf, ctx->cb_arg);

    return 0;
}