#include "msgpack.h"

namespace msgpack {

result<body_info> get_body_info(const uint8_t *p, int size) {
  if (size < 1) return std::unexpected(error_code::empty);
  uint8_t b = p[0];

  using namespace format;

  if (is_positive_fixint(b))  return body_info{1, 0};
  if (is_negative_fixint(b))  return body_info{1, 0};
  if (is_fixmap(b))           return body_info{1, 0};
  if (is_fixarray(b))         return body_info{1, 0};
  if (is_fixstr(b))           return body_info{1, static_cast<uint32_t>(b & 0x1F)};

  switch (b) {
  case NIL: case FALSE: case TRUE:
    return body_info{1, 0};
  case NEVER_USED:
    return std::unexpected(error_code::invalid);

  case BIN8:  { auto n = body_number<uint8_t>(p, size);  if (!n) return std::unexpected(n.error()); return body_info{1+1, *n}; }
  case BIN16: { auto n = body_number<uint16_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+2, *n}; }
  case BIN32: { auto n = body_number<uint32_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+4, *n}; }

  case EXT8:  { auto n = body_number<uint8_t>(p, size);  if (!n) return std::unexpected(n.error()); return body_info{1+1+1, *n}; }
  case EXT16: { auto n = body_number<uint16_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+2+1, *n}; }
  case EXT32: { auto n = body_number<uint32_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+4+1, *n}; }

  case FLOAT32: return body_info{1, 4};
  case FLOAT64: return body_info{1, 8};
  case UINT8:   return body_info{1, 1};
  case UINT16:  return body_info{1, 2};
  case UINT32:  return body_info{1, 4};
  case UINT64:  return body_info{1, 8};
  case INT8:    return body_info{1, 1};
  case INT16:   return body_info{1, 2};
  case INT32:   return body_info{1, 4};
  case INT64:   return body_info{1, 8};

  case FIXEXT1:  return body_info{1+1, 1};
  case FIXEXT2:  return body_info{1+1, 2};
  case FIXEXT4:  return body_info{1+1, 4};
  case FIXEXT8:  return body_info{1+1, 8};
  case FIXEXT16: return body_info{1+1, 16};

  case STR8:  { auto n = body_number<uint8_t>(p, size);  if (!n) return std::unexpected(n.error()); return body_info{1+1, *n}; }
  case STR16: { auto n = body_number<uint16_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+2, *n}; }
  case STR32: { auto n = body_number<uint32_t>(p, size); if (!n) return std::unexpected(n.error()); return body_info{1+4, *n}; }

  case ARRAY16: case ARRAY32:
  case MAP16: case MAP32:
    return body_info{1 + (b == ARRAY16 || b == MAP16 ? 2 : 4), 0};

  default:
    return std::unexpected(error_code::invalid);
  }
}

packer::pack_result packer::pack_nil() {
  m_buf.push_back(format::NIL);
  return *this;
}

packer::pack_result packer::pack_bool(bool v) {
  m_buf.push_back(v ? format::TRUE : format::FALSE);
  return *this;
}

packer::pack_result packer::pack_uint32_fixed(uint32_t n) {
  m_buf.push_back(format::UINT32);
  push_big_endian(n);
  return *this;
}

packer::pack_result packer::pack_float(float n) {
  m_buf.push_back(format::FLOAT32);
  push_big_endian(n);
  return *this;
}

packer::pack_result packer::pack_double(double n) {
  m_buf.push_back(format::FLOAT64);
  push_big_endian(n);
  return *this;
}

packer::pack_result packer::pack_str(const char *s) {
  return pack_str(std::string_view(s));
}

packer::pack_result packer::pack_array(size_t n) {
  if (n <= 15) {
    m_buf.push_back(format::FIXARRAY_MIN | static_cast<uint8_t>(n));
  } else if (n <= std::numeric_limits<uint16_t>::max()) {
    m_buf.push_back(format::ARRAY16);
    push_big_endian(static_cast<uint16_t>(n));
  } else if (n <= std::numeric_limits<uint32_t>::max()) {
    m_buf.push_back(format::ARRAY32);
    push_big_endian(static_cast<uint32_t>(n));
  } else {
    return std::unexpected(error_code::overflow);
  }
  return *this;
}

packer::pack_result packer::pack_map(size_t n) {
  if (n <= 15) {
    m_buf.push_back(format::FIXMAP_MIN | static_cast<uint8_t>(n));
  } else if (n <= std::numeric_limits<uint16_t>::max()) {
    m_buf.push_back(format::MAP16);
    push_big_endian(static_cast<uint16_t>(n));
  } else if (n <= std::numeric_limits<uint32_t>::max()) {
    m_buf.push_back(format::MAP32);
    push_big_endian(static_cast<uint32_t>(n));
  } else {
    return std::unexpected(error_code::overflow);
  }
  return *this;
}

packer::pack_result packer::pack_ext16_header(char type, uint16_t len) {
  m_buf.push_back(format::EXT16);
  push_big_endian(len);
  m_buf.push_back(static_cast<uint8_t>(type));
  return *this;
}

packer::pack_result packer::pack_bin16_header(uint16_t len) {
  m_buf.push_back(format::BIN16);
  push_big_endian(len);
  return *this;
}

packer::pack_result packer::pack(bool v)        { return pack_bool(v); }
packer::pack_result packer::pack(float v)       { return pack_float(v); }
packer::pack_result packer::pack(double v)      { return pack_double(v); }
packer::pack_result packer::pack(const char *v) { return pack_str(v); }
packer::pack_result packer::pack(std::string_view v) { return pack_str(v); }
packer::pack_result packer::pack(const std::string &v) { return pack_str(v); }
packer::pack_result packer::pack(const std::vector<uint8_t> &v) { return pack_bin(v); }

result<parser> parser::next() const {
  auto hdr = header_byte();
  if (!hdr) return std::unexpected(hdr.error());

  if (is_array()) {
    auto info = get_body_info(m_p, m_size);
    if (!info) return std::unexpected(info.error());
    auto cnt = count();
    if (!cnt) return std::unexpected(cnt.error());
    auto cur = advance(info->header);
    if (!cur) return std::unexpected(cur.error());
    for (uint32_t i = 0; i < *cnt; ++i) {
      auto n = cur->next();
      if (!n) return std::unexpected(n.error());
      cur = *n;
    }
    return *cur;
  } else if (is_map()) {
    auto info = get_body_info(m_p, m_size);
    if (!info) return std::unexpected(info.error());
    auto cnt = count();
    if (!cnt) return std::unexpected(cnt.error());
    auto cur = advance(info->header);
    if (!cur) return std::unexpected(cur.error());
    for (uint32_t i = 0; i < *cnt; ++i) {
      auto k = cur->next();
      if (!k) return std::unexpected(k.error());
      cur = *k;
      auto v = cur->next();
      if (!v) return std::unexpected(v.error());
      cur = *v;
    }
    return *cur;
  } else {
    auto info = get_body_info(m_p, m_size);
    if (!info) return std::unexpected(info.error());
    auto total = info->header + static_cast<int>(info->body);
    return advance(total);
  }
}

bool parser::is_nil() const {
  auto h = header_byte();
  return h && *h == format::NIL;
}

bool parser::is_bool() const {
  auto h = header_byte();
  return h && (*h == format::TRUE || *h == format::FALSE);
}

bool parser::is_number() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  if (format::is_positive_fixint(b)) return true;
  if (format::is_negative_fixint(b)) return true;
  return b >= format::FLOAT32 && b <= format::INT64;
}

bool parser::is_string() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  if (format::is_fixstr(b)) return true;
  return b == format::STR8 || b == format::STR16 || b == format::STR32;
}

bool parser::is_binary() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  return b == format::BIN8 || b == format::BIN16 || b == format::BIN32;
}

bool parser::is_ext() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  return (b >= format::FIXEXT1 && b <= format::FIXEXT16) ||
         b == format::EXT8 || b == format::EXT16 || b == format::EXT32;
}

bool parser::is_array() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  if (format::is_fixarray(b)) return true;
  return b == format::ARRAY16 || b == format::ARRAY32;
}

bool parser::is_map() const {
  auto h = header_byte();
  if (!h) return false;
  uint8_t b = *h;
  if (format::is_fixmap(b)) return true;
  return b == format::MAP16 || b == format::MAP32;
}

result<bool> parser::get_bool() const {
  auto h = header_byte();
  if (!h) return std::unexpected(h.error());
  if (*h == format::TRUE) return true;
  if (*h == format::FALSE) return false;
  return std::unexpected(error_code::type_error);
}

result<std::string_view> parser::get_string() const {
  auto h = header_byte();
  if (!h) return std::unexpected(h.error());
  uint8_t b = *h;
  size_t offset, len;
  if (format::is_fixstr(b)) {
    len = b & 0x1F;
    offset = 1;
  } else if (b == format::STR8) {
    auto n = body_number<uint8_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 1;
  } else if (b == format::STR16) {
    auto n = body_number<uint16_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 2;
  } else if (b == format::STR32) {
    auto n = body_number<uint32_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 4;
  } else {
    return std::unexpected(error_code::type_error);
  }
  if (static_cast<int>(offset + len) > m_size) {
    return std::unexpected(error_code::lack);
  }
  return std::string_view(reinterpret_cast<const char *>(m_p + offset), len);
}

result<std::string_view> parser::get_binary_view() const {
  auto h = header_byte();
  if (!h) return std::unexpected(h.error());
  uint8_t b = *h;
  size_t offset, len;

  if (b == format::BIN8) {
    auto n = body_number<uint8_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 1;
  } else if (b == format::BIN16) {
    auto n = body_number<uint16_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 2;
  } else if (b == format::BIN32) {
    auto n = body_number<uint32_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    len = *n; offset = 1 + 4;
  } else {
    return std::unexpected(error_code::type_error);
  }
  if (static_cast<int>(offset + len) > m_size) {
    return std::unexpected(error_code::lack);
  }
  return std::string_view(reinterpret_cast<const char *>(m_p + offset), len);
}

result<std::tuple<int8_t, std::string_view>> parser::get_ext() const {
  auto h = header_byte();
  if (!h) return std::unexpected(h.error());
  uint8_t b = *h;
  int8_t ext_type;
  size_t data_offset, data_len;

  switch (b) {
  case format::FIXEXT1:  ext_type = m_p[1]; data_offset = 2; data_len = 1;  break;
  case format::FIXEXT2:  ext_type = m_p[1]; data_offset = 2; data_len = 2;  break;
  case format::FIXEXT4:  ext_type = m_p[1]; data_offset = 2; data_len = 4;  break;
  case format::FIXEXT8:  ext_type = m_p[1]; data_offset = 2; data_len = 8;  break;
  case format::FIXEXT16: ext_type = m_p[1]; data_offset = 2; data_len = 16; break;
  case format::EXT8: {
    auto n = body_number<uint8_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    ext_type = m_p[2]; data_offset = 3; data_len = *n;
    break;
  }
  case format::EXT16: {
    auto n = body_number<uint16_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    ext_type = m_p[3]; data_offset = 4; data_len = *n;
    break;
  }
  case format::EXT32: {
    auto n = body_number<uint32_t>(m_p, m_size);
    if (!n) return std::unexpected(n.error());
    ext_type = m_p[5]; data_offset = 6; data_len = *n;
    break;
  }
  default:
    return std::unexpected(error_code::type_error);
  }
  if (static_cast<int>(data_offset + data_len) > m_size) {
    return std::unexpected(error_code::lack);
  }
  return std::tuple{ext_type,
                    std::string_view(reinterpret_cast<const char *>(m_p + data_offset), data_len)};
}

result<uint32_t> parser::count() const {
  auto h = header_byte();
  if (!h) return std::unexpected(h.error());
  uint8_t b = *h;

  if (format::is_fixarray(b)) return static_cast<uint32_t>(b & 0x0F);
  if (format::is_fixmap(b))   return static_cast<uint32_t>(b & 0x0F);

  switch (b) {
  case format::ARRAY16: { auto n = body_number<uint16_t>(m_p, m_size); if (!n) return std::unexpected(n.error()); return static_cast<uint32_t>(*n); }
  case format::ARRAY32: { auto n = body_number<uint32_t>(m_p, m_size); if (!n) return std::unexpected(n.error()); return *n; }
  case format::MAP16:   { auto n = body_number<uint16_t>(m_p, m_size); if (!n) return std::unexpected(n.error()); return static_cast<uint32_t>(*n); }
  case format::MAP32:   { auto n = body_number<uint32_t>(m_p, m_size); if (!n) return std::unexpected(n.error()); return *n; }
  default:
    return std::unexpected(error_code::type_error);
  }
}

result<parser> parser::first_item() const {
  if (!is_array() && !is_map()) return std::unexpected(error_code::type_error);
  auto info = get_body_info(m_p, m_size);
  if (!info) return std::unexpected(info.error());
  return advance(info->header);
}

parser parser::operator[](int index) const {
  auto cur = first_item();
  if (!cur) return {};
  for (int i = 0; i < index; ++i) {
    auto n = cur->next();
    if (!n) return {};
    cur = *n;
  }
  return *cur;
}

result<parser> unpack(const parser &p, bool &out) {
  auto v = p.get_bool();
  if (!v) return std::unexpected(v.error());
  out = *v;
  return p.next();
}

result<parser> unpack(const parser &p, std::string_view &out) {
  auto v = p.get_string();
  if (!v) return std::unexpected(v.error());
  out = *v;
  return p.next();
}

result<parser> unpack(const parser &p, std::string &out) {
  auto v = p.get_string();
  if (!v) return std::unexpected(v.error());
  out = std::string(v->data(), v->size());
  return p.next();
}

result<parser> unpack(const parser &p, std::vector<uint8_t> &out) {
  auto v = p.get_binary_view();
  if (!v) return std::unexpected(v.error());
  out.assign(v->begin(), v->end());
  return p.next();
}

result<parser> unpack(const parser &p, std::span<const uint8_t> &out) {
  auto v = p.get_binary_view();
  if (!v) return std::unexpected(v.error());
  out = std::span<const uint8_t>(reinterpret_cast<const uint8_t*>(v->data()), v->size());
  return p.next();
}

} // namespace msgpack