#include "input_buffer.h"

#include <assert.h>
#include <string.h>

typedef struct {
  i64               ref_count;
  kit_is_handle_t   upstream;
  kit_allocator_t  *alloc;
  kit_str_builder_t data;
} internal_buffer_t;

static internal_buffer_t *kit_buf_init_(kit_is_handle_t  upstream,
                                        kit_allocator_t *alloc) {
  internal_buffer_t *buf = kit_alloc_dispatch(alloc, KIT_ALLOCATE,
                                              sizeof *buf, 0, NULL);

  if (buf != NULL) {
    memset(buf, 0, sizeof *buf);
    buf->ref_count = 1;
    buf->upstream  = upstream;
    buf->alloc     = alloc;
    DA_INIT(buf->data, 0, alloc);
  }

  return buf;
}

static kit_allocator_t *kit_buf_alloc_(void *p) {
  assert(p != NULL);

  return ((internal_buffer_t *) p)->alloc;
}

static void kit_buf_acquire_(void *p) {
  assert(p != NULL);

  ((internal_buffer_t *) p)->ref_count++;
}

static void kit_buf_release_(void *p) {
  if (p == NULL)
    return;

  internal_buffer_t *buf = (internal_buffer_t *) p;

  if (--buf->ref_count == 0) {
    DA_DESTROY(buf->data);
    kit_alloc_dispatch(buf->alloc, KIT_DEALLOCATE, 0, 0, buf);
  }
}

static void kit_buf_adjust_(void *p, i64 size) {
  assert(p != NULL);
  assert(size >= 0);

  internal_buffer_t *buf    = (internal_buffer_t *) p;
  i64                offset = buf->data.size;

  if (offset < size) {
    DA_RESIZE(buf->data, size);
    kit_str_t destination = { .size   = size - offset,
                              .values = buf->data.values + offset };
    i64       n           = KIT_IS_READ(buf->upstream, destination);
    DA_RESIZE(buf->data, offset + n);
  }
}

static i64 kit_buf_read_(void *p, i64 offset, kit_str_t destination) {
  internal_buffer_t *buf = (internal_buffer_t *) p;
  i64                n   = destination.size < buf->data.size - offset
                               ? destination.size
                               : buf->data.size - offset;
  memcpy(destination.values, buf->data.values + offset, n);
  return n;
}

kit_ib_t kit_ib_wrap(kit_is_handle_t  upstream,
                     kit_allocator_t *alloc) {
  kit_ib_t buf;
  memset(&buf, 0, sizeof buf);
  buf.status = KIT_OK;
  DA_INIT(buf.data, 0, alloc);
  buf.internal = kit_buf_init_(upstream, alloc);
  if (buf.internal == NULL)
    buf.status = KIT_ERROR_BAD_ALLOC;
  return buf;
}

kit_ib_t kit_ib_copy(kit_ib_t buf) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);

  next.offset   = buf.offset;
  next.internal = buf.internal;

  return next;
}

kit_ib_t kit_ib_read(kit_ib_t buf, i64 size) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);
  kit_buf_adjust_(buf.internal, buf.offset + size);

  DA_INIT(next.data, size, kit_buf_alloc_(buf.internal));
  if (next.data.size != size)
    next.status = KIT_ERROR_BAD_ALLOC;

  kit_str_t destination = { .size   = next.data.size,
                            .values = next.data.values };
  i64       n = kit_buf_read_(buf.internal, buf.offset, destination);
  next.offset = buf.offset + n;
  next.internal = buf.internal;

  DA_RESIZE(next.data, n);
  if (next.data.size != n)
    next.status = KIT_ERROR_BAD_ALLOC;

  return next;
}

kit_ib_t kit_ib_any(kit_ib_t buf, kit_str_t data) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);

  DA_INIT(next.data, 0, kit_buf_alloc_(buf.internal));

  i64 size = 0;

  for (;; ++size) {
    kit_buf_adjust_(buf.internal, buf.offset + size + 1);

    DA_RESIZE(next.data, size + 1);

    assert(next.data.size == size + 1);
    if (next.data.size != size + 1) {
      next.status = KIT_ERROR_BAD_ALLOC;
      return next;
    }

    kit_str_t destination = { .size   = 1,
                              .values = next.data.values + size };
    i64       n = kit_buf_read_(buf.internal, buf.offset + size,
                                destination);

    if (n != 1)
      break;

    i8 found = 0;

    for (i64 i = 0; i < data.size; i++)
      if (data.values[i] == destination.values[0]) {
        found = 1;
        break;
      }

    if (!found)
      break;
  }

  next.offset   = buf.offset + size;
  next.internal = buf.internal;

  DA_RESIZE(next.data, size);
  if (next.data.size != size)
    next.status = KIT_ERROR_BAD_ALLOC;

  return next;
}

kit_ib_t kit_ib_none(kit_ib_t buf, kit_str_t data) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);

  DA_INIT(next.data, 0, kit_buf_alloc_(buf.internal));

  i64 size = 0;

  for (;; ++size) {
    kit_buf_adjust_(buf.internal, buf.offset + size + 1);

    DA_RESIZE(next.data, size + 1);

    assert(next.data.size == size + 1);
    if (next.data.size != size + 1) {
      next.status = KIT_ERROR_BAD_ALLOC;
      return next;
    }

    kit_str_t destination = { .size   = 1,
                              .values = next.data.values + size };
    i64       n = kit_buf_read_(buf.internal, buf.offset + size,
                                destination);

    if (n != 1)
      break;

    i8 found = 0;

    for (i64 i = 0; i < data.size; i++)
      if (data.values[i] == destination.values[0]) {
        found = 1;
        break;
      }

    if (found)
      break;
  }

  next.offset   = buf.offset + size;
  next.internal = buf.internal;

  DA_RESIZE(next.data, size);
  if (next.data.size != size)
    next.status = KIT_ERROR_BAD_ALLOC;

  return next;
}

kit_ib_t kit_ib_exact(kit_ib_t buf, kit_str_t data) {
  kit_ib_t res = kit_ib_read(buf, data.size);
  if (!AR_EQUAL(res.data, data))
    res.status = KIT_ERROR_INTERNAL;
  return res;
}

kit_ib_t kit_ib_until(kit_ib_t buf, kit_str_t data) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);

  DA_INIT(next.data, 0, kit_buf_alloc_(buf.internal));

  i64 size = 0;

  for (;; ++size) {
    kit_buf_adjust_(buf.internal, buf.offset + size + 1);

    DA_RESIZE(next.data, size + 1);

    assert(next.data.size == size + 1);
    if (next.data.size != size + 1) {
      next.status = KIT_ERROR_BAD_ALLOC;
      return next;
    }

    kit_str_t destination = { .size   = 1,
                              .values = next.data.values + size };
    i64       n = kit_buf_read_(buf.internal, buf.offset + size,
                                destination);

    if (n != 1)
      break;

    if (size + 1 >= data.size &&
        AR_EQUAL(kit_str(data.size,
                         next.data.values + (size + 1 - data.size)),
                 data)) {
      size -= data.size - 1;
      break;
    }
  }

  next.offset   = buf.offset + size;
  next.internal = buf.internal;

  DA_RESIZE(next.data, size);
  if (next.data.size != size)
    next.status = KIT_ERROR_BAD_ALLOC;

  return next;
}

kit_ib_t kit_ib_while(kit_ib_t                 buf,
                      kit_ib_read_condition_fn condition,
                      void                    *context) {
  kit_ib_t next;
  memset(&next, 0, sizeof next);

  if (buf.status != KIT_OK) {
    next.status = buf.status;
    return next;
  }

  kit_buf_acquire_(buf.internal);

  DA_INIT(next.data, 0, kit_buf_alloc_(buf.internal));

  i64 size = 0;

  for (;; ++size) {
    kit_buf_adjust_(buf.internal, buf.offset + size + 1);

    DA_RESIZE(next.data, size + 1);

    assert(next.data.size == size + 1);
    if (next.data.size != size + 1) {
      next.status = KIT_ERROR_BAD_ALLOC;
      return next;
    }

    kit_str_t destination = { .size   = 1,
                              .values = next.data.values + size };
    i64       n = kit_buf_read_(buf.internal, buf.offset + size,
                                destination);

    kit_str_t data = { .size = size + 1, .values = next.data.values };
    if (n != 1 || condition == NULL || condition(data, context) == 0)
      break;
  }

  next.offset   = buf.offset + size;
  next.internal = buf.internal;

  DA_RESIZE(next.data, size);
  if (next.data.size != size)
    next.status = KIT_ERROR_BAD_ALLOC;

  return next;
}

void kit_ib_destroy(kit_ib_t buf) {
  kit_buf_release_(buf.internal);
  DA_DESTROY(buf.data);
}