#pragma once

#include "descriptions/header.hpp"

#include "coding/dd_vector.hpp"
#include "coding/string_utf8_multilang.hpp"
#include "coding/text_storage.hpp"

#include "base/assert.hpp"
#include "base/buffer_vector.hpp"
#include "base/stl_helpers.hpp"

#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <vector>

namespace descriptions
{
using FeatureIndex = uint32_t;
using StringIndex = uint32_t;
using LangCode = int8_t;
using LangMeta = buffer_vector<std::pair<LangCode, StringIndex>, 8>;
using LangMetaOffset = uint32_t;

enum class Version : uint8_t
{
  V0 = 0,
  Latest = V0
};

struct FeatureDescription
{
  FeatureIndex m_ftIndex = 0;
  LangMeta m_strIndices;
};

struct DescriptionsCollection
{
  std::vector<FeatureDescription> m_features;
  std::vector<std::string> m_strings;

  size_t GetFeaturesCount() const { return m_features.size(); }
};

/// \brief
/// Section name: "descriptions".
/// Description: keeping text descriptions of features in different languages.
/// Section tables:
/// * version
/// * header
/// * sorted feature ids vector
/// * vector of unordered maps with language codes and string indices of corresponding translations of a description
/// * vector of maps offsets for each feature id (and one additional dummy offset in the end)
/// * BWT-compressed strings grouped by language.
class Serializer
{
public:
  /// \param descriptions A non-empty unsorted collection of feature descriptions.
  ///                     FeatureDescription::m_description must contain non-empty translations.
  explicit Serializer(DescriptionsCollection && descriptions) : m_collection(std::move(descriptions))
  {
    std::sort(m_collection.m_features.begin(), m_collection.m_features.end(),
              base::LessBy(&FeatureDescription::m_ftIndex));
  }

  template <typename Sink>
  void Serialize(Sink & sink)
  {
    WriteToSink(sink, static_cast<uint8_t>(Version::Latest));

    auto const startPos = sink.Pos();

    HeaderV0 header;
    header.Serialize(sink);

    header.m_featuresOffset = sink.Pos() - startPos;
    SerializeFeaturesIndices(sink);

    std::vector<LangMetaOffset> offsets;
    header.m_langMetaOffset = sink.Pos() - startPos;
    SerializeLangMetaCollection(sink, offsets);

    header.m_indexOffset = sink.Pos() - startPos;
    SerializeLangMetaIndex(sink, offsets);

    header.m_stringsOffset = sink.Pos() - startPos;
    SerializeStrings(sink);

    header.m_eosOffset = sink.Pos() - startPos;
    sink.Seek(startPos);
    header.Serialize(sink);
    sink.Seek(startPos + header.m_eosOffset);
  }

  // Serializes a vector of 32-bit sorted feature ids.
  template <typename Sink>
  void SerializeFeaturesIndices(Sink & sink)
  {
    for (auto const & index : m_collection.m_features)
      WriteToSink(sink, index.m_ftIndex);
  }

  template <typename Sink>
  void SerializeLangMetaCollection(Sink & sink, std::vector<LangMetaOffset> & offsets)
  {
    auto const startPos = sink.Pos();
    for (auto const & meta : m_collection.m_features)
    {
      offsets.push_back(static_cast<LangMetaOffset>(sink.Pos() - startPos));
      for (auto const & pair : meta.m_strIndices)
      {
        WriteToSink(sink, pair.first);
        WriteVarUint(sink, pair.second);
      }
    }
    offsets.push_back(static_cast<LangMetaOffset>(sink.Pos() - startPos));
  }

  template <typename Sink>
  void SerializeLangMetaIndex(Sink & sink, std::vector<LangMetaOffset> const & offsets)
  {
    for (auto const & offset : offsets)
      WriteToSink(sink, offset);
  }

  // Serializes strings in a compressed storage with block access.
  template <typename Sink>
  void SerializeStrings(Sink & sink)
  {
    coding::BlockedTextStorageWriter<Sink> writer(sink, 200000 /* blockSize */);
    for (auto const & s : m_collection.m_strings)
      writer.Append(s);
  }

private:
  DescriptionsCollection m_collection;
};

class Deserializer
{
public:
  using LangPriorities = LangsBufferT;

  template <typename ReaderT>
  std::string Deserialize(ReaderT & reader, FeatureIndex featureIndex, LangPriorities const & langPriority)
  {
    InitializeIfNeeded(reader);

    auto const it = std::lower_bound(m_featureIds->begin(), m_featureIds->end(), featureIndex);
    if (it == m_featureIds->end() || *it != featureIndex)
      return {};

    auto const d = static_cast<uint32_t>(std::distance(m_featureIds->begin(), it));
    ASSERT_LESS(d, m_langMetaOffsets->size(), ());
    ASSERT_LESS(d + 1, m_langMetaOffsets->size(), ());

    auto const startOffset = (*m_langMetaOffsets)[d];
    auto const endOffset = (*m_langMetaOffsets)[d + 1];

    LangMeta langMeta;
    {
      auto const pos = m_header.m_langMetaOffset + startOffset;
      NonOwningReaderSource source(*m_dataReader, pos, pos + endOffset - startOffset);
      while (source.Size() > 0)
      {
        auto const lang = ReadPrimitiveFromSource<LangCode>(source);
        auto const stringIndex = ReadVarUint<StringIndex>(source);
        langMeta.emplace_back(lang, stringIndex);
      }
    }

    for (LangCode const lang : langPriority)
    {
      for (auto const & meta : langMeta)
        if (lang == meta.first)
          return m_stringsReader.ExtractString(*m_stringsSubReader, meta.second);
    }

    return {};
  }

private:
  template <typename ReaderT>
  void InitializeIfNeeded(ReaderT & reader)
  {
    if (m_initialized)
      return;

    {
      NonOwningReaderSource source(reader);
      auto const version = static_cast<Version>(ReadPrimitiveFromSource<uint8_t>(source));
      CHECK(version == Version::V0, ());
      m_dataReader = reader.CreateSubReader(source.Pos(), source.Size());
    }

    {
      NonOwningReaderSource source(*m_dataReader);
      m_header.Deserialize(source);
    }

    m_featureIds.emplace(ReaderPtr<Reader>(m_dataReader->CreateSubReader(
        m_header.m_featuresOffset, m_header.m_langMetaOffset - m_header.m_featuresOffset)));

    m_langMetaOffsets.emplace(ReaderPtr<Reader>(
        m_dataReader->CreateSubReader(m_header.m_indexOffset, m_header.m_stringsOffset - m_header.m_indexOffset)));

    m_stringsSubReader =
        m_dataReader->CreateSubReader(m_header.m_stringsOffset, m_header.m_eosOffset - m_header.m_stringsOffset);

    m_initialized = true;
  }

  bool m_initialized = false;
  HeaderV0 m_header;
  std::unique_ptr<Reader> m_dataReader;
  std::optional<DDVector<FeatureIndex, ReaderPtr<Reader>>> m_featureIds;
  std::optional<DDVector<LangMetaOffset, ReaderPtr<Reader>>> m_langMetaOffsets;
  std::unique_ptr<Reader> m_stringsSubReader;
  coding::BlockedTextStorageReader m_stringsReader;
};
}  // namespace descriptions
