#pragma once

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/secure_config.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_object/include/config.h>

#include <maps/wikimap/mapspro/services/mrc/libs/common/include/exif.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/opencv.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/load.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/recognition_gateway.h>

#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/io.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/schema.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/serialization.h>

#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.h>

#include <optional>
#include <string>
#include <string_view>

namespace maps::mrc::eye {

template<class Detector, class Segmentator, class MakeRecognition>
class DetectObjectWithFilteringWorker: public yt::Mapper {

public:
    DetectObjectWithFilteringWorker() = default;

    DetectObjectWithFilteringWorker(const FrameLoader& frameLoader)
        : frameLoader_(frameLoader)
    {}

    void Do(yt::Reader* reader, yt::Writer* writer) override
    {
        Detector detector;
        Segmentator segmentator;
        MakeRecognition makeRecognition;

        for (; reader->IsValid(); reader->Next()) {
            const auto frame = yt::deserialize<db::eye::Frame>(reader->GetRow());
            const cv::Mat image = frameLoader_.load(frame);

            const auto size = common::sizeOf(image);
            ASSERT(size.width and size.height);

            const auto detections = detector.detect(image);

            // Achtung! Mask may be empty!
            cv::Mat mask;
            if (not detections.empty()) {
                mask = segmentator.segment(image);
                ASSERT(mask.dims == 2);
            }

            writer->AddRow(yt::serialize(makeRecognition(frame, detections, mask)));
        }
    }

    Y_SAVELOAD_JOB(frameLoader_);

private:
    FrameLoader frameLoader_;
};

NYT::TNode makeSimpleDetectObjectSpec(
        std::string_view title,
        bool useGpu,
        size_t concurrency,
        const std::optional<std::string>& portoLayerPath,
        std::optional<yt::PoolType> poolType);

template<class DetectObjectWorker>
db::eye::Recognitions simpleDetectObject(
        const SimpleDetectObjectConfig& config,
        std::string_view title,
        const db::eye::Frames& frames)
{
    {   // check config
        REQUIRE(isValid(config.yt), "Invalid YT config");
    }

    if (frames.empty()) {
        INFO() << "No frames for detection!";
        return {};
    }

    INFO() << "Run detection...";
    NYT::ITransactionPtr ytTxn = config.yt.client->StartTransaction();

    const TString root(config.yt.rootPath);
    if (!ytTxn->Exists(root)) {
        INFO() << "Create node " << root;
        ytTxn->Create(root, NYT::NT_MAP, NYT::TCreateOptions().Recursive(true));
    }

    const TString input = root + "/frame";
    const TString output = root + "/recognition";

    INFO() << "Save to table " << input;
    yt::saveToTable(
        *ytTxn,
        NYT::TRichYPath(input)
            .Schema(yt::getSchemaOf<db::eye::Frame>()),
        frames
    );

    REQUIRE(SecureConfig::isInitialized(), "Secure config was not initialized");

    ytTxn->Map(
        NYT::TMapOperationSpec()
            .AddInput<NYT::TNode>(input)
            .AddOutput<NYT::TNode>(
                NYT::TRichYPath(output)
                    .Schema(yt::getSchemaOf<db::eye::Recognition>())
            )
            .JobCount(std::max(frames.size() / config.yt.partitionSize, 1ul)),
        new DetectObjectWorker(*(config.yt.frameLoader)),
        NYT::TOperationOptions().Spec(
            makeSimpleDetectObjectSpec(
                title,
                config.yt.useGpu,
                config.yt.concurrency,
                config.yt.portoLayerPath,
                config.yt.poolType
            )
        )
        .SecureVault(SecureConfig::instance())
    );

    INFO() << "Load table " << output;
    auto result = yt::loadFromTable<db::eye::Recognitions>(*ytTxn, output);

    if (config.yt.commit) {
        INFO() << "Yt commited!";
        ytTxn->Commit();
    }

    return result;
}

} // namespace maps::mrc::eye
