#pragma once
#include "coding/bit_streams.hpp"
#include "coding/huffman.hpp"
#include "coding/reader.hpp"
#include "coding/varint.hpp"

#include "base/assert.hpp"
#include "base/bits.hpp"
#include "base/checked_cast.hpp"
#include "base/macros.hpp"
#include "base/string_utils.hpp"

#include "std/unique_ptr.hpp"
#include "std/utility.hpp"
#include "std/vector.hpp"

#include "3party/succinct/rs_bit_vector.hpp"

namespace trie
{
// This class is used to read the data that is common to every node
// of the trie: the trie topology and the offsets into the data buffer.
// The topology can then be used to navigate the trie and the offsets
// can be used to extract the values associated with the key strings.
template <class TReader, class TValueReader>
class TopologyAndOffsets
{
public:
  using TValue = typename TValueReader::ValueType;

  TopologyAndOffsets(TReader const & reader, TValueReader const & valueReader)
    : m_reader(reader), m_valueReader(valueReader)
  {
    Parse();
  }

  // Returns the topology of the trie in the external node representation.
  // Arrange the nodes in the level order and write two bits for every
  // node: the first bit for its left child and the second bit for its right child.
  // Write '1' if the child is present and '0' if it is not.
  // Prepend the resulting bit string with '1' (for a super-root) and you will get
  // the external node representaiton.
  succinct::rs_bit_vector const & GetTopology() const { return m_trieTopology; }

  // Returns the number of trie nodes.
  uint32_t NumNodes() const { return m_numNodes; }

  // Returns the Huffman encoding that was used to encode the strings
  // before adding them to this trie.
  coding::HuffmanCoder const & GetEncoding() const { return m_huffman; }

  // Returns true if and only if the node is the last node of a string
  // that was used as a key when building this trie.
  bool NodeIsFinal(uint32_t id) const { return m_finalNodeIndex[id] >= 0; }

  // Returns the offset relative to the beginning of the start of the
  // value buffer where the data for node number |id| starts.
  uint32_t Offset(uint32_t id) const
  {
    CHECK(NodeIsFinal(id), ());
    return m_offsetTable[m_finalNodeIndex[id]];
  }

  // Returns the number of values associated with a final node |id|.
  uint32_t ValueListSize(uint32_t id)
  {
    CHECK(NodeIsFinal(id), ());
    int finalId = m_finalNodeIndex[id];
    return m_offsetTable[finalId + 1] - m_offsetTable[finalId];
  }

  // Returns the reader from which the value buffer can be read.
  TValueReader const & GetValueReader() const { return m_valueReader; }

  // Returns the reader from which the trie can be read.
  TReader const & GetReader() const { return m_reader; }

private:
  void Parse()
  {
    ReaderSource<TReader> src(m_reader);

    m_huffman.ReadEncoding(src);
    m_numNodes = ReadVarUint<uint32_t, ReaderSource<TReader>>(src);

    BitReader<ReaderSource<TReader>> bitReader(src);
    vector<bool> bv(2 * m_numNodes + 1);
    // Convert the level order representation into the external node representation.
    bv[0] = 1;
    for (size_t i = 0; i < 2 * m_numNodes; ++i)
      bv[i + 1] = bitReader.Read(1) == 1;

    succinct::rs_bit_vector(bv).swap(m_trieTopology);

    uint32_t const numFinalNodes = ReadVarUint<uint32_t, ReaderSource<TReader>>(src);
    m_finalNodeIndex.assign(m_numNodes, -1);
    m_offsetTable.resize(numFinalNodes + 1);
    uint32_t id = 0;
    uint32_t offset = 0;
    for (size_t i = 0; i < numFinalNodes; ++i)
    {
      // ids and offsets are delta-encoded
      id += ReadVarUint<uint32_t, ReaderSource<TReader>>(src);
      offset += ReadVarUint<uint32_t, ReaderSource<TReader>>(src);
      m_finalNodeIndex[id] = base::asserted_cast<int>(i);
      m_offsetTable[i] = offset;
    }
    m_offsetTable[numFinalNodes] = base::checked_cast<uint32_t>(src.Size());
    m_reader = m_reader.SubReader(src.Pos(), src.Size());
  }

  TReader m_reader;

  // todo(@pimenov) Why do we even need an instance? Type name is enough.
  TValueReader const & m_valueReader;
  uint32_t m_numNodes;

  coding::HuffmanCoder m_huffman;
  succinct::rs_bit_vector m_trieTopology;

  // m_finalNodeIndex[i] is the 0-based index of the i'th
  // node in the list of all final nodes, or -1 if the
  // node is not final.
  vector<int> m_finalNodeIndex;
  vector<uint32_t> m_offsetTable;
};

template <class TReader, class TValueReader>
class SuccinctTrieIterator
{
public:
  using TValue = typename TValueReader::ValueType;
  using TCommonData = TopologyAndOffsets<TReader, TValueReader>;

  SuccinctTrieIterator(TReader const & reader, shared_ptr<TCommonData> common,
                       uint32_t nodeBitPosition)
    : m_reader(reader), m_common(common), m_nodeBitPosition(nodeBitPosition), m_valuesRead(false)
  {
  }

  unique_ptr<SuccinctTrieIterator> Clone() const
  {
    return make_unique<SuccinctTrieIterator>(m_reader, m_common, m_nodeBitPosition);
  }

  unique_ptr<SuccinctTrieIterator> GoToEdge(size_t i) const
  {
    ASSERT_LESS(i, 2, ("Bad edge id of a binary trie."));
    ASSERT_EQUAL(m_common->GetTopology()[m_nodeBitPosition - 1], 1, (m_nodeBitPosition));

    // rank(x) returns the number of ones in [0, x) but we count bit positions from 1
    uint32_t childBitPosition =
        base::asserted_cast<uint32_t>(2 * m_common->GetTopology().rank(m_nodeBitPosition));
    if (i == 1)
      ++childBitPosition;
    if (childBitPosition > 2 * m_common->NumNodes() ||
        m_common->GetTopology()[childBitPosition - 1] == 0)
    {
      return nullptr;
    }
    return make_unique<SuccinctTrieIterator>(m_reader, m_common, childBitPosition);
  }

  template <typename TEncodingReader>
  unique_ptr<SuccinctTrieIterator> GoToString(TEncodingReader & bitEncoding, uint32_t numBits)
  {
    ReaderSource<TEncodingReader> src(bitEncoding);

    // String size. Useful when the string is not alone in the writer to know where it
    // ends but useless for our purpose here.
    ReadVarUint<uint32_t, ReaderSource<TEncodingReader>>(src);

    BitReader<ReaderSource<TEncodingReader>> bitReader(src);

    auto ret = Clone();
    for (uint32_t i = 0; i < numBits; ++i)
    {
      uint8_t bit = bitReader.Read(1);
      auto nxt = ret->GoToEdge(bit);
      if (!nxt)
        return nullptr;
      ret = move(nxt);
    }
    return ret;
  }

  unique_ptr<SuccinctTrieIterator> GoToString(strings::UniString const & s)
  {
    vector<uint8_t> buf;
    uint32_t numBits;
    {
      MemWriter<vector<uint8_t>> writer(buf);
      numBits = m_common->GetEncoding().EncodeAndWrite(writer, s);
    }
    MemReader reader(&buf[0], buf.size());
    return GoToString(reader, numBits);
  }

  size_t NumValues()
  {
    if (!m_valuesRead)
      ReadValues();
    return m_values.size();
  }

  TValue GetValue(size_t i)
  {
    if (!m_valuesRead)
      ReadValues();
    ASSERT_LESS(i, m_values.size(), ());
    return m_values[i];
  }

private:
  void ReadValues()
  {
    if (m_valuesRead)
      return;
    m_valuesRead = true;
    // Back to 0-based indices.
    uint32_t m_nodeId =
        base::checked_cast<uint32_t>(m_common->GetTopology().rank(m_nodeBitPosition) - 1);
    if (!m_common->NodeIsFinal(m_nodeId))
      return;
    uint32_t offset = m_common->Offset(m_nodeId);
    uint32_t size = m_common->ValueListSize(m_nodeId);
    ReaderPtr<TReader> subReaderPtr(unique_ptr<TReader>(static_cast<TReader*>(m_reader.CreateSubReader(offset, size).release())));
    ReaderSource<ReaderPtr<TReader>> src(subReaderPtr);
    while (src.Size() > 0)
    {
      TValue val;
      m_common->GetValueReader()(src, val);
      m_values.push_back(val);
    }
  }

  TReader const & m_reader;
  shared_ptr<TCommonData> m_common;

  // The bit with this 1-based index represents this node
  // in the external node representation of binary trie.
  uint32_t m_nodeBitPosition;

  vector<TValue> m_values;
  bool m_valuesRead;
};

template <typename TReader, typename TValueReader>
unique_ptr<SuccinctTrieIterator<TReader, TValueReader>> ReadSuccinctTrie(
    TReader const & reader, TValueReader valueReader = TValueReader())
{
  using TCommonData = TopologyAndOffsets<TReader, TValueReader>;
  using TIter = SuccinctTrieIterator<TReader, TValueReader>;

  shared_ptr<TCommonData> common(new TCommonData(reader, valueReader));
  return make_unique<TIter>(common->GetReader(), common, 1 /* bitPosition */);
}

}  // namespace trie