#pragma once

#include "archive_config.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>

#include <google/protobuf/dynamic_message.h>

namespace NRTYT {

class TProtoReader : public NYT::IProtoReaderImpl {
public:
    using TReaderType = NYT::IProtoReaderImpl;

    TProtoReader(const TString& path,
                const ::google::protobuf::Descriptor* descriptor,
                const TMaybe<NYT::TSortColumns>& groupedColumns = Nothing())
        : ArchivePath(path)
        , MessageFactory(new ::google::protobuf::DynamicMessageFactory())
        , Descriptor(descriptor)
        , RowId(0)
        , RangeId(0)
        , RangeSwitched(false) {
        NRTYArchive::TMultipartConfig config = DefaultArchiveConfig;
        Archive = TArchiveOwner::Create(path, config);
        ReadFromArchive();

        if (groupedColumns.Defined()) {
            CurrentRange = CopyMessage(CurrentRow);
            KeyColumns = GetFieldDescriptorsByName(
                Descriptor,
                groupedColumns->EnsureAscending().GetNames()
            );
        }
    }

    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 ReadFromArchive();
    bool CheckCurrentRange() const;

private:
    TArchiveOwner::TPtr Archive;
    TString ArchivePath;
    TAtomicSharedPtr<::google::protobuf::DynamicMessageFactory> MessageFactory;
    
    const ::google::protobuf::Descriptor* Descriptor;
    TMaybe<TVector<const ::google::protobuf::FieldDescriptor*>> KeyColumns;
    THolder<NYT::Message> CurrentRow;
    THolder<NYT::Message> CurrentRange;
    ui32 RowId;
    ui32 RangeId;
    bool RangeSwitched;
};

} // namespace NRTYT
