#pragma once

#include "stream_class.h"
#include <cstddef>
#include <cstring>
#include <limits>
#include <sstream>
#include <stdexcept>
#include <stdint.h>
#include <tuple>
#include <type_traits>
#include <vector>

using namespace onnx_light_helpers;

namespace ONNX_LIGHT_NAMESPACE {

inline int64_t checked_size_to_int64(size_t value) {
  EXT_ENFORCE(value <= static_cast<size_t>(std::numeric_limits<int64_t>::max()),
              "Serialized size exceeds int64_t.");
  return static_cast<int64_t>(value);
}

inline int64_t checked_uint64_to_int64(uint64_t value) {
  EXT_ENFORCE(value <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max()),
              "Serialized size exceeds int64_t.");
  return static_cast<int64_t>(value);
}

inline SerializeSizeResult make_proto_size(uint64_t proto_size) {
  return {0, 0, checked_uint64_to_int64(proto_size)};
}

// Computes the SerializeSize of a sub-message and stores it in the stream's
// size cache (keyed by the sub-message address). The subsequent write pass
// uses write_with_cache_size with the same key and so reuses the cached
// result without re-traversing the entire subtree.
template <typename T>
inline SerializeSizeResult size_with_cache(utils::BinaryWriteStream &stream, const T &field,
                                           SerializeOptions &options) {
  SerializeSizeResult s;
  if (!stream.GetCachedSize(reinterpret_cast<const void *>(&field), s)) {
    s = field.SerializeSize(stream, options);
    stream.CacheSize(reinterpret_cast<const void *>(&field), s);
  }
  return s;
}

template <typename T>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const T &field,
                               SerializeOptions &options) {
  auto s = size_with_cache(stream, field, options);
  return {s.small_data_size, s.big_data_size,
          checked_uint64_to_int64(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                                  stream.VarintSize(static_cast<uint64_t>(s.proto_size)) +
                                  static_cast<uint64_t>(s.proto_size))};
}

template <typename T>
SerializeSizeResult size_optional_proto_field(utils::BinaryWriteStream &stream, int order,
                                              const utils::OptionalField<T> &field,
                                              SerializeOptions &options) {
  if (field.has_value()) {
    auto s = size_with_cache(stream, *field, options);
    return {s.small_data_size, s.big_data_size,
            checked_uint64_to_int64(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                                    stream.VarintSize(static_cast<uint64_t>(s.proto_size)) +
                                    static_cast<uint64_t>(s.proto_size))};
  }
  return {};
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::String &field, SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.size_string(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::OptionalString &field, SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.size_string(utils::RefString(field.data(), field.size())));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::OptionalField<uint64_t> &field, SerializeOptions &) {
  if (field.has_value()) {
    return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                           stream.size_variant_uint64(*field));
  }
  return {};
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::OptionalField<int64_t> &field, SerializeOptions &) {
  if (field.has_value()) {
    return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                           stream.size_int64(*field));
  }
  return {};
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::OptionalField<int32_t> &field, SerializeOptions &) {
  if (field.has_value()) {
    return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                           stream.size_int32(*field));
  }
  return {};
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::OptionalField<float> &field, SerializeOptions &) {
  if (field.has_value()) {
    return make_proto_size(stream.size_field_header(order, FIELD_FIXED32) +
                           stream.size_float(*field));
  }
  return {};
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const int64_t &field,
                               SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_VARINT) + stream.size_int64(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const uint64_t &field,
                               SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                         stream.size_variant_uint64(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const int32_t &field,
                               SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_VARINT) + stream.size_int32(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const double &field,
                               SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.size_double(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order, const float &field,
                               SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED32) + stream.size_float(field));
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const std::vector<uint8_t> &field, SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.VarintSize(field.size()) + field.size());
}

template <>
SerializeSizeResult size_field(utils::BinaryWriteStream &stream, int order,
                               const utils::ByteSpan &field, SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.VarintSize(field.size()) + field.size());
}

SerializeSizeResult size_field_limit(utils::BinaryWriteStream &stream, int order,
                                     const std::vector<uint8_t> &field, SerializeOptions &options) {
  if (!options.skip_raw_data || field.size() < static_cast<size_t>(options.raw_data_threshold)) {
    if (stream.ExternalWeights() &&
        static_cast<int64_t>(field.size()) >= options.raw_data_threshold) {
      SerializeSizeResult size;
      size.add_tensor_data_size(checked_size_to_int64(field.size()), options.raw_data_threshold);
      return size;
    } else {
      return size_field(stream, order, field, options);
    }
  }
  return {};
}

SerializeSizeResult size_field_limit(utils::BinaryWriteStream &stream, int order,
                                     const utils::ByteSpan &field, SerializeOptions &options) {
  const size_t sz = field.size();
  if (!options.skip_raw_data || sz < static_cast<size_t>(options.raw_data_threshold)) {
    if (stream.ExternalWeights() && static_cast<int64_t>(sz) >= options.raw_data_threshold) {
      SerializeSizeResult size;
      size.add_tensor_data_size(checked_size_to_int64(sz), options.raw_data_threshold);
      return size;
    } else {
      return size_field(stream, order, field, options);
    }
  }
  return {};
}

template <typename T>
SerializeSizeResult size_enum_field(utils::BinaryWriteStream &stream, int order, const T &field,
                                    SerializeOptions &) {
  return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                         stream.VarintSize(static_cast<uint64_t>(field)));
}

// repeated fields

template <typename T>
SerializeSizeResult size_repeated_field(utils::BinaryWriteStream &stream, int order,
                                        const utils::RepeatedField<T> &field, bool is_packed,
                                        SerializeOptions &options) {
  return size_repeated_field(stream, order, field.values(), is_packed, options);
}

template <typename T>
SerializeSizeResult size_repeated_field(utils::BinaryWriteStream &stream, int order,
                                        const utils::RepeatedProtoField<T> &field, bool is_packed,
                                        SerializeOptions &options) {
  EXT_ENFORCE(!is_packed, "option is_packed is not implemented for field order ", order);
  SerializeSizeResult size;
  for (size_t i = 0; i < field.size(); ++i) {
    auto s = size_with_cache(stream, field[i], options);
    size += {s.small_data_size, s.big_data_size,
             checked_uint64_to_int64(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                                     stream.VarintSize(static_cast<uint64_t>(s.proto_size)) +
                                     static_cast<uint64_t>(s.proto_size))};
  }
  return size;
}

template <typename T>
SerializeSizeResult size_repeated_field(utils::BinaryWriteStream &stream, int order,
                                        const std::vector<T> &field, bool is_packed,
                                        SerializeOptions &options) {
  EXT_ENFORCE(!is_packed, "option is_packed is not implemented for field order ", order);
  SerializeSizeResult size;
  for (const auto &d : field) {
    auto s = size_with_cache(stream, d, options);
    size += {s.small_data_size, s.big_data_size,
             checked_uint64_to_int64(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                                     stream.VarintSize(static_cast<uint64_t>(s.proto_size)) +
                                     static_cast<uint64_t>(s.proto_size))};
  }
  return size;
}

template <>
SerializeSizeResult size_repeated_field(utils::BinaryWriteStream &stream, int order,
                                        const std::vector<utils::String> &field, bool is_packed,
                                        SerializeOptions &) {
  EXT_ENFORCE(!is_packed, "option is_packed is not implemented for field order ", order);
  SerializeSizeResult size;
  for (const auto &d : field) {
    size += make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                            stream.VarintSize(d.size()) + d.size());
  }
  return size;
}

// unpacked number

template <typename T>
SerializeSizeResult size_unpacked_number_float(utils::BinaryWriteStream &stream, int order,
                                               const T &) {
  return make_proto_size(
      stream.size_field_header(order, sizeof(T) == 4 ? FIELD_FIXED32 : FIELD_FIXED64) + sizeof(T));
}

template <typename T>
SerializeSizeResult size_unpacked_number_int(utils::BinaryWriteStream &stream, int order,
                                             const T &value) {
  return make_proto_size(stream.size_field_header(order, FIELD_VARINT) +
                         stream.size_variant_uint64(static_cast<uint64_t>(value)));
}

template <typename T>
SerializeSizeResult size_unpacked_number(utils::BinaryWriteStream &stream, int order,
                                         const T &value);

#define SIZE_UNPACKED_NUMBER_FLOAT(type)                                                           \
  template <>                                                                                      \
  SerializeSizeResult size_unpacked_number(utils::BinaryWriteStream &stream, int order,            \
                                           const type &value) {                                    \
    return size_unpacked_number_float(stream, order, value);                                       \
  }

SIZE_UNPACKED_NUMBER_FLOAT(float)
SIZE_UNPACKED_NUMBER_FLOAT(double)

#define SIZE_UNPACKED_NUMBER_INT(type)                                                             \
  template <>                                                                                      \
  SerializeSizeResult size_unpacked_number(utils::BinaryWriteStream &stream, int order,            \
                                           const type &value) {                                    \
    return size_unpacked_number_int(stream, order, value);                                         \
  }

SIZE_UNPACKED_NUMBER_INT(uint64_t)
SIZE_UNPACKED_NUMBER_INT(int64_t)
SIZE_UNPACKED_NUMBER_INT(int32_t)

// repeated packed numbers

template <typename T>
SerializeSizeResult
size_repeated_field_numerical_numbers_float(utils::BinaryWriteStream &stream, int order,
                                            const std::vector<T> &field, bool, SerializeOptions &) {
  uint64_t size = field.size() * sizeof(T);
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.size_variant_uint64(size) + size);
}

template <typename T>
SerializeSizeResult
size_repeated_field_numerical_numbers_int(utils::BinaryWriteStream &stream, int order,
                                          const std::vector<T> &field, bool, SerializeOptions &) {
  utils::StringWriteStream local;
  uint64_t size = 0;
  for (const T &d : field) {
    size += local.size_variant_uint64(static_cast<uint64_t>(d));
  }
  return make_proto_size(stream.size_field_header(order, FIELD_FIXED_SIZE) +
                         stream.size_variant_uint64(size) + size);
}

template <typename T>
SerializeSizeResult size_repeated_field_numerical_numbers(utils::BinaryWriteStream &stream,
                                                          int order, const std::vector<T> &field,
                                                          bool is_packed,
                                                          SerializeOptions &options);

#define SIZE_REPEATED_FIELD_NUMERICAL_FLOAT(type)                                                  \
  template <>                                                                                      \
  SerializeSizeResult size_repeated_field_numerical_numbers(                                       \
      utils::BinaryWriteStream &stream, int order, const std::vector<type> &field, bool is_packed, \
      SerializeOptions &options) {                                                                 \
    return size_repeated_field_numerical_numbers_float(stream, order, field, is_packed, options);  \
  }

SIZE_REPEATED_FIELD_NUMERICAL_FLOAT(float)
SIZE_REPEATED_FIELD_NUMERICAL_FLOAT(double)

#define SIZE_REPEATED_FIELD_NUMERICAL_INT(type)                                                    \
  template <>                                                                                      \
  SerializeSizeResult size_repeated_field_numerical_numbers(                                       \
      utils::BinaryWriteStream &stream, int order, const std::vector<type> &field, bool is_packed, \
      SerializeOptions &options) {                                                                 \
    return size_repeated_field_numerical_numbers_int(stream, order, field, is_packed, options);    \
  }

SIZE_REPEATED_FIELD_NUMERICAL_INT(int64_t)
SIZE_REPEATED_FIELD_NUMERICAL_INT(uint64_t)
SIZE_REPEATED_FIELD_NUMERICAL_INT(int32_t)

// main function to calculate size of repeated numbers

template <typename T>
SerializeSizeResult size_repeated_field_numerical(utils::BinaryWriteStream &stream, int order,
                                                  const std::vector<T> &field, bool is_packed,
                                                  SerializeOptions &options) {
  if (is_packed) {
    return size_repeated_field_numerical_numbers(stream, order, field, is_packed, options);
  } else {
    SerializeSizeResult size;
    for (const T &d : field) {
      size += size_unpacked_number(stream, order, d);
    }
    return size;
  }
}

#define SIZE_REPEATED_FIELD_IMPL(type)                                                             \
  template <>                                                                                      \
  SerializeSizeResult size_repeated_field(utils::BinaryWriteStream &stream, int order,             \
                                          const std::vector<type> &field, bool is_packed,          \
                                          SerializeOptions &options) {                             \
    return size_repeated_field_numerical(stream, order, field, is_packed, options);                \
  }

SIZE_REPEATED_FIELD_IMPL(double)
SIZE_REPEATED_FIELD_IMPL(float)
SIZE_REPEATED_FIELD_IMPL(uint64_t)
SIZE_REPEATED_FIELD_IMPL(int64_t)
SIZE_REPEATED_FIELD_IMPL(int32_t)

} // namespace ONNX_LIGHT_NAMESPACE
