#include "input_buffer.h"

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

enum { KIT_IB_CACHE_SIZE = 256 };

static s32 kit_buf_adjust_(kit_input_buffer_t *buf, i64 size) {
  assert(buf != NULL);
  assert(size >= 0);

  if (buf == NULL)
    return KIT_ERROR_INTERNAL;

  i64 offset = buf->data.size;

  if (offset >= size)
    return KIT_OK;

  DA_RESIZE(buf->data, size);
  if (buf->data.size != size)
    return KIT_ERROR_OUT_OF_MEMORY;

  str_t destination = { .size   = size - offset,
                        .values = buf->data.values + offset };
  i64   n           = IS_READ(buf->upstream, destination);

  DA_RESIZE(buf->data, offset + n);
  if (buf->data.size != offset + n)
    return KIT_ERROR_OUT_OF_MEMORY;

  return KIT_OK;
}

kit_input_buffer_t kit_ib_init(is_handle_t      upstream,
                               kit_allocator_t *alloc) {
  kit_input_buffer_t buf;
  memset(&buf, 0, sizeof buf);

  buf.upstream = upstream;
  DA_INIT(buf.data, 0, alloc);

  return buf;
}

void kit_ib_destroy(kit_input_buffer_t *buf) {
  assert(buf != NULL);
  if (buf == NULL)
    return;

  DA_DESTROY(buf->data);
  memset(buf, 0, sizeof *buf);
}

kit_ib_token_t kit_ib_token(kit_input_buffer_t *buf) {
  return (kit_ib_token_t) {
    .status = KIT_OK, .offset = 0, .size = 0, .buffer = buf
  };
}

kit_str_t kit_ib_str(kit_ib_token_t tok) {
  assert(tok.buffer != NULL);

  if (tok.buffer == NULL || tok.buffer->data.values == NULL)
    return (str_t) { .size = 0, .values = NULL };

  return (str_t) { .size   = tok.size,
                   .values = tok.buffer->data.values + tok.offset };
}

kit_ib_token_t kit_ib_read(kit_ib_token_t tok, i64 size) {
  assert(tok.buffer != NULL);

  kit_ib_token_t next = tok;

  next.offset = tok.offset + tok.size;
  next.size   = 0;

  if (tok.status != KIT_OK)
    return next;

  if (tok.buffer == NULL) {
    next.status = KIT_ERROR_INVALID_ARGUMENT;
    return next;
  }

  s32 s = kit_buf_adjust_(tok.buffer, tok.offset + tok.size + size);

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

  assert(tok.buffer->data.values != NULL);

  if (tok.buffer->data.values == NULL) {
    next.status = KIT_ERROR_INTERNAL;
    return next;
  }

  next.size = size < tok.buffer->data.size - tok.offset - tok.size
                  ? size
                  : tok.buffer->data.size - tok.offset - tok.size;

  return next;
}

#define IB_CACHE_INIT_(res_)                                       \
  str_builder_t cache_dynamic;                                     \
  memset(&cache_dynamic, 0, sizeof cache_dynamic);                 \
  char cache_static[KIT_IB_CACHE_SIZE];                            \
                                                                   \
  do {                                                             \
    if (data.size > 0) {                                           \
      if (data.size < KIT_IB_CACHE_SIZE) {                         \
        memcpy(cache_static, data.values, data.size);              \
        data.values = cache_static;                                \
      } else {                                                     \
        DA_INIT(cache_dynamic, data.size, tok.buffer->data.alloc); \
        if (cache_dynamic.size != data.size) {                     \
          (res_).status |= KIT_ERROR_OUT_OF_MEMORY;                \
          return (res_);                                           \
        }                                                          \
        memcpy(cache_dynamic.values, data.values, data.size);      \
        data.values = cache_dynamic.values;                        \
      }                                                            \
    }                                                              \
  } while (0)

#define IB_CACHE_CLEANUP_()           \
  do {                                \
    if (cache_dynamic.values != NULL) \
      DA_DESTROY(cache_dynamic);      \
  } while (0)

kit_ib_token_t kit_ib_any(kit_ib_token_t tok, str_t data) {
  assert(tok.buffer != NULL);

  kit_ib_token_t next = tok;

  next.offset = tok.offset + tok.size;
  next.size   = 0;

  if (tok.status != KIT_OK)
    return next;

  if (tok.buffer == NULL) {
    next.status = KIT_ERROR_INVALID_ARGUMENT;
    return next;
  }

  IB_CACHE_INIT_(next);

  for (;; ++next.size) {
    s32 s = kit_buf_adjust_(tok.buffer,
                            tok.offset + tok.size + next.size + 1);

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

    assert(tok.buffer->data.values != NULL);

    if (tok.buffer->data.values == NULL) {
      next.status = KIT_ERROR_INTERNAL;
      return next;
    }

    if (tok.offset + tok.size + next.size >= tok.buffer->data.size)
      break;

    i8 found = 0;

    for (i64 i = 0; i < data.size; i++)
      if (data.values[i] ==
          tok.buffer->data.values[next.offset + next.size]) {
        found = 1;
        break;
      }

    if (!found)
      break;
  }

  IB_CACHE_CLEANUP_();

  return next;
}

kit_ib_token_t kit_ib_none(kit_ib_token_t tok, str_t data) {
  assert(tok.buffer != NULL);

  kit_ib_token_t next = tok;

  next.offset = tok.offset + tok.size;
  next.size   = 0;

  if (tok.status != KIT_OK)
    return next;

  if (tok.buffer == NULL) {
    next.status = KIT_ERROR_INVALID_ARGUMENT;
    return next;
  }

  IB_CACHE_INIT_(next);

  for (;; ++next.size) {
    s32 s = kit_buf_adjust_(tok.buffer,
                            tok.offset + tok.size + next.size + 1);

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

    assert(tok.buffer->data.values != NULL);

    if (tok.buffer->data.values == NULL) {
      next.status = KIT_ERROR_INTERNAL;
      return next;
    }

    if (tok.offset + tok.size + next.size >= tok.buffer->data.size)
      break;

    i8 found = 0;

    for (i64 i = 0; i < data.size; i++)
      if (data.values[i] ==
          tok.buffer->data.values[next.offset + next.size]) {
        found = 1;
        break;
      }

    if (found)
      break;
  }

  IB_CACHE_CLEANUP_();

  return next;
}

kit_ib_token_t kit_ib_exact(kit_ib_token_t tok, str_t data) {
  kit_ib_token_t res = tok;

  res.offset = tok.offset + tok.size;
  res.size   = 0;

  IB_CACHE_INIT_(res);

  res = kit_ib_read(tok, data.size);
  if (!AR_EQUAL(kit_ib_str(res), data))
    res.status = KIT_PARSING_FAILED;

  IB_CACHE_CLEANUP_();

  return res;
}

kit_ib_token_t kit_ib_until(kit_ib_token_t tok, str_t data) {
  assert(tok.buffer != NULL);

  kit_ib_token_t next = tok;

  next.offset = tok.offset + tok.size;
  next.size   = 0;

  if (tok.status != KIT_OK)
    return next;

  if (tok.buffer == NULL) {
    next.status = KIT_ERROR_INVALID_ARGUMENT;
    return next;
  }

  IB_CACHE_INIT_(next);

  for (;; ++next.size) {
    s32 s = kit_buf_adjust_(tok.buffer,
                            tok.offset + tok.size + next.size + 1);

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

    assert(tok.buffer->data.values != NULL);

    if (tok.buffer->data.values == NULL) {
      next.status = KIT_ERROR_INTERNAL;
      return next;
    }

    if (tok.offset + tok.size + next.size >= tok.buffer->data.size)
      break;

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

  IB_CACHE_CLEANUP_();

  return next;
}

kit_ib_token_t kit_ib_while(kit_ib_token_t           tok,
                            kit_ib_read_condition_fn condition,
                            void                    *context) {
  assert(tok.buffer != NULL);

  kit_ib_token_t next = tok;

  next.offset = tok.offset + tok.size;
  next.size   = 0;

  if (tok.status != KIT_OK)
    return next;

  if (tok.buffer == NULL) {
    next.status = KIT_ERROR_INVALID_ARGUMENT;
    return next;
  }

  for (;; ++next.size) {
    s32 s = kit_buf_adjust_(tok.buffer,
                            tok.offset + tok.size + next.size + 1);

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

    assert(tok.buffer->data.values != NULL);

    if (tok.buffer->data.values == NULL) {
      next.status = KIT_ERROR_INTERNAL;
      return next;
    }

    if (tok.offset + tok.size + next.size >= tok.buffer->data.size)
      break;

    if (condition == NULL ||
        !condition(kit_str(next.size + 1,
                           tok.buffer->data.values + next.offset),
                   context))
      break;
  }

  return next;
}