#include "processor_mode.h"

#include <saas/tools/remap_for_yt_pull/lib/common_context.h>

#include <saas/api/mr_client/processors/processors.h>
#include <saas/api/yt_pull_client/saas_yt_writer.h>

#include <search/session/compression/compression.h>

#include <mapreduce/yt/interface/client.h>

#include <library/cpp/getopt/last_getopt.h>

#include <google/protobuf/messagext.h>
#include <google/protobuf/text_format.h>

#include <util/ysaveload.h>
#include <util/string/vector.h>
#include <util/generic/set.h>

namespace {
    class TProcessorConversionMap: public NYT::IMapper<NYT::TTableReader<NYT::TNode>, NYT::TTableWriter<NYT::TYaMRRow>> {
    private:
        NSaas::ShardingType ShardingType = NSaas::UrlHash;
        bool WritePrototext = false;
        TString Processor;
        NSaas::TRowProcessorOptions ProcessorOptions;

        THolder<NSaas::IRowProcessor> RowProcessor;
        THolder<NSaas::NYTPull::TSaasSlaveTableWriter> SlaveWriter;

    public:
        TProcessorConversionMap(NSaas::ShardingType shardingType, TString processor, const NSaas::TRowProcessorOptions& options, bool writePrototext)
            : ShardingType(shardingType)
            , WritePrototext(writePrototext)
            , Processor(processor)
            , ProcessorOptions(options)
        {
        }

        TProcessorConversionMap() {
        }

        void Save(IOutputStream& stream) const final {
            ::Save(&stream, ToString(ShardingType));
            ::Save(&stream, WritePrototext);
            ::Save(&stream, Processor);
            ProcessorOptions.Save(&stream);
        }

        void Load(IInputStream& stream) final {
            TString shardingTypeStr;
            ::Load(&stream, shardingTypeStr);
            ShardingType = FromString(shardingTypeStr);
            ::Load(&stream, WritePrototext);
            ::Load(&stream, Processor);
            ProcessorOptions.Load(&stream);
        }

        void Start(NYT::TTableWriter<NYT::TYaMRRow>* writer) final {
            NSaas::TShardsDispatcher::TPtr shardDispatcherPtr = MakeAtomicShared<NSaas::TShardsDispatcher>(NSaas::TShardsDispatcher::TContext(ShardingType));
            RowProcessor.Reset(NSaas::IRowProcessor::TFactory::Construct(Processor, ProcessorOptions));
            SlaveWriter.Reset(new NSaas::NYTPull::TSaasSlaveTableWriter(writer, shardDispatcherPtr, false, WritePrototext));
        }

        void Do(NYT::TTableReader<NYT::TNode>* reader, NYT::TTableWriter<NYT::TYaMRRow>* /*writer*/) final {
            while(reader->IsValid()) {
                const auto& row = reader->GetRow();
                NRTYServer::TMessage message = RowProcessor->ProcessRowSingle(row).ToProtobuf();
                if (Y_UNLIKELY(!SlaveWriter->WriteMessage(message))) {
                    Cerr << "Failed to serialize message for " << message.GetDocument().GetUrl() << '\n';
                }
                reader->Next();
            }
        }

        void Finish(NYT::TTableWriter<NYT::TYaMRRow>* /*writer*/) final {
            SlaveWriter.Reset();
        }
    };

    REGISTER_MAPPER(TProcessorConversionMap);

} // namespace

int NSaas::RunProcessorConversion(int argc, const char **argv) {
    TCommonContext commonContext;
    TString processorKey;
    TRowProcessorOptions options;

    NLastGetopt::TOpts opts;
    opts.AddHelpOption('h');
    opts.SetFreeArgsMax(0);

    AddCommonOpts(opts, commonContext, true);


    TSet<TString> availableProcessors;
    IRowProcessor::TFactory::GetRegisteredKeys(availableProcessors);
    const TString availableProcessorsStr = JoinStrings(availableProcessors.cbegin(), availableProcessors.cend(), "|");
    opts.AddLongOption("processor", "Processor name: {" + availableProcessorsStr + "}")
        .DefaultValue("simple")
        .StoreResult(&processorKey);

    opts.AddLongOption("keyprefix", "key prefix")
        .DefaultValue("0")
        .StoreResult(&options.Prefix);

    opts.AddLongOption("document-timestamp", "document timestamp")
        .StoreResultT<ui32>(&options.Timestamp);

    opts.AddLongOption("document-deadline-minutes-utc", "document deadline")
        .StoreResultT<ui32>(&options.DeadlineMinutesUTC);

    opts.AddLongOption("compression-method", "property compression method")
        .StoreResult(&options.Compression);

    std::function<NRTYServer::TMessage::TMessageType(const TString&)> func = [](const TString& val) -> NRTYServer::TMessage::TMessageType {
        NRTYServer::TMessage::TMessageType res;
        Y_ENSURE(NRTYServer::TMessage::TMessageType_Parse(val, &res), "Failed to parse");
        return res;
    };
    const NProtoBuf::EnumDescriptor* descriptor = NRTYServer::TMessage::TMessageType_descriptor();
    TVector<TString> allowedMessageTypes;
    allowedMessageTypes.reserve(descriptor->value_count());
    for (i32 i = 0; i < descriptor->value_count(); ++i) {
        const auto* item = descriptor->FindValueByNumber(i);
        if (item != nullptr) {
            allowedMessageTypes.push_back(item->name());
        }
    }

    NRTYServer::TMessage::TMessageType messageType;
    const TString allowedMessageTypesStr = JoinStrings(allowedMessageTypes.cbegin(), allowedMessageTypes.cend(), "|");
    opts.AddLongOption("message-type", "Message type: {" + allowedMessageTypesStr + "}")
        .DefaultValue(NRTYServer::TMessage::TMessageType_Name(NRTYServer::TMessage::MODIFY_DOCUMENT))
        .StoreMappedResultT<TString>(&messageType, func);

    NLastGetopt::TOptsParseResult parseOpt(&opts, argc, argv);

    Y_ENSURE(commonContext.ShardingType != NSaas::Broadcast, "Brodacast type is unsupported");
    Y_ENSURE(availableProcessors.contains(processorKey), "Unknown processor name: " + processorKey);
    Y_ENSURE(NRTYServer::TMessage_TMessageType_IsValid(messageType), "Invalid message type");
    options.ActionType = enum_cast<NSaas::TAction::TActionType>(messageType);

    NYT::IClientPtr client = NYT::CreateClient(commonContext.MRServer);
    NYT::ITransactionPtr txClient = client->StartTransaction();
    try {
        NYT::TMapOperationSpec spec;
        spec.AddInput<NYT::TNode>(commonContext.SrcTable)
            .AddOutput<NYT::TYaMRRow>(commonContext.DstTable);
        TIntrusivePtr<TProcessorConversionMap> mapper = new TProcessorConversionMap(
                commonContext.ShardingType, processorKey, options, commonContext.ProtoText);
        txClient->Map(spec, mapper);
        mapper.Reset();

        NSaas::NYTPull::SortSlaveTable(commonContext.DstTable, txClient);
        if (!txClient->Exists(commonContext.MasterTable)) {
            NSaas::NYTPull::CreateMasterTable(commonContext.MasterTable, txClient);
        }
        NSaas::NYTPull::RegisterSlaveTable(commonContext.MasterTable, commonContext.DstTable, commonContext.SlaveTableTimestamp, txClient);
        txClient->Commit();
        txClient.Drop();
        client.Drop();
    } catch(...) {
        Cerr << "Error " << CurrentExceptionMessage() << Endl;
        txClient->Abort();
        txClient.Drop();
        client.Drop();
        return 1;
    }
    return 0;
}
