#pragma once

#include "archive_config.h"

#include <mapreduce/yt/interface/io.h>
#include <kernel/multipart_archive/multipart.h>

namespace NRTYT {

class TProtoWriter : public NYT::IProtoWriterImpl {
public:
    TProtoWriter(const TVector<TString>& paths, const TVector<const ::google::protobuf::Descriptor*>& descriptors)
        : ArchivePaths(paths)
        , Descriptors(descriptors) {
        NRTYArchive::TMultipartConfig config = DefaultArchiveConfig;
        for (const auto& path : paths) {
            Archives.push_back(TArchiveOwner::Create(path, config));
            RowIds.push_back(Archives.back()->GetDocsCount(/*withDeleted =*/ true));
        }
    }

    TProtoWriter(const TString& path, const ::google::protobuf::Descriptor* descriptor)
        : TProtoWriter(
            TVector<TString>{path},
            TVector<const ::google::protobuf::Descriptor*>{descriptor}
        )
        {
    }

    virtual void AddRow(const NYT::Message& row, size_t tableIndex) override;
    virtual void AddRow(NYT::Message&& row, size_t tableIndex) override;

public: // IWriterImplBase
    size_t GetTableCount() const override;
    void FinishTable(size_t) override;

private:
    TVector<TArchiveOwner::TPtr> Archives;
    TVector<TString> ArchivePaths;
    
    TVector< const ::google::protobuf::Descriptor*> Descriptors;
    TVector<ui32> RowIds;
};

} // namespace NRTYT
