#pragma once

#include "capped_buffered_output.h"

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

#include <util/generic/maybe.h>
#include <util/stream/output.h>
#include <util/system/rwlock.h>
#include <util/system/yassert.h>


class INamedOutput : public IOutputStream {
public:
    INamedOutput(TString name)
        : Name(std::move(name))
    {
    }
    virtual ~INamedOutput() = default;

    const TString& GetName() const {
        return Name;
    }
private:
    const TString Name;
};

class IChunkedOutput : public INamedOutput {
public:
    IChunkedOutput(TString name)
        : INamedOutput(std::move(name))
        , CurrentChunkId(0)
    {
    }
    virtual ~IChunkedOutput() = default;

protected:
    void DoWrite(const void* buf, size_t len) override final {
        DoWriteChunk(buf, len, CurrentChunkId++, GetName());
    }

    void DoFinish() override final {
        if (CurrentChunkId == 0) {
            DoWriteChunk(nullptr, 0, CurrentChunkId++, GetName());
        }
        INamedOutput::DoFinish();
    }

    virtual void DoWriteChunk(const void* buf, size_t len, ui32 chunkId, const TString& name) = 0;

private:
    ui32 CurrentChunkId;
};

class TYTChunkedOutputBase : public IChunkedOutput {
public:
    TYTChunkedOutputBase(TString name, NYT::TTableWriter<google::protobuf::Message>* tableWriter, TShardId shardId,
                        ui32 segmentId = 0, size_t tableIndex = 0, EYTBlobFormat format = EYTBlobFormat::Old)
        : IChunkedOutput(std::move(name))
        , TableWriter(tableWriter)
        , ShardId(shardId)
        , SegmentId(segmentId)
        , TableIndex(tableIndex)
        , Format(format)
    {
        Y_ASSERT(TableWriter);
    }

protected:
    virtual void DoWriteChunk(const void* buf, size_t len, ui32 chunkId, const TString& name) override;

private:
    NYT::TTableWriter<google::protobuf::Message>* const TableWriter;
    const TShardId ShardId;
    const ui32 SegmentId;
    const size_t TableIndex;
    const EYTBlobFormat Format;
};


class TYTChunkedOutput : public TCappedBufferedOutput<TYTChunkedOutputBase> {
    using TBase = TCappedBufferedOutput<TYTChunkedOutputBase>;

public:
    constexpr static size_t DefaultChunkSize = 8ul * 1024 * 1024;
    struct TConstructionContext {
        NYT::TTableWriter<google::protobuf::Message>* TableWriter = nullptr;
        TShardId ShardId = TShardId();
        size_t TableIndex = 0;
        TMaybe<ui32> SegmentId = Nothing();
        size_t ChunkSize = DefaultChunkSize;
        EYTBlobFormat Format = EYTBlobFormat::Old;
    };

public:
    using TBase::TBase;
    TYTChunkedOutput(const TConstructionContext& context, TString name)
        : TBase(context.ChunkSize, std::move(name), context.TableWriter, context.ShardId, context.SegmentId.GetOrElse(0), context.TableIndex, context.Format)
    {
    }
};

class TYTCommonChunkedOutput : private TReadGuard, public TYTChunkedOutput {
public:
    static void SetCommonContext(const TConstructionContext& context) {
        TWriteGuard g(CommonContextMutex);
        CommonContext = context;
    }

    static TConstructionContext GetCommonContext() {
        TReadGuard g(CommonContextMutex);
        return CommonContext;
    }

    TYTCommonChunkedOutput(TString Name)
        : TReadGuard(CommonContextMutex)
        , TYTChunkedOutput(CommonContext, std::move(Name))
    {
        TReadGuard::Release();
    }

private:
    static TConstructionContext CommonContext;
    static TRWMutex CommonContextMutex;
};

// this class is used for debugging, e.g. to dump data into Cout or TStringOutput
class TTextChunkedOutputBase : public IChunkedOutput {
public:
    TTextChunkedOutputBase(IOutputStream* slave, TString name, bool hex = true)
        : IChunkedOutput(std::move(name))
        , Slave(slave)
        , Hex(hex)
    {
    }

protected:
    void DoWriteChunk(const void* buf, size_t len, ui32 chunkId, const TString& name) override;

private:
    IOutputStream* Slave;
    bool Hex;
};

using TTextChunkedOutput = TCappedBufferedOutput<TTextChunkedOutputBase>;

// Drop-in debug replacement class for TYTCommonChunkedOutput
class TCoutChunkedOutput : public TTextChunkedOutput {
    static const size_t DefaultChunkSize = 32;

public:
    TCoutChunkedOutput(TString name)
        : TTextChunkedOutput(DefaultChunkSize, &Cout, name)
    {
    }
};
