#include "amduat/enc/fcs1.h"
#include "amduat/enc/asl1_core_codec.h"
#include "amduat/fps.h"

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static bool encode_varint(uint64_t value, uint8_t **out, size_t *out_len) {
  uint8_t buf[10];
  size_t len = 0u;

  do {
    uint8_t byte = (uint8_t)(value & 0x7fu);
    value >>= 7u;
    if (value != 0u) {
      byte |= 0x80u;
    }
    buf[len++] = byte;
  } while (value != 0u);

  *out = (uint8_t *)malloc(len);
  if (*out == NULL) {
    return false;
  }
  memcpy(*out, buf, len);
  *out_len = len;
  return true;
}

static bool append_bytes(uint8_t **buf, size_t *len, size_t *cap,
                         const void *data, size_t data_len) {
  if (*len + data_len > *cap) {
    size_t next = *cap == 0u ? 64u : *cap * 2u;
    while (next < *len + data_len) {
      next *= 2u;
    }
    uint8_t *next_buf = (uint8_t *)realloc(*buf, next);
    if (next_buf == NULL) {
      return false;
    }
    *buf = next_buf;
    *cap = next;
  }
  memcpy(*buf + *len, data, data_len);
  *len += data_len;
  return true;
}

static bool append_varint(uint8_t **buf, size_t *len, size_t *cap,
                          uint64_t value) {
  uint8_t *enc = NULL;
  size_t enc_len = 0u;
  bool ok = encode_varint(value, &enc, &enc_len) &&
            append_bytes(buf, len, cap, enc, enc_len);
  free(enc);
  return ok;
}

static bool build_pcb1_bytes(amduat_octets_t *out_bytes,
                             amduat_reference_t *out_cid) {
  uint8_t *buf = NULL;
  size_t len = 0u;
  size_t cap = 0u;
  uint8_t manifest[1 + 1 + 1 + 1 + 32];
  size_t manifest_len = 0u;
  uint8_t slot_data[] = {0x01, 0x02, 0x03};
  amduat_octets_t payload;
  amduat_octets_t cid_bytes;

  manifest[manifest_len++] = 0x01; /* slot count */
  manifest[manifest_len++] = 0x01; /* index */
  manifest[manifest_len++] = 0x00; /* name len */
  manifest[manifest_len++] = 0x00; /* type */
  manifest[manifest_len++] = 0x20; /* digest len */
  memset(manifest + manifest_len, 0, 32);
  manifest_len += 32;

  if (!append_bytes(&buf, &len, &cap, "PCB1", 4u)) {
    return false;
  }
  {
    uint8_t header[3] = {0x01, 0x00, 0x00};
    if (!append_bytes(&buf, &len, &cap, header, sizeof(header))) {
      free(buf);
      return false;
    }
  }

  if (!append_bytes(&buf, &len, &cap, (uint8_t[]){0x50}, 1u) ||
      !append_varint(&buf, &len, &cap, manifest_len) ||
      !append_bytes(&buf, &len, &cap, manifest, manifest_len)) {
    free(buf);
    return false;
  }

  if (!append_bytes(&buf, &len, &cap, (uint8_t[]){0x51}, 1u) ||
      !append_varint(&buf, &len, &cap, sizeof(slot_data)) ||
      !append_bytes(&buf, &len, &cap, slot_data, sizeof(slot_data))) {
    free(buf);
    return false;
  }

  payload = amduat_octets(buf, len);
  if (!amduat_fps_cid_for_payload(AMDUAT_HASH_ASL1_ID_SHA256,
                                  payload,
                                  out_cid)) {
    free(buf);
    return false;
  }
  if (!amduat_enc_asl1_core_encode_reference_v1(*out_cid, &cid_bytes)) {
    amduat_reference_free(out_cid);
    free(buf);
    return false;
  }
  amduat_octets_free(&cid_bytes);

  *out_bytes = payload;
  return true;
}

static bool build_fcs1_bytes(amduat_octets_t *out_bytes,
                             amduat_reference_t function_ref,
                             amduat_reference_t pcb_ref,
                             uint32_t arity) {
  uint8_t *buf = NULL;
  size_t len = 0u;
  size_t cap = 0u;
  amduat_octets_t func_bytes;
  amduat_octets_t pcb_bytes;

  if (!amduat_enc_asl1_core_encode_reference_v1(function_ref, &func_bytes) ||
      !amduat_enc_asl1_core_encode_reference_v1(pcb_ref, &pcb_bytes)) {
    return false;
  }

  if (!append_bytes(&buf, &len, &cap, "FCS1", 4u)) {
    return false;
  }
  {
    uint8_t header[3] = {0x01, 0x00, 0x00};
    if (!append_bytes(&buf, &len, &cap, header, sizeof(header))) {
      free(buf);
      return false;
    }
  }

  if (!append_bytes(&buf, &len, &cap, (uint8_t[]){0x30}, 1u) ||
      !append_varint(&buf, &len, &cap, func_bytes.len) ||
      !append_bytes(&buf, &len, &cap, func_bytes.data, func_bytes.len)) {
    free(buf);
    amduat_octets_free(&func_bytes);
    amduat_octets_free(&pcb_bytes);
    return false;
  }

  if (!append_bytes(&buf, &len, &cap, (uint8_t[]){0x31}, 1u) ||
      !append_varint(&buf, &len, &cap, pcb_bytes.len) ||
      !append_bytes(&buf, &len, &cap, pcb_bytes.data, pcb_bytes.len)) {
    free(buf);
    amduat_octets_free(&func_bytes);
    amduat_octets_free(&pcb_bytes);
    return false;
  }

  if (!append_bytes(&buf, &len, &cap, (uint8_t[]){0x32}, 1u) ||
      !append_varint(&buf, &len, &cap, arity)) {
    free(buf);
    amduat_octets_free(&func_bytes);
    amduat_octets_free(&pcb_bytes);
    return false;
  }

  amduat_octets_free(&func_bytes);
  amduat_octets_free(&pcb_bytes);
  *out_bytes = amduat_octets(buf, len);
  return true;
}

typedef struct {
  amduat_reference_t pcb_ref;
  amduat_octets_t pcb_bytes;
} resolver_ctx_t;

static bool resolve_ref(void *ctx,
                        amduat_reference_t ref,
                        amduat_octets_t *out_bytes) {
  resolver_ctx_t *state = (resolver_ctx_t *)ctx;

  if (state == NULL || out_bytes == NULL) {
    return false;
  }
  if (!amduat_reference_eq(ref, state->pcb_ref)) {
    return false;
  }
  return amduat_octets_clone(state->pcb_bytes, out_bytes);
}

int main(void) {
  amduat_octets_t pcb_bytes;
  amduat_reference_t pcb_ref;
  amduat_octets_t fcs_bytes;
  amduat_reference_t func_ref;
  uint8_t func_digest[32];
  resolver_ctx_t ctx;
  amduat_fcs1_error_t err;

  if (!build_pcb1_bytes(&pcb_bytes, &pcb_ref)) {
    return 1;
  }
  memset(func_digest, 0x11, sizeof(func_digest));
  func_ref = amduat_reference(AMDUAT_HASH_ASL1_ID_SHA256,
                              amduat_octets(func_digest,
                                            sizeof(func_digest)));
  if (!build_fcs1_bytes(&fcs_bytes, func_ref, pcb_ref, 1u)) {
    amduat_reference_free(&pcb_ref);
    amduat_octets_free(&pcb_bytes);
    return 1;
  }

  ctx.pcb_ref = pcb_ref;
  ctx.pcb_bytes = pcb_bytes;
  err = amduat_fcs1_validate_v1(fcs_bytes, resolve_ref, &ctx);
  if (err != AMDUAT_FCS1_OK) {
  amduat_reference_free(&pcb_ref);
  amduat_octets_free(&pcb_bytes);
  amduat_octets_free(&fcs_bytes);
  return 1;
}

  amduat_reference_free(&pcb_ref);
  amduat_octets_free(&pcb_bytes);
  amduat_octets_free(&fcs_bytes);
  return 0;
}