#include <saas/tools/patch_dh_requests/params.pb.h>

#include <search/idl/meta.pb.h>
#include <search/session/compression/report.h>

#include <kernel/urlid/doc_handle.h>

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

#include <util/datetime/base.h>
#include <util/generic/array_ref.h>
#include <util/stream/buffer.h>
#include <util/stream/file.h>
#include <util/stream/null.h>
#include <util/string/split.h>

TVector<TString> ReadLines(const TString& path) {
    if (path.empty()) {
        return {};
    }
    auto text = TFileInput(path).ReadAll();
    return StringSplitter(text).Split('\n').SkipEmpty().ToList<TString>();
}

TStringBuf Shorten(TStringBuf s, size_t maxLen) {
    return s.substr(0, maxLen);
}

TStringBuf ShortReq(TStringBuf s) {
    return Shorten(s, 60);
}

TString SendNehRequest(const TString& request) {
    NNeh::TResponseRef resp = NNeh::Request(request)->Wait(TDuration::Seconds(10));
    Y_ENSURE(resp, "Got no response");
    Y_ENSURE(!resp->IsError(), "Got an error response: " << resp->GetErrorText());
    return resp->Data;
}

TVector<TString> SendSearchRequest(const TString& hostWithPath, const TString& request) {
    TVector<TString> result;
    try {
        Y_ENSURE(request.StartsWith('?'));

        auto response = SendNehRequest(hostWithPath + request);
        if (response.empty()) {
            return result;
        }

        NMetaProtocol::TReport report;
        Y_VERIFY(report.ParseFromString(response));
        NMetaProtocol::Decompress(report);
        for (auto& grouping : report.GetGrouping()) {
            for (auto& group : grouping.GetGroup()) {
                for (auto& doc : group.GetDocument()) {
                    if (doc.GetDocId().empty()) {
                        result.push_back(TDocHandle(doc.GetDocHash(), doc.GetRoute()).ToString());
                    } else {
                        result.push_back(doc.GetDocId());
                    }
                }
            }
        }

    } catch (...) {
        ERROR_LOG << "Cannot send request " << ShortReq(request) << ". Error: " << CurrentExceptionMessage() << Endl;
    }
    return result;
}

THolder<IOutputStream> CreateFileOrNull(const TString& path) {
    if (path.empty()) {
        return MakeHolder<TNullOutput>();
    }
    return MakeHolder<TFileOutput>(path);
}

bool PatchDocids(IOutputStream& output, TStringBuf request, TArrayRef<const TString> docIds, i32 extraDocs) {
    Y_ENSURE(request.StartsWith('?'));
    TBufferStream buffer;
    buffer << '?';
    i32 srcCount = 0;
    for (auto& item : StringSplitter(request.After('?')).Split('&').SkipEmpty()) {
        if (!item.Token().StartsWith(TStringBuf("dh="))) {
            buffer << '&' << item.Token();
        } else {
            ++srcCount;
        }
    }
    Y_ENSURE(docIds.size() < static_cast<size_t>(Max<i32>()));
    const i32 nReplacements = static_cast<i32>(docIds.size());
    const i32 dstCount = Min(srcCount + extraDocs, nReplacements);
    if (dstCount <= 0) {
        return 0;
    }
    buffer.ReadAll(output);
    for (i32 i = 0; i < dstCount; ++i) {
        output << "&dh=" << docIds[i] << ':';
    }
    return true;
}

struct TRequestPatcher {
    TRequestPatcher(const TString& input, const TString& output, const TString& stageName, bool writeEmptyLines, i32 extraDocs)
        : WriteEmptyLines(writeEmptyLines)
        , ExtraDocs(extraDocs)
    {
        Y_ENSURE(!input.empty() || output.empty(), "cannot process " << stageName << " queries without input file specified");
        InputRequests = ReadLines(input);
        Output = CreateFileOrNull(output);
    }

    void Process(size_t i, TArrayRef<const TString> docIds) {
        if (i >= InputRequests.size()) {
            return;
        }
        const bool patched = PatchDocids(*Output, InputRequests[i], docIds, ExtraDocs);
        if (patched || WriteEmptyLines) {
            *Output << '\n';
        }
    }

private:
    const bool WriteEmptyLines;
    const i32 ExtraDocs;
    TVector<TString> InputRequests;
    THolder<IOutputStream> Output;
};

int main(int argc, const char* argv[]) {
    NGetoptPb::TGetoptPbSettings getoptSettings;
    getoptSettings.DumpConfig = false;
    const NUtils::TParams params = NGetoptPb::GetoptPbOrAbort(argc, argv, getoptSettings);

    InitGlobalLog2Console(params.GetVerbose() ? TLOG_DEBUG : TLOG_INFO);
    try {
        auto searchRequests = ReadLines(params.GetSearchInput());

        TRequestPatcher factorPatcher(params.GetFactorInput(), params.GetFactorOutput(), "factor", params.GetWriteEmptyLines(), params.GetExtraDocs());
        TRequestPatcher snippetPatcher(params.GetSnippetInput(), params.GetSnippetOutput(), "snippet", params.GetWriteEmptyLines(), params.GetExtraDocs());

        const auto hostWithPath = TString::Join("http2://", params.GetBasesearch(), "/", params.GetCollection());

        for (size_t i = 0; i < searchRequests.size(); ++i) {
            DEBUG_LOG << '#' << i << " " << ShortReq(searchRequests[i]) << Endl;
            auto docIds = SendSearchRequest(hostWithPath, searchRequests[i]);
            DEBUG_LOG << "Got " << docIds.size() << " documents" << Endl;
            factorPatcher.Process(i, docIds);
            snippetPatcher.Process(i, docIds);
        }

        return 0;

    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
        return 1;
    }
}
