module;

#include <algorithm>
#include <bit>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <print>
#include <span>
#include <stack>
#include <string_view>
#include <vector>

export module nbtpp;

namespace nbtpp {

export enum class tag : std::uint8_t {
  end = 0,
  i8 = 1,
  i16 = 2,
  i32 = 3,
  i64 = 4,
  f32 = 5,
  f64 = 6,
  i8_arr = 7,
  string = 8,
  list = 9,
  compound = 10,
  i32_arr = 11,
  i64_arr = 12,
};

void verify(tag required, tag t) {
  if (required != t) {
    throw std::runtime_error("unexpected type on access");
  }
}

template <typename T> void verify_type(tag t) {}

export struct node {
  std::string_view name = {};
  tag type = tag::end;

  std::uint32_t next = std::numeric_limits<std::uint32_t>::max();

private:
  [[nodiscard]] const std::byte *get_data() const {
    return reinterpret_cast<const std::byte *>(name.data() + name.size());
  }

  template <std::integral T> [[nodiscard]] T read() const {
    verify_type<T>(type);
    auto value = T();
    std::memcpy(&value, get_data(), sizeof(T));
    return std::byteswap(value);
  }

  template <std::floating_point T> [[nodiscard]] T read() const {
    verify_type<T>(type);
    auto data = std::array<std::byte, sizeof(T)>();
    std::memcpy(data.data(), get_data(), sizeof(T));
    std::reverse(data.begin(), data.end());
    return std::bit_cast<T>(data);
  }

public:
  template <typename T> [[nodiscard]] T as() const;

  [[nodiscard]] std::size_t advance_size() const {
    using enum tag;
    switch (type) {
    case i8:
      return 1;
    case i16:
      return 2;
    case i32:
      return 4;
    case i64:
      return 8;
    case f32:
      return 4;
    case f64:
      return 8;
    case string: {
      const auto size = read<std::uint16_t>();
      return size + sizeof(std::uint16_t);
    }
    case i8_arr:
    case i32_arr:
    case i64_arr: {
      const auto size = read<std::int32_t>();
      return size + sizeof(std::int32_t);
    };
    default:
      throw std::runtime_error("not expected to work");
    }
  }
};

template <>
std::vector<std::int8_t> node::as<std::vector<std::int8_t>>() const {
  auto length = std::int32_t();
  std::memcpy(&length, get_data(), 4);
  length = std::byteswap(length);

  auto list = std::vector<std::int8_t>(length);
  std::memcpy(list.data(), get_data() + 4, length);
  return list;
}

template <> std::int8_t node::as<std::int8_t>() const {
  return read<std::int8_t>();
}

template <> std::int16_t node::as<std::int16_t>() const {
  return read<std::int16_t>();
}

template <> std::int32_t node::as<std::int32_t>() const {
  return read<std::int32_t>();
}

template <> std::int64_t node::as<std::int64_t>() const {
  return read<std::int64_t>();
}

template <> std::string node::as<std::string>() const {
  const auto length = read<std::uint16_t>();
  auto str = std::string(length, ' ');
  std::memcpy(str.data(), get_data() + 2, length);
  return str;
}

template <> float node::as<float>() const { return read<float>(); }

template <> double node::as<double>() const { return read<double>(); }

template <> void verify_type<std::int8_t>(tag t) { verify(tag::i8, t); }

template <> void verify_type<std::int16_t>(tag t) { verify(tag::i16, t); }

template <> void verify_type<std::int32_t>(tag t) { verify(tag::i32, t); }

template <> void verify_type<std::int64_t>(tag t) { verify(tag::i64, t); }

template <> void verify_type<float>(tag t) { verify(tag::f32, t); }

template <> void verify_type<double>(tag t) { verify(tag::f64, t); }

template <> void verify_type<std::vector<std::int8_t>>(tag t) {
  verify(tag::i8_arr, t);
}

template <> void verify_type<std::string_view>(tag t) {
  verify(tag::string, t);
}

template <> void verify_type<std::vector<node>>(tag t) { verify(tag::list, t); }

template <> void verify_type<std::vector<std::int32_t>>(tag t) {
  verify(tag::i32_arr, t);
}

template <> void verify_type<std::vector<std::int64_t>>(tag t) {
  verify(tag::i64_arr, t);
}

void parse_node_header(std::span<node> nodes, std::int64_t &cursor,
                       const std::byte *bytes) {
  nodes[cursor].type = std::bit_cast<tag>(*bytes);

  if (nodes[cursor].type == tag::end) {
    auto name_length = std::uint16_t();
    std::memcpy(&name_length, bytes + 1, sizeof(std::uint16_t));
    name_length = std::byteswap(name_length);
    nodes[cursor].name = std::string_view(
        reinterpret_cast<const char *>(bytes + 3), name_length);
  }
}

export [[nodiscard]] std::vector<node> parse(std::span<const std::byte> nbt) {
  const auto node_count = nbt.size() / 4;
  auto nodes = std::make_unique<node[]>(node_count);
  auto span = std::span(nodes.get(), nodes.get() + node_count);

  auto cursor = std::int64_t();
  auto node_index = std::int64_t(0);

  const auto read_name = [&] -> std::string_view {
    auto name_length = std::uint16_t();
    std::memcpy(&name_length, &nbt[cursor], sizeof(std::uint16_t));
    cursor += sizeof(std::uint16_t);
    name_length = std::byteswap(name_length);

    const auto string_start = reinterpret_cast<const char *>(&nbt[cursor]);
    cursor += name_length;

    return std::string_view(string_start, name_length);
  };

  auto stack = std::array<std::uint32_t, 128>();
  auto stack_ptr = std::int16_t();

  do {
    const auto node_type = std::bit_cast<tag>(nbt[cursor++]);
    if (node_type == tag::end) {
      const auto parent_index = stack[stack_ptr--];
      nodes[parent_index].next = node_index;
      continue;
    }

    const auto node_name = read_name();

    if (node_type == tag::compound) {
      stack[stack_ptr++] = node_index;
      nodes[node_index++] = {
          .name = node_name,
          .type = node_type,
      };
      continue;
    } else {
      // Primitive node
      nodes[node_index++] = {
          .name = node_name,
          .type = node_type,
      };
      const auto size = nodes[node_index - 1].advance_size();
      std::println("advancing by {} for node {}", size, node_name);
      cursor += size;
    }

  } while (stack_ptr > 0);

  return {nodes.get(), nodes.get() + node_index};
}

} // namespace nbtpp