#include "fa_dumper.h"
#include "dump_kiwi_object.h"

#include <saas/library/hash_to_docid/hash_to_docid.h> // NSaas::IDocHashStorage
#include <kernel/urlid/urlid.h>
#include <saas/api/clientapi.h>
#include <yweb/protos/indexeddoc.pb.h>

#include <library/cpp/getopt/last_getopt.h>
#include <library/cpp/getopt/modchooser.h>
#include <library/cpp/logger/global/global.h>
#include <library/cpp/sighandler/async_signals_handler.h>
#include <library/cpp/protobuf/json/proto2json.h>
#include <library/cpp/json/json_writer.h>


#include <google/protobuf/text_format.h>
#include <util/generic/noncopyable.h>
#include <util/generic/hash.h>
#include <util/stream/file.h>
#include <util/string/cast.h>
#include <saas/library/hash_to_docid/hash_to_docid.h>
using namespace NSaas;
using namespace NRTYServer;
using namespace NLastGetopt;


namespace {
    class IAction : public TFADumper::THandler {
    public:
        virtual void AddOpts(TOpts& /*opts*/) {
        }

        virtual void Init(TOptsParseResult& /*args*/) {
        }
    };

    using TId2Hash = THashMap<ui32, NSaas::TDocHash>;
    THolder<TId2Hash> LoadDocHashMapping(TString path) {
        THolder<TId2Hash> result;
        THolder<THashToDocId> storage = MakeHolder<THashToDocId>(/*useGlobalMapping=*/false);
        if (!storage->Open(TFsPath(path) / "url2docid"))
            return result;

        result = MakeHolder<TId2Hash>();
        for(auto iter = storage->CreateDocHashIterator(); iter->IsValid(); iter->Next()) {
            auto kv = iter->GetHashWithId();
            result->emplace(kv.second, kv.first);
        }
        if (result->empty())
            result.Destroy();

        storage->Close();
        return result;
    }

    class TPrintUrlAction final : public IAction {
    public:
        void AddOpts(TOpts& opts) override {
            opts.AddCharOption('D', "print numerical docids").Optional().NoArgument();
            opts.AddCharOption('H', "print url hashes").Optional().NoArgument();
        }

        void Init(TOptsParseResult& args) override {
            PrintIds = args.Has('D');
            if (args.Has('H')) {
                DocHashMapping = LoadDocHashMapping(args.Get<TString>("input"));
            }
        }

        bool OnParsedDoc(ui32 realDocId, const NRTYServer::TParsedDoc&) override {
            if (PrintIds)
                Cout << realDocId << "\t";
            if (DocHashMapping) {
                TDocHash value(/*zeroed=*/true);
                auto iter = DocHashMapping->find(realDocId);
                if (iter != DocHashMapping->end()) {
                    value = iter->second;
                }
                TString buffer;
                const NUrlId::TUrlId* hash = reinterpret_cast<const NUrlId::TUrlId*>(&value.Ptr()[8]); //assume it's the first 8 bytes (this may change, of course)
                Cout << NUrlId::Hash2StrZ(*hash, buffer) << "\t";
            }
            return true;
        }
    private:
        bool PrintIds = false;
        THolder<TId2Hash> DocHashMapping;
    };

    class TDebugDumpKiwiObjectAction final : public IAction {
    public:
        void Init(TOptsParseResult&) override {
            JsonOptions.ValidateUtf8 = false;
            JsonOptions.FormatOutput = true;
            JsonOptions.SortKeys = true;
        }

        bool OnParsedDoc(ui32 realDocId, const NRTYServer::TParsedDoc& doc) override {
            if (!doc.HasDocument()) {
                return true; // skip
            }

            if (!doc.GetDocument().HasIndexedDoc() || !doc.GetDocument().GetIndexedDoc().HasKiwiObject()) {
                ERROR_LOG << "Missing Kiwi data for finalDocId=" << realDocId << Endl;
                return true;
            }

            const TString binKiwiObject = doc.GetDocument().GetIndexedDoc().GetKiwiObject();

            NJson::TJsonValue c = DeserializeKiwiObjectToJson(TBlob::FromString(binKiwiObject));
            NJson::TJsonWriter wr(&Cout, JsonOptions);
            wr.Write(&c);
            wr.Flush();

            return true;
        }
    private:
        NJson::TJsonWriterConfig JsonOptions;
    };

    static TAtomic s_StopSignal;
}


static void SetSigintHandler() {
    ::SetAsyncSignalFunction(SIGINT, [] (int) {
        if (!AtomicGet(s_StopSignal)) {
            WARNING_LOG << "Ctrl-C received, exiting" << Endl;
        } else {
            WARNING_LOG << "Force exit, terminate now" << Endl;
            exit(1);
        }
        AtomicSet(s_StopSignal, true);
    });
}

int run(int argc, const char **argv, IAction& action) {
    TOpts opts = TOpts::Default();
    opts.AddHelpOption();
    opts.AddVersionOption();
    opts.AddLongOption('i', "input", "source index path").RequiredArgument("STRING").Required();
    action.AddOpts(opts);

    TOptsParseResult args(&opts, argc, argv);
    action.Init(args);

    THolder<TFADumper> dumper = TFADumper::Create({
            args.Get<TString>("input"),
        });

    SetSigintHandler();
    dumper->Open();
    dumper->Run(action, &s_StopSignal);
    dumper->Close();

    return EXIT_SUCCESS;
}

int main_urls(int argc, const char** argv) {
    TPrintUrlAction action;
    return run(argc, argv, action);
}

int main_kiwi(int argc, const char** argv) {
    TDebugDumpKiwiObjectAction action;
    return run(argc, argv, action);
}

int main(int argc, const char** argv) {
    DoInitGlobalLog("cerr", TLOG_NOTICE, false, false);

    TModChooser modChooser;
    modChooser.AddMode("urls", main_urls, "dump document URLs to stdout");
    modChooser.AddMode("kiwi", main_kiwi, "dump KiwiObject to stdout in text form");


    try {
        return modChooser.Run(argc, argv);
    } catch (...) {
        Cerr << "An exception has occurred: " << CurrentExceptionMessage() << Endl;
        return  EXIT_FAILURE;
    }
}
