#pragma once

#include "range_comparators.h"
#include "proto_table_reader.h"
#include <google/protobuf/descriptor.h>
#include <saas/library/rtyt/lib/util/proto_utils.h>
#include <mapreduce/yt/interface/io.h>
#include <kernel/multipart_archive/multipart.h>

namespace NRTYT {

class TMultiProtoReader : public NYT::IProtoReaderImpl {
public:
    TMultiProtoReader(
        const TVector<TString>& paths,
        const TVector<const ::google::protobuf::Descriptor*>& descriptors,
        const NYT::TSortColumns groupedColumns = {},
        const NYT::TSortColumns sortedColumns = {})
            : MessageFactory(new ::google::protobuf::DynamicMessageFactory())
            , RowComparator(descriptors, sortedColumns.EnsureAscending().GetNames(), /*CompareIndexes*/true)
            , RangesComparator(descriptors, groupedColumns.EnsureAscending().GetNames(), /*CompareIndexes*/false)
            , RowId(0)
            , RangeId(0)
            , TableIndex(0)
            , RangeSwitched(false) {
        for (ui32 i = 0; i < paths.size(); i++) {
            Readers.emplace_back(paths[i], descriptors[i], groupedColumns);
            MessageBuffer.emplace_back(MessageFactory->GetPrototype(descriptors[i])->New());
            if (Readers.back().IsValid()) {
                Readers[i].ReadRow(MessageBuffer[i].Get());
            }
        }
        UpdateTableIndex();
    }

    virtual void ReadRow(NYT::Message* row) override;

public: // IReaderImplBase
    virtual bool IsValid() const override;
    virtual void Next() override;
    virtual ui32 GetTableIndex() const override;
    virtual ui32 GetRangeIndex() const override;
    virtual ui64 GetRowIndex() const override;
    virtual void NextKey() override;

private:
    void UpdateTableIndex();
    /*
    * We suggest that TableIndex always points on reader we can read from
    */
    TAtomicSharedPtr<::google::protobuf::DynamicMessageFactory> MessageFactory;
    TVector<NRTYT::TProtoReader> Readers;
    TVector<TAtomicSharedPtr<::google::protobuf::Message>> MessageBuffer;
    TProtoRangesComparator RowComparator;
    TProtoRangesComparator RangesComparator;
    ui64 RowId;
    ui32 RangeId;
    ui32 TableIndex;
    bool RangeSwitched;
};

} // namespace NRTYT
