#include <saas/tools/querysearch/qd_validator/params.pb.h>

#include <kernel/querydata/saas/qd_saas_key.h>
#include <kernel/querydata/saas_yt/qd_saas_yt.h>
#include <kernel/querydata/saas_yt/idl/qd_saas_input_record.pb.h>
#include <kernel/querydata/saas_yt/idl/qd_saas_snapshot_record.pb.h>

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

#include <library/cpp/getoptpb/getoptpb.h>

#include <util/string/split.h>

using TQdInputRecord = NQueryDataSaaS::TQDSaaSInputRecord;
using NQueryDataSaaS::TQDSaaSSnapshotRecord;

struct TQdInputValidator : NYT::IMapper<NYT::TTableReader<TQdInputRecord>, NYT::TTableWriter<NYT::TNode>> {
    TQdInputValidator() = default;
    TQdInputValidator(const TVector<TString>& tablePaths)
        : TablePaths(std::move(tablePaths))
    {
    }

    Y_SAVELOAD_JOB(TablePaths)

    void Do(NYT::TTableReader<TQdInputRecord>* input, NYT::TTableWriter<NYT::TNode>* output) final {
        NQueryDataSaaS::TSnapshotRecordFillParams params;
        params.Namespace = "test_namespace";
        params.TimestampMicro = TInstant::Now().MicroSeconds();

        for (; input->IsValid(); input->Next()) {
            const auto& inputRow = input->GetRow();

            try {
                TQDSaaSSnapshotRecord snapRecord;
                NQueryDataSaaS::FillSnapshotRecordFromInputRecord(snapRecord, inputRow, params);
            } catch (...) {
                NYT::TNode node;
                node["table_path"] = TablePaths.at(input->GetTableIndex());
                node["row_index"] = input->GetRowIndex();
                node["doc"] = inputRow.Utf8DebugString();
                node["error"] = CurrentExceptionMessage();
                output->AddRow(node);
            }
        }
    }

private:
    TVector<TString> TablePaths;
};

REGISTER_MAPPER(TQdInputValidator);

void WaitForMultipleLocks(const TVector<NYT::ILockPtr>& locks, TDuration timeout) {
    TVector<NThreading::TFuture<void>> futures(Reserve(locks.size()));
    for (auto& lock : locks) {
        futures.emplace_back(lock->GetAcquiredFuture());
    }
    WaitExceptionOrAll(futures).GetValue(timeout);
}

NYT::TRichYPath GetLockedPath(NYT::ILock& lock) {
    NYT::TRichYPath path;
    path.Path("#" + GetGuidAsString(lock.GetLockedNodeId()));
    return path;
}

void ValidateTable(NYT::ITransaction& tx, const TVector<TString>& input, const TString& output, bool overwriteOutput) {
    TVector<NYT::ILockPtr> locks(Reserve(input.size()));
    for (auto& inp : input) {
        locks.emplace_back(tx.Lock(inp, NYT::LM_SNAPSHOT, NYT::TLockOptions().Waitable(true)));
    }
    WaitForMultipleLocks(locks, TDuration::Minutes(1));

    tx.Create(output, NYT::NT_TABLE, NYT::TCreateOptions().Recursive(true).Force(overwriteOutput));

    NYT::TMapOperationSpec mapSpec;
    for (auto& lock : locks) {
        mapSpec.AddInput<TQdInputRecord>(GetLockedPath(*lock));
    }
    mapSpec.AddOutput<NYT::TNode>(output);
    tx.Map(mapSpec, new TQdInputValidator(input));
}

int main(int argc, const char* argv[]) {
    NYT::Initialize(argc, argv);
    NUtils::TParams params = NGetoptPb::GetoptPbOrAbort(argc, argv);

    try {
        Y_ENSURE(params.InputSize() > 0, "At least one input table should be specified");
        auto client = NYT::CreateClient(params.GetYt());
        auto tx = client->StartTransaction();

        const auto& inp = params.GetInput();
        ValidateTable(*tx, {inp.begin(), inp.end()}, params.GetOutput(), params.GetForce());

        tx->Commit();

    } catch (...) {
        Cerr << "ERROR: " << CurrentExceptionMessage() << Endl;
        return 1;
    }
    return 0;
}

