#include "ops_common.h"

#include <saas/tools/rty_ops/lib/args.pb.h>

#include <kernel/factor_storage/factors_reader.h>
#include <kernel/urlid/doc_handle.h>

#include <library/cpp/http/io/stream.h>

#include <quality/query_pool/prs_ops/lib/parsed_report.h>

#include <search/idl/meta.pb.h>

#include <util/generic/hash.h>
#include <util/network/socket.h>
#include <util/stream/file.h>

using namespace NSaas;

namespace NRTYServer {

class TReportParserForPool {
private:
    const NMetaProtocol::TReport& MetaReport;
    TIntrusivePtr<TParsedReport> Result;
    TGatheredDocumentPtr Doc;
    bool HasFactors = true;

public:
    TReportParserForPool(TRequestId qid, const NMetaProtocol::TReport& metaReport)
        : MetaReport(metaReport)
        , Result(MakeIntrusive<TParsedReport>())
    {
        Result->Qid = qid;
    }

    static TString GetUrl(const NMetaProtocol::TDocument& document) {
        if (document.HasUrl()) {
            return document.GetUrl();
        } else if (document.GetArchiveInfo().HasUrl()) {
            return document.GetArchiveInfo().GetUrl();
        } else {
            for (const auto& attrib : document.GetFirstStageAttribute()) {
                if (attrib.GetKey() == "_Url") {
                    const auto& attribValue = attrib.GetValue();
                    return attribValue;
                }
            }
            return TString();
        }
    }

    void BeginDoc() {
        Doc = new TGatheredDocument;
        HasFactors = false;
    }

    void OnError(const TString msg) {
        ERROR_LOG << "Failed to parse:" << msg << Endl;
    }

    static void SetFactorsFromCompressed(NFeaturePool::TLine& dest, const TString& compressed) {
        TString data;
        Base64Decode(compressed, data);
        TStringInput inp_stream(data);
        TVector<ui8> tmpBuf;

        auto reader = NFSSaveLoad::CreateCompressedFactorsReader(&inp_stream, tmpBuf);
        TFactorStorage storage;
        reader->ReadTo(storage);

        TFactorView view = storage.CreateViewFor(EFactorSlice::ALL);
        dest.SetFactorSlices(SerializeFactorBorders(storage.GetBorders(), NFactorSlices::ESerializationMode::LeafOnly));
        for (size_t i = 0; i < view.Size(); ++i) {
            const auto& feature = view[i];
            if (!IsValidFloat(feature)) {
                dest.AddFeature(-1000);
            } else {
                dest.AddFeature(feature);
            }
        }
    }

    void OnGta(const NMetaProtocol::TPairBytesBytes& attr) {
        auto& line = Doc->Document;
        const auto& key = attr.GetKey();
        if (!HasFactors && key == "_CompressedAllFactors") {
            SetFactorsFromCompressed(line, attr.GetValue());
        } else if (key == "_DocOmniTitle" || key == "_DocZoneTitle") {
            line.SetOmniTitle(attr.GetValue());
        } else if (key == "_DocOmniUrl") {
            line.SetOmniUrl(attr.GetValue());
        }
    }

    TIntrusivePtr<TParsedReport> ParseReport() {
        // this is simplfied from TPRSHarvester::ReportToParsedReport
        Result->UnanswersCount = MetaReport.GetDebugInfo().GetBaseSearchNotRespondCount();
        Result->SearchersCount = MetaReport.GetDebugInfo().GetBaseSearchCount();

        size_t groupId = 0; // counter
        for (const auto& grouping : MetaReport.GetGrouping()) {
            for (const auto& group : grouping.GetGroup()) {
                for (const auto& document : group.GetDocument()) {
                    const TString url = GetUrl(document);
                    if (!url) {
                        OnError("Document has no url");
                        continue;
                    }

                    BeginDoc();
                    for (const auto& attrib : document.GetArchiveInfo().GetGtaRelatedAttribute()) {
                        OnGta(attrib);
                    }
                    for (const auto& attrib : document.GetFirstStageAttribute()) {
                        OnGta(attrib);
                    }

                    {
                        auto& line = Doc->Document;
                        line.SetRequestId(Result->Qid);
                        line.SetMainRobotUrl(url);
                        line.SetRelevance(document.GetSRelevance());
                        line.SetFiltrationType(NFeaturePool::FT_NOT_FILTERED);
                        line.SetGrouping(ToString(groupId));
                    }

                    Result->Documents.emplace_back(std::move(Doc));
                }
                groupId++;
            }
        }

        return std::move(Result);
    }
};

//
// TPoolOp
//
class TPoolOp: public TQueryOpBase {
protected:
    NRTYServer::TParamsOpPool Params;

public:
    using TBase = TQueryOpBase;

public:
    void Init(IRtyOpHost& host) override {
        TBase::Init(host);
        Params = GetExtraParams<NRTYServer::TParamsOpPool>(host.GetConfig().VirtualArgs);
        Y_ENSURE(!Params.HasExtraCgi() || Params.GetExtraCgi().StartsWith("&"));
    }

    void Run(IRtyOpHost& host) override {
        TVector<TDolbilkaPlan::TQueryData> plan;
        LoadPlan(plan, host.GetConfig(), Params.GetMaxQueries());
        PatchQueries(plan, Params.GetExtraCgi() + "&waitall=da&timeout=10000000&fsgta=_CompressedAllFactors&gta=_Url&fsgta=_Url");

        for (const auto& req : plan) {
            TString result;
            const ui32 httpCode = DoQuery(req.Query, &result);
            Y_ENSURE(httpCode == 200);

            NMetaProtocol::TReport report;

            const bool parsed = report.ParseFromString(result) && !report.HasCompressedReport();
            Y_ENSURE(parsed, "failed to parse protobin answer for qid=" << req.RequestId);

            TIntrusivePtr<TParsedReport> processedReport = TReportParserForPool(req.RequestId, report).ParseReport();

            PrintTsvLines(*processedReport);
        }
    }

    void PrintTsvLines(const TParsedReport& processedReport) {
        TStringStream m;
        for (const auto& doc : processedReport.Documents) {
            const auto& line = doc->Document;
            m << line.GetRequestId() << '\t'
              << line.GetRating() << '\t'
              << line.GetMainRobotUrl() << '\t'
              << '\t';

            for (size_t i = 0; i != line.FeatureSize(); ++i) {
                m << '\t' << line.GetFeature(i);
            }

            Cout << m.Str() << Endl;
            m.Clear();
        }
    }

private:
    static IRtyOp::TFactory::TRegistrator<TPoolOp> Registrator;
};

IRtyOp::TFactory::TRegistrator<TPoolOp> TPoolOp::Registrator("Pool");

}

namespace NRTYOpsApp {
    void RegisterModulePool() {
        // this call ensures that TRegistrators in the classes here are called
    }
}
