#include <stdint.h>
#include <stddef.h>
#include "FastLZ/fastlz.h"
#include "std/string.h"

#define SIXPACK_OK                0
#define SIXPACK_ERR_MAGIC        -1
#define SIXPACK_ERR_CHECKSUM     -2
#define SIXPACK_ERR_DECOMPRESS   -3
#define SIXPACK_ERR_BAD_FORMAT   -4
#define SIXPACK_ERR_NO_SPACE     -5

#define ADLER32_BASE 65521
#define BLOCK_SIZE   65536

static const uint8_t sixpack_magic[8] = {137, '6', 'P', 'K', 13, 10, 26, 10};

static uint32_t update_adler32(uint32_t checksum, const void *buf, int len) {
  const uint8_t *ptr = (const uint8_t*)buf;
  uint32_t s1 = checksum & 0xffff;
  uint32_t s2 = (checksum >> 16) & 0xffff;

  while (len > 0) {
    unsigned k = len < 5552 ? len : 5552;
    len -= k;
    while (k--) {
      s1 += *ptr++;
      if (s1 >= ADLER32_BASE) s1 -= ADLER32_BASE;
      s2 += s1;
      if (s2 >= ADLER32_BASE) s2 -= ADLER32_BASE;
    }
  }
  return (s2 << 16) + s1;
}

static uint32_t readU16(const uint8_t *p) {
  return (uint32_t)p[0] | ((uint32_t)p[1] << 8);
}
static uint32_t readU32(const uint8_t *p) {
  return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void read_chunk_header(const uint8_t *p, int *id, int *opts,
                              uint32_t *size, uint32_t *checksum, uint32_t *extra) {
  *id = (int)readU16(p);
  *opts = (int)readU16(p + 2);
  *size = readU32(p + 4);
  *checksum = readU32(p + 8);
  *extra = readU32(p + 12);
}

int sixpack_decompress_mem(const uint8_t *input, size_t in_size,
                           uint8_t *output, size_t out_cap)
{
  if (!input || in_size < 8)
    return SIXPACK_ERR_BAD_FORMAT;

  /* Check magic */
  for (int i = 0; i < 8; i++)
    if (input[i] != sixpack_magic[i])
      return SIXPACK_ERR_MAGIC;

  size_t pos = 8; /* skip magic */
  size_t written = 0;

  while (pos + 16 <= in_size) {
    int chunk_id, opts;
    uint32_t size, checksum, extra;

    read_chunk_header(input + pos, &chunk_id, &opts, &size, &checksum, &extra);
    pos += 16;

    if (pos + size > in_size)
      return SIXPACK_ERR_BAD_FORMAT;

    const uint8_t *chunk_data = input + pos;
    pos += size;

    /* File header chunk */
    if (chunk_id == 1) {
      uint32_t decomp_size = readU32(chunk_data);
      (void)decomp_size; /* not strictly needed here */
      continue;
    }

    /* Data chunk */
    if (chunk_id == 17) {
      uint32_t cksum = update_adler32(1L, chunk_data, size);
      if (cksum != checksum)
        return SIXPACK_ERR_CHECKSUM;

      if (opts == 0) {
        /* Stored (uncompressed) */
        if (written + size > out_cap)
          return SIXPACK_ERR_NO_SPACE;
        memcpy(output + written, chunk_data, size);
        written += size;
      } else if (opts == 1) {
        /* FastLZ compressed */
        if (written + extra > out_cap)
          return SIXPACK_ERR_NO_SPACE;

        int dec = fastlz_decompress(chunk_data, size, output + written, extra);
        if (dec != (int)extra)
          return SIXPACK_ERR_DECOMPRESS;

        written += extra;
      } else {
        return SIXPACK_ERR_BAD_FORMAT;
      }
    }
  }

  return (int)written;
}