#include "pipeline.h"

#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
#include <assert.h>
#include <limits.h>

struct ptx__pipeline_stage_tag
{
  struct ptx_pipeline_stage_funcs impl;
  char *name;
  void *user;
};

struct ptx__pipeline_ctx_tag
{
  ptx_pipeline_t *pipeline;
  size_t curidx;
};

struct ptx__pipeline_tag
{
  size_t stages_len, stages_cap;
  struct ptx__pipeline_stage_tag *stages;
  bool handling;
};

/* returns:
 * -1 : allocation failed
 *  0 : nothing was allocated
 *  1 : pipeline grown */
static int ptx__pipeline_reserve(ptx_pipeline_t *pl, size_t req)
{
  assert(pl);
  assert(req < SSIZE_MAX);

  if (req < pl->stages_cap) return 0;

  struct ptx__pipeline_stage_tag *temp = realloc(pl->stages, req * sizeof(struct ptx__pipeline_stage_tag));
  if (!temp) return -1;

  pl->stages = temp;
  pl->stages_cap = req;

  return 1;
}

/* allocates a new pipeline, with space for `initial_size' new stages. */
PTX_API ptx_pipeline_t *ptx_pipeline_new(size_t initial_size)
{
  ptx_pipeline_t *pl = malloc(sizeof(ptx_pipeline_t));
  if (!pl) return NULL;

  pl->stages_len = 0;
  pl->stages_cap = 0;
  pl->stages = NULL;
  pl->handling = false;

  ptx__pipeline_reserve(pl, initial_size < 1 ? 1 : initial_size);

  return pl;
}

static void ptx__cleanup_pipeline_stage(struct ptx__pipeline_stage_tag *stage)
{
  assert(stage);

  if (stage->impl.cleanup) {
    (*stage->impl.cleanup)(stage->name, stage->user);
  }

  free(stage->name);

  /* clobber the structure just cuz */
  memset(stage, 0, sizeof(*stage));
}

/* frees an existing pipeline stage. note that this will call the cleanup function in the stages as well. */
PTX_API void ptx_pipeline_free(ptx_pipeline_t *pl)
{
  if (!pl) return;

  for (size_t idx = 0; idx < pl->stages_len; ++idx) {
    ptx__cleanup_pipeline_stage(pl->stages + idx);
  }

  free(pl->stages);
  free(pl);
}

static int ptx__init_pipeline_stage(
    struct ptx__pipeline_stage_tag *PTX_RESTRICT pstage,
    const struct ptx_pipeline_stage_funcs *PTX_RESTRICT stage_impl,
    const char *name,
    va_list setup_args)
{
  assert(pstage);
  assert(stage_impl);
  assert(name);

  char *namedup = NULL;
  void *user = NULL;

  /* I'd rather try to duplicate the name first. Otherwise, we'd have to clean up the
   * stage's user data right after initializing it. */
  namedup = strdup(name);
  if (!namedup) goto cleanup;

  if (stage_impl->init && (*stage_impl->init)(name, &user, setup_args) < 0) {
    /* setup failed :((( */
    goto cleanup;
  }

  /* yay it worked :) */
  pstage->user = user;
  pstage->name = namedup;
  memcpy(&pstage->impl, stage_impl, sizeof(*stage_impl));

  return 0;

cleanup:
  free(namedup);
  return -1;
}

/* this will grow a pipeline to fit a new stage.
 * this function also shifts entries in `pl->stages' such that `to_add' is an empty stage. */
static int ptx__pipeline_prepare_insert_stage_at(ptx_pipeline_t *pipeline, size_t to_add)
{
  assert(to_add <= pipeline->stages_len);
  if (ptx__pipeline_reserve(pipeline, pipeline->stages_len + 1) < 0) {
    return -1;
  }

  /* don't perform the move if we wouldn't be moving anything. */
  if (to_add < pipeline->stages_len) {
    memmove(pipeline->stages + to_add + 1,
        pipeline->stages + to_add,
        sizeof(struct ptx__pipeline_stage_tag) * (pipeline->stages_len - to_add));
  }

  /* increment the length of the array here for consistency.
   * NOTE: this function still leaves the pipeline in an inconsistent state:
   * there is an uninitialized stage at `to_add' :( */
  ++pipeline->stages_len;

  return 0;
}

static ssize_t ptx__pipeline_find_stage(ptx_pipeline_t *ptx, const char *name)
{
  assert(ptx);
  assert(name);

  for (size_t idx = 0; idx < ptx->stages_len; ++idx) {
    if (!strcmp(ptx->stages[idx].name, name)) {
      assert(idx <= SSIZE_MAX);
      return (ssize_t)idx;
    }
  }

  return -1;
}

static int ptx__pipeline_add_at(
    ptx_pipeline_t *pipeline,
    const struct ptx_pipeline_stage_funcs *stage_impl,
    const char *name,
    size_t idx,
    va_list init_args)
{
  struct ptx__pipeline_stage_tag temp_stage;
  memset(&temp_stage, 0, sizeof(temp_stage));

  /* first, initialize the stage itself */
  if (ptx__init_pipeline_stage(&temp_stage, stage_impl, name, init_args) < 0) {
    return -1;
  }

  /* next, grow the pipeline to fit the new stage */
  if (ptx__pipeline_prepare_insert_stage_at(pipeline, idx) < 0) {
    /* uh oh, it failed. clean up the stage we just initialized */
    ptx__cleanup_pipeline_stage(&temp_stage);
    return -1;
  }

  /* finally, put the stage into the array */
  memcpy(pipeline->stages + idx, &temp_stage, sizeof(temp_stage));
  return 0;
}

/* adds `stage' with `name' to `pipeline' after `after' *
 * returns: number of stages added to pipeline (0 or 1), or -1 if there was an allocation or initialization error. */
PTX_API int ptx_pipeline_add_after(
    ptx_pipeline_t *pipeline,
    const struct ptx_pipeline_stage_funcs *stage,
    const char *name,
    const char *after,
    ...)
{
  assert(pipeline);
  assert(!pipeline->handling);

  size_t idx;

  if (after) {
    ssize_t sidx = ptx__pipeline_find_stage(pipeline, after);
    if (sidx < 0) return 0;

    idx = (size_t)(sidx + 1); /* after */
  } else {
    idx = pipeline->stages_len;
  }

  va_list init_args;
  va_start(init_args, after);

  int ret = ptx__pipeline_add_at(pipeline, stage, name, idx, init_args);

  va_end(init_args);
  return ret;
}

/* adds `stage' with `name' to `pipeline' before `before'
 * returns: number of stages added to pipeline (0 or 1), or -1 if there was an allocation or initialization error. */
PTX_API int ptx_pipeline_add_before(
    ptx_pipeline_t *pipeline,
    const struct ptx_pipeline_stage_funcs *stage,
    const char *name,
    const char *before,
    ...)
{
  assert(pipeline);
  assert(!pipeline->handling);

  size_t idx;

  if (before) {
    ssize_t sidx = ptx__pipeline_find_stage(pipeline, before);
    if (sidx < 0) return 0;

    idx = (size_t)sidx; /* before */
  } else {
    idx = 0;
  }

  va_list init_args;
  va_start(init_args, before);

  int ret = ptx__pipeline_add_at(pipeline, stage, name, idx, init_args);

  va_end(init_args);
  return ret;
}

/* returns: number of stages removed from pipeline */
PTX_API int ptx_pipeline_remove_stage(ptx_pipeline_t *pipeline, const char *name)
{
  assert(pipeline);
  assert(!pipeline->handling);

  ssize_t idx = ptx__pipeline_find_stage(pipeline, name);
  if (idx < 0) return 0;

  ptx__cleanup_pipeline_stage(pipeline->stages + idx);
  memmove(pipeline->stages + idx, pipeline->stages + idx + 1, sizeof(pipeline->stages[0]) * (pipeline->stages_len - (size_t)idx - 1));

  --pipeline->stages_len;
  return 1;
}

PTX_API void **ptx_pipeline_ctx_get_user(const ptx_pipeline_ctx_t *ctx)
{
  assert(ctx);
  return &ctx->pipeline->stages[ctx->curidx].user;
}

PTX_API const char *ptx_pipeline_ctx_get_name(const ptx_pipeline_ctx_t *ctx)
{
  assert(ctx);
  return ctx->pipeline->stages[ctx->curidx].name;
}

PTX_API int ptx_pipeline_ctx_next(const ptx_pipeline_ctx_t *ctx, const void *nextdata, size_t nextsize)
{
  assert(ctx->curidx + 1 < ctx->pipeline->stages_len);

  ptx_pipeline_ctx_t new_ctx;
  new_ctx.pipeline = ctx->pipeline;
  new_ctx.curidx = ctx->curidx + 1;
  
  assert(new_ctx.pipeline->stages[new_ctx.curidx].impl.handler);

  return (*new_ctx.pipeline->stages[new_ctx.curidx].impl.handler)(&new_ctx, nextdata, nextsize);
}

PTX_API int ptx_pipeline_handle(ptx_pipeline_t *pipeline, const void *data, size_t sz)
{
  assert(pipeline);
  assert(!pipeline->handling);

  if (pipeline->stages_len == 0) return 0; /* pipeline does nothing... */

  pipeline->handling = true;

  ptx_pipeline_ctx_t ctx;
  ctx.pipeline = pipeline;
  ctx.curidx = 0;

  assert(pipeline->stages[0].impl.handler);
  int ret = (*pipeline->stages[0].impl.handler)(&ctx, data, sz);

  pipeline->handling = 0;
  return ret;
}