#include "amduat/enc/fcs1.h"

#include "amduat/enc/asl1_core_codec.h"

#include <string.h>

enum {
  AMDUAT_FCS1_HEADER_LEN = 7,
  AMDUAT_FCS1_TAG_FUNCTION = 0x30,
  AMDUAT_FCS1_TAG_PARAMS = 0x31,
  AMDUAT_FCS1_TAG_ARITY = 0x32,
  AMDUAT_FCS1_TAG_LEGACY_MIN = 0x33,
  AMDUAT_FCS1_TAG_LEGACY_MAX = 0x36
};

typedef struct {
  const uint8_t *data;
  size_t len;
  size_t offset;
} amduat_fcs1_cursor_t;

static bool amduat_fcs1_read_varint(amduat_fcs1_cursor_t *cur,
                                    uint64_t *out,
                                    bool *out_minimal) {
  uint64_t value = 0u;
  uint32_t shift = 0u;
  size_t bytes = 0u;

  if (cur == NULL || out == NULL || out_minimal == NULL) {
    return false;
  }
  while (cur->offset < cur->len) {
    uint8_t byte = cur->data[cur->offset++];
    value |= ((uint64_t)(byte & 0x7fu)) << shift;
    bytes++;
    if ((byte & 0x80u) == 0u) {
      break;
    }
    shift += 7u;
    if (shift >= 64u) {
      return false;
    }
  }
  if (bytes == 0u || cur->offset > cur->len) {
    return false;
  }
  {
    size_t minimal_len = 1u;
    uint64_t tmp = value;
    while (tmp >= 0x80u) {
      tmp >>= 7u;
      minimal_len++;
    }
    *out_minimal = (minimal_len == bytes);
  }
  *out = value;
  return true;
}

static bool amduat_fcs1_read_bytes(amduat_fcs1_cursor_t *cur,
                                   amduat_octets_t *out,
                                   bool *out_minimal) {
  uint64_t len;
  bool minimal;

  if (!amduat_fcs1_read_varint(cur, &len, &minimal)) {
    return false;
  }
  if (len > SIZE_MAX || cur->len - cur->offset < (size_t)len) {
    return false;
  }
  *out = amduat_octets(cur->data + cur->offset, (size_t)len);
  cur->offset += (size_t)len;
  *out_minimal = minimal;
  return true;
}

static bool amduat_fcs1_has_header(amduat_octets_t bytes) {
  static const uint8_t magic[4] = {'F', 'C', 'S', '1'};

  if (bytes.len < AMDUAT_FCS1_HEADER_LEN || bytes.data == NULL) {
    return false;
  }
  if (memcmp(bytes.data, magic, sizeof(magic)) != 0) {
    return false;
  }
  if (bytes.data[4] != 0x01u || bytes.data[5] != 0x00u ||
      bytes.data[6] != 0x00u) {
    return false;
  }
  return true;
}

void amduat_fcs1_descriptor_free(amduat_fcs1_descriptor_t *desc) {
  if (desc == NULL) {
    return;
  }
  amduat_reference_free(&desc->function_ptr);
  amduat_reference_free(&desc->parameter_block);
  desc->arity = 0u;
}

amduat_fcs1_error_t amduat_fcs1_decode_v1(amduat_octets_t bytes,
                                          amduat_fcs1_descriptor_t *out_desc) {
  amduat_fcs1_cursor_t cur;
  uint8_t last_tag = 0u;
  bool seen_function = false;
  bool seen_params = false;
  bool seen_arity = false;

  if (out_desc == NULL) {
    return AMDUAT_FCS1_ERR_INVALID;
  }
  out_desc->function_ptr = amduat_reference(0, amduat_octets(NULL, 0u));
  out_desc->parameter_block = amduat_reference(0, amduat_octets(NULL, 0u));
  out_desc->arity = 0u;

  if (!amduat_fcs1_has_header(bytes)) {
    return AMDUAT_FCS1_ERR_HEADER_INVALID;
  }

  cur.data = bytes.data;
  cur.len = bytes.len;
  cur.offset = AMDUAT_FCS1_HEADER_LEN;

  while (cur.offset < cur.len) {
    uint8_t tag = cur.data[cur.offset++];
    bool minimal = true;

    if (tag <= last_tag) {
      return AMDUAT_FCS1_ERR_TAG_ORDER;
    }
    last_tag = tag;

    if (tag == AMDUAT_FCS1_TAG_FUNCTION ||
        tag == AMDUAT_FCS1_TAG_PARAMS) {
      amduat_octets_t cid_bytes;
      if (!amduat_fcs1_read_bytes(&cur, &cid_bytes, &minimal)) {
        return AMDUAT_FCS1_ERR_INVALID;
      }
      if (!minimal) {
        return AMDUAT_FCS1_ERR_VARINT_NON_MINIMAL;
      }
      if (tag == AMDUAT_FCS1_TAG_FUNCTION) {
        if (seen_function) {
          return AMDUAT_FCS1_ERR_TAG_ORDER;
        }
        if (!amduat_enc_asl1_core_decode_reference_v1(cid_bytes,
                                                      &out_desc->function_ptr)) {
          return AMDUAT_FCS1_ERR_INVALID;
        }
        seen_function = true;
      } else {
        if (seen_params) {
          return AMDUAT_FCS1_ERR_TAG_ORDER;
        }
        if (!amduat_enc_asl1_core_decode_reference_v1(cid_bytes,
                                                      &out_desc->parameter_block)) {
          return AMDUAT_FCS1_ERR_INVALID;
        }
        seen_params = true;
      }
    } else if (tag == AMDUAT_FCS1_TAG_ARITY) {
      uint64_t arity_raw = 0u;
      if (seen_arity) {
        return AMDUAT_FCS1_ERR_TAG_ORDER;
      }
      if (!amduat_fcs1_read_varint(&cur, &arity_raw, &minimal)) {
        return AMDUAT_FCS1_ERR_INVALID;
      }
      if (!minimal) {
        return AMDUAT_FCS1_ERR_VARINT_NON_MINIMAL;
      }
      if (arity_raw > UINT32_MAX) {
        return AMDUAT_FCS1_ERR_INVALID;
      }
      out_desc->arity = (uint32_t)arity_raw;
      seen_arity = true;
    } else if (tag >= AMDUAT_FCS1_TAG_LEGACY_MIN &&
               tag <= AMDUAT_FCS1_TAG_LEGACY_MAX) {
      return AMDUAT_FCS1_ERR_UNKNOWN_TAG;
    } else {
      return AMDUAT_FCS1_ERR_UNKNOWN_TAG;
    }
  }

  if (!seen_function || !seen_params || !seen_arity) {
    return AMDUAT_FCS1_ERR_TAG_ORDER;
  }
  if (cur.offset != cur.len) {
    return AMDUAT_FCS1_ERR_TRAILING_BYTES;
  }
  return AMDUAT_FCS1_OK;
}

static amduat_fcs1_error_t amduat_fcs1_validate_cycle(
    amduat_reference_t ref,
    amduat_fcs1_resolve_fn resolver,
    void *resolver_ctx,
    amduat_reference_t *stack,
    size_t stack_len,
    size_t stack_cap) {
  amduat_octets_t bytes = amduat_octets(NULL, 0u);
  amduat_fcs1_descriptor_t desc;
  amduat_fcs1_error_t err;

  if (resolver == NULL) {
    return AMDUAT_FCS1_ERR_INVALID;
  }
  if (!resolver(resolver_ctx, ref, &bytes)) {
    return AMDUAT_FCS1_OK;
  }
  if (!amduat_fcs1_has_header(bytes)) {
    amduat_octets_free(&bytes);
    return AMDUAT_FCS1_OK;
  }

  for (size_t i = 0u; i < stack_len; ++i) {
    if (amduat_reference_eq(stack[i], ref)) {
      amduat_octets_free(&bytes);
      return AMDUAT_FCS1_ERR_CYCLE_DETECTED;
    }
  }

  memset(&desc, 0, sizeof(desc));
  err = amduat_fcs1_decode_v1(bytes, &desc);
  amduat_octets_free(&bytes);
  if (err != AMDUAT_FCS1_OK) {
    amduat_fcs1_descriptor_free(&desc);
    return err;
  }

  if (stack_len >= stack_cap) {
    amduat_fcs1_descriptor_free(&desc);
    return AMDUAT_FCS1_ERR_INVALID;
  }
  stack[stack_len] = desc.function_ptr;
  err = amduat_fcs1_validate_cycle(desc.function_ptr,
                                   resolver,
                                   resolver_ctx,
                                   stack,
                                   stack_len + 1u,
                                   stack_cap);
  amduat_fcs1_descriptor_free(&desc);
  return err;
}

amduat_fcs1_error_t amduat_fcs1_validate_v1(
    amduat_octets_t bytes,
    amduat_fcs1_resolve_fn resolver,
    void *resolver_ctx) {
  amduat_fcs1_descriptor_t desc;
  amduat_octets_t pcb_bytes = amduat_octets(NULL, 0u);
  amduat_pcb1_manifest_t manifest;
  amduat_octets_t slot_data;
  amduat_fcs1_error_t err;
  amduat_pcb1_error_t pcb_err;
  amduat_reference_t stack[16];

  memset(&desc, 0, sizeof(desc));
  err = amduat_fcs1_decode_v1(bytes, &desc);
  if (err != AMDUAT_FCS1_OK) {
    return err;
  }

  if (resolver == NULL ||
      !resolver(resolver_ctx, desc.parameter_block, &pcb_bytes)) {
    amduat_fcs1_descriptor_free(&desc);
    return AMDUAT_FCS1_ERR_PARAMETER_FORMAT;
  }

  pcb_err = amduat_pcb1_decode_v1(pcb_bytes, &manifest, &slot_data);
  amduat_octets_free(&pcb_bytes);
  if (pcb_err != AMDUAT_PCB1_OK) {
    amduat_fcs1_descriptor_free(&desc);
    return AMDUAT_FCS1_ERR_PARAMETER_FORMAT;
  }
  (void)slot_data;

  if (desc.arity != manifest.slot_count) {
    amduat_fcs1_descriptor_free(&desc);
    return AMDUAT_FCS1_ERR_PCB_ARITY_MISMATCH;
  }

  stack[0] = desc.function_ptr;
  err = amduat_fcs1_validate_cycle(desc.function_ptr,
                                   resolver,
                                   resolver_ctx,
                                   stack,
                                   1u,
                                   sizeof(stack) / sizeof(stack[0]));
  amduat_fcs1_descriptor_free(&desc);
  return err;
}