#include "arc_merge.h"

#include "docid_remapper.h"
#include "pudge.h"
#include "arc_verify.h"

#include <kernel/doom/wad/mega_wad_writer.h>
#include <kernel/tarc/iface/tarcio.h>
#include <kernel/tarc/merge/builder.h>

#include <library/cpp/logger/global/global.h>

#include <robot/jupiter/tools/shardmerge_utils/lib/document_index/arc_merge.h>
#include <robot/jupiter/library/pudge/config/tier_config.h>
#include <robot/jupiter/tools/shardmerge_utils/lib/document_index/chunked_index_writer.h>
#include <robot/jupiter/tools/shardmerge_utils/lib/document_index/arc_chunk_traits.h>
#include <kernel/doom/chunked_wad/doc_chunk_mapping_writer.h>

namespace NRtDoc {
    const TStringBuf ARC_SUFFIX = "arc";
    const TStringBuf ARC_WAD_SUFFIX = "arc.wad";

    namespace {
        const TStringBuf TMP_SUFFIX = "_tmp";

        TStringBuf ExtractArcPrefix(TStringBuf path) {
            Y_ENSURE(path.EndsWith(ARC_SUFFIX));
            return path.substr(0, path.size() - ARC_SUFFIX.size());
        }

        THolder<ITextArchiveBuilder> CreateArchiveBuilder(TStringBuf arcPath) {
            const TString pathPrefix = TString{ExtractArcPrefix(arcPath)};
            ITextArchiveBuilder::TConstructContext builderCtx;
            builderCtx.DirConfig.TempPrefix = pathPrefix + TMP_SUFFIX;
            builderCtx.DirConfig.NewPrefix = pathPrefix;
            THolder<ITextArchiveBuilder> builder = THolder(ITextArchiveBuilder::TFactory::Construct(EArchiveType::AT_FLAT, builderCtx));
            builder->Start();
            return builder;
        }
    }

    TArcMerger::TArcMerger(const TString& fileName, TString mappingFileName)
        : MappingFileName(std::move(mappingFileName))
    {
        Init(fileName);
    }

    void TArcMerger::Init(const TString& fileName) {
        FileName = fileName;
    }

    void TArcMerger::Add(const TString& wadPath, THolder<TDocIdMap> docIdMap, IWadPatcher::TPtr, const TString& /*name*/, bool /*isDelta*/) {
        Chunks.emplace_back(wadPath, std::move(docIdMap));
    }

    bool TArcMerger::Empty() const {
        return Chunks.empty();
    }

    void TArcMerger::CheckMergeable() {
    }

    void TArcMerger::Finish() {
        Y_ENSURE(FileName);
        Y_ENSURE(!Chunks.empty(), "No input files");

        TDocIdRemapper docidRemapper(MappingFileName, Chunks.size());

        TVector<THolder<TArchiveIterator>> inputs;
        for (const auto& chunk : Chunks) {
            inputs.emplace_back(MakeHolder<TArchiveIterator>());
            inputs.back()->Open(chunk.Path.c_str());
            docidRemapper.AddSegmentSize(inputs.back()->Size());
        }
        docidRemapper.InitFinalDocIds();

        auto writer = CreateArchiveBuilder(FileName);
        ui32 nextDocId = 0;
        for (size_t i = 0; i < Chunks.size(); ++i) {
            docidRemapper.ProcessSegment(*Chunks[i].DocIdMap, [&](const ui32 docId, const ui32 outputDocId) {
                TArchiveHeader* doc = inputs[i]->SeekToDoc(docId);
                writer->IndexDoc(outputDocId, TBlob::NoCopy(reinterpret_cast<const char*>(doc), doc->DocLen), 0);
                nextDocId = outputDocId + 1;
            });
        }

        writer->Stop();
        writer->Close(nullptr);

        docidRemapper.Finish();
    }

    void ConvertArcToWad(const TString& arcPath, const TString& wadPath, NDoom::TWadLumpId lump) {
        INFO_LOG << "Restore " << arcPath << " from " << wadPath << Endl;
        TArchiveIterator input;
        input.Open(arcPath.c_str());
        NDoom::TMegaWadWriter output(wadPath);
        output.RegisterDocLumpType(lump);
        for (auto item = input.NextAuto(); item != nullptr; item = input.NextAuto()) {
            Y_ENSURE(item->DocLen >= sizeof(*item));
            auto data = reinterpret_cast<const char*>(item);
            output.StartDocLump(item->DocId, lump)->Write(data + sizeof(*item), item->DocLen - sizeof(*item));
        }
        output.Finish();
    }

    void ConvertArcToChunkedWad(const TString& arcPath, const TString& outputDir, const TVector<NDoom::TDocChunkMapping>& docChunkMapping, const TString& filePrefix) {
        INFO_LOG << "Restore chunked arc wad from " << arcPath << Endl;

        TArchiveIterator input;
        input.Open(arcPath.c_str());

        using TTraits = TWadChunkTraits;
        using ArcWriter = typename TTraits::TArcChunkWriter;
        using ArcMover = typename TTraits::TArcChunkMover;

        using TChunkedArcWriter = NJupiter::TChunkedIndexWriter<ArcWriter, ArcMover>;

        INFO_LOG << "Build full chunked text arc" << Endl;
        const NJupiter::TTierConfig tierConfig = NFusion::TPudgeData::LoadTierConfig();

        INFO_LOG << "Write docs" << Endl;

        const TString outputStorageDir = outputDir + "/" + NJupiter::TArcMergeOpts::RemoteStorageDirName;
        TChunkedArcWriter writer(
            outputStorageDir + ".tmp",
            outputStorageDir,
            filePrefix,
            tierConfig.GetChunks(),
            &docChunkMapping);

        TBuffer buffer;

        for (auto item = input.NextAuto(); item != nullptr; item = input.NextAuto()) {
            Y_ENSURE(item->DocLen >= sizeof(*item));
            auto data = reinterpret_cast<const char*>(item);

            auto text = TBlob::NoCopy(data + sizeof(*item), item->DocLen - sizeof(*item));

            writer.Write(item->DocId, [&](ui32 localDocId, ArcWriter* chunkWriter) {
                chunkWriter->WriteDoc(localDocId, text);
            });
        }

        writer.Finish();

        for (size_t i = 0; ; ++i) {
            const TFsPath treePath = TStringBuilder() << outputStorageDir << "/" << i << "/" << filePrefix << "arc.wad";
            if (!NFs::Exists(treePath)) {
                break;
            }
            const TString plainPath = TStringBuilder() << outputDir << "/" << filePrefix << "arc." << i << ".wad";
            INFO_LOG << "output link " << treePath << ", " << plainPath << Endl;
            Y_ENSURE(NFs::HardLink(treePath, plainPath));
        }

        NFs::RemoveRecursive(outputStorageDir);
        NFs::RemoveRecursive(outputStorageDir + ".tmp");

        // dummy
        NDoom::TMegaWadWriter(TFsPath(outputDir) / (filePrefix + "arc.global.wad")).Finish();

        INFO_LOG << "Write doc chunk mapping" << Endl;
        {
            const auto path = TFsPath(outputDir) / (filePrefix + "arc.mapping.wad");
            NDoom::TDocChunkMappingWriter writer(path.GetPath());
            for (const NDoom::TDocChunkMapping& mapping : docChunkMapping) {
                writer.Write(mapping);
            }
            writer.Finish();
        }
        INFO_LOG << "DONE!" << Endl;
    }

}

namespace NFusion {
    void TArcWadExtension::PreprocessInputs(const NRtDoc::IBuilderInputs& inputs) {
        for (size_t i = 0; i < inputs.InputsSize(); ++i) {
            auto input = inputs.GetInputs(i);
            TFsPath srcDir(input.GetSrcDir());
            auto indexPath = srcDir / ConstructFileName(OutputPrefix);
            if (input.GetIsFinalIndex() && !indexPath.Exists()) {
                auto arcPath = srcDir / OldName;
                Y_ENSURE(arcPath.Exists());
                NRtDoc::ConvertArcToWad(arcPath, indexPath, LumpId);
            }
        }
    }

    void TChunkedArcWadExtension::PreprocessInputs(const NRtDoc::IBuilderInputs& inputs) {
        for (size_t i = 0; i < inputs.InputsSize(); ++i) {
            const auto& input = inputs.GetInputs(i);
            TFsPath srcDir(input.GetSrcDir());
            auto indexPath = srcDir / "indexarc.global.wad";
            if (input.GetIsFinalIndex() && !indexPath.Exists()) {
                auto arcPath = srcDir / (OutputPrefix + "arc");
                Y_ENSURE(arcPath.Exists());
                NRtDoc::ConvertArcToChunkedWad(arcPath, srcDir, NFusion::TPudgeData::LoadDocChunkMapping(srcDir), OutputPrefix);
                NFusion::VerifyPlainArcEqualToChunked(arcPath, srcDir, OutputPrefix);
            }
        }
    }
}

