#pragma once

#include <mapreduce/yt/interface/client.h>
#include <saas/library/yt/common/yt_blob.h>

#include <library/cpp/streams/special/buffered_throttled_file.h>
#include <util/folder/path.h>
#include <util/generic/map.h>
#include <util/string/printf.h>
#include <util/generic/yexception.h>
#include <util/stream/file.h>
#include <util/system/file.h>
#include <util/system/yassert.h>


class IYTChunkedInputProcessor {
public:
    using TTableReader = NYT::TTableReader<NYT::Message>;

public:
    IYTChunkedInputProcessor(TTableReader* tableReader, EYTBlobFormat format = EYTBlobFormat::Old)
        : TableReader(tableReader)
        , NamedConsumersMap()
        , CurrentConsumer(NamedConsumersMap.end())
        , Format(format)
        , State(EState::Started)
    {
        Y_ASSERT(TableReader);
    }
    virtual ~IYTChunkedInputProcessor();
    void ProcessAll();

    TShardId GetShardId() const {
        return ShardId;
    }

    void Finish();

protected:
    virtual THolder<IOutputStream> CreateOutputStream(TString name) = 0;
    virtual THolder<IOutputStream> ResumeOutputStream(TString name, i64 offset) {
        Y_UNUSED(name);
        Y_UNUSED(offset);
        ythrow yexception() << "ResumeOutputStream not supported";
    }

private:
    enum EState {
        Started,
        Finished,
        Failed
    };

    class TChunkConsumer {
    public:
        TChunkConsumer(THolder<IOutputStream> output)
            : Output(std::move(output))
            , NextChunkId(0)
            , Offset(0)
        {
        }

        void ConsumeChunk(ui32 chunkId, TString data) {
            Y_ENSURE(chunkId == NextChunkId);
            Y_ENSURE(Output);
            Output->Write(data.data(), data.size());
            NextChunkId++;
            Offset += data.size();
        }

        void Close() {
            Y_ENSURE(Output);
            // NB: If you remove the following line (explicit call of Finish) then ChunkedInputFailedFlushSuite
            // and ChunkedInputFailedFinishSuite should fail because destructor hides errors.
            Output->Finish();
            Output.Destroy();
        }

        bool IsClosed() const {
            return !Output;
        }

        void Resume(THolder<IOutputStream> output) {
            Y_ENSURE(!Output);
            Output = std::move(output);
        }

        i64 GetOffset() const {
            return Offset;
        }

    private:
        THolder<IOutputStream> Output;
        ui32 NextChunkId;
        i64 Offset;
    };

private:
    void ProcessRow(const TYTBlob& row);
    TChunkConsumer& GetConsumerByName(const TString& name);
    void DoProcessAll();

private:
    TTableReader* const TableReader;
    TShardId ShardId;
    TMap<TString, TChunkConsumer> NamedConsumersMap;
    TMap<TString, TChunkConsumer>::iterator CurrentConsumer;
    EYTBlobFormat Format = EYTBlobFormat::Old;
    EState State;
};

class TYTChunkedInputToFiles : public IYTChunkedInputProcessor {
public:
    TYTChunkedInputToFiles(TTableReader* tableReader, TFsPath root = TFsPath(), EYTBlobFormat format = EYTBlobFormat::Old)
        : IYTChunkedInputProcessor(tableReader, format)
        , Root(root)
    {
    }
    ~TYTChunkedInputToFiles() = default;

protected:
    THolder<IOutputStream> CreateOutputStream(TString name) override {
        TFsPath outPath = Root / name;
        outPath.Parent().MkDirs();
        return THolder(new TUnbufferedFileOutput(outPath));
    }

    THolder<IOutputStream> ResumeOutputStream(TString name, i64 offset) override {
        TFsPath outPath = Root / name;
        TFile f(outPath, OpenExisting | ForAppend | WrOnly | Seq);
        Y_ENSURE(f.GetLength() == offset, Sprintf("error in ResumeOutputStream: %ld != %ld", f.GetLength(), offset));
        return THolder(new TUnbufferedFileOutput(f));
    }

private:
    TFsPath Root;
};

class TYTChunkedInputToThrottledFiles : public IYTChunkedInputProcessor {
public:
    TYTChunkedInputToThrottledFiles(TTableReader* tableReader, ui32 maxBytesInSeconds, TFsPath root = TFsPath(), EYTBlobFormat format = EYTBlobFormat::Old)
        : IYTChunkedInputProcessor(tableReader, format)
        , WriteBytesPerSec(maxBytesInSeconds)
        , Root(root)
    {
    }
    ~TYTChunkedInputToThrottledFiles() = default;

protected:
    THolder<IOutputStream> CreateOutputStream(TString name) override {
        TFsPath outPath = Root / name;
        outPath.Parent().MkDirs();
        return THolder(new TBufferedThrottledFileOutputStream(outPath, WriteBytesPerSec));
    }

private:
    ui32 WriteBytesPerSec;
    TFsPath Root;
};
