#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/impl/context.h>

#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/recognition_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/toloka/task_gateway.h>

namespace maps::mrc::eye {

namespace {

std::map<db::TId, db::eye::Frame> loadFrameIdToFrame(
    pqxx::transaction_base& txn,
    const db::eye::Recognitions& recognitions)
{
    db::TIds frameIds;
    for (const db::eye::Recognition& recognition : recognitions) {
        frameIds.push_back(recognition.frameId());
    }

    db::eye::Frames frames = db::eye::FrameGateway(txn).loadByIds(frameIds);

    std::map<db::TId, db::eye::Frame> frameIdToFrame;
    for (const db::eye::Frame& frame : frames) {
        frameIdToFrame.emplace(frame.id(), frame);
    }

    return frameIdToFrame;
}

std::map<db::TId, db::FeaturePrivacy> loadFrameIdToPrivacy(
    pqxx::transaction_base& txn,
    const db::eye::Recognitions& recognitions)
{
    db::TIds frameIds;
    for (const db::eye::Recognition& recognition : recognitions) {
        frameIds.push_back(recognition.frameId());
    }

    db::eye::FramePrivacies privacies = db::eye::FramePrivacyGateway(txn).load(
        db::eye::table::FramePrivacy::frameId.in(frameIds)
    );

    std::map<db::TId, db::FeaturePrivacy> frameIdToPrivacy;
    for (const db::eye::FramePrivacy& privacy : privacies) {
        frameIdToPrivacy.emplace(privacy.frameId(), privacy.type());
    }

    return frameIdToPrivacy;
}

std::map<db::TId, db::eye::RecognitionTasks> loadRecognitionIdToRecognitionTasks(
    pqxx::transaction_base& txn,
    const db::eye::Recognitions& recognitions)
{
    db::TIds recognitionIds;
    for (const db::eye::Recognition& recognition : recognitions) {
        recognitionIds.push_back(recognition.id());
    }

    db::eye::RecognitionTasks recognitionTasks = db::eye::RecognitionTaskGateway(txn).load(
        db::eye::table::RecognitionTask::recognitionId.in(recognitionIds)
    );

    std::map<db::TId, db::eye::RecognitionTasks> recognitionIdToRecognitionTasks;
    for (const db::eye::RecognitionTask& recognitionTask : recognitionTasks) {
        recognitionIdToRecognitionTasks[recognitionTask.recognitionId()].push_back(recognitionTask);
    }

    return recognitionIdToRecognitionTasks;
}

std::map<db::TId, db::toloka::Task> loadRecognitionTaskIdToTolokaTask(
    pqxx::transaction_base& txn,
    const std::map<db::TId, db::eye::RecognitionTasks>& recognitionIdToRecognitionTasks)
{
    db::TIds taskIds;
    std::map<db::TId, db::TId> taskIdToRecognitionTaskId;
    for (const auto& [recognitionId, recognitionTasks] : recognitionIdToRecognitionTasks) {
        for (const db::eye::RecognitionTask& recognitionTask : recognitionTasks) {
            taskIds.push_back(recognitionTask.taskId());
            taskIdToRecognitionTaskId.emplace(recognitionTask.taskId(), recognitionTask.id());
        }
    }

    db::toloka::Tasks tasks = db::toloka::TaskGateway(txn).loadByIds(taskIds);

    std::map<db::TId, db::toloka::Task> recognitionTaskIdToTask;
    for (db::toloka::Task& task : tasks) {
        recognitionTaskIdToTask.emplace(taskIdToRecognitionTaskId[task.id()], std::move(task));
    }

    return recognitionTaskIdToTask;
}

} // namespace

TolokaRecognitionContext::TolokaRecognitionContext(
    pqxx::transaction_base& txn,
    const db::eye::Recognitions& recognitions)
{
    frameIdToFrame_ = loadFrameIdToFrame(txn, recognitions);
    frameIdToPrivacy_ = loadFrameIdToPrivacy(txn, recognitions);
    recognitionIdToRecognitionTasks_
        = loadRecognitionIdToRecognitionTasks(txn, recognitions);
    recognitionTaskIdToTolokaTask_
        = loadRecognitionTaskIdToTolokaTask(txn, recognitionIdToRecognitionTasks_);
}

const db::eye::Frame& TolokaRecognitionContext::getFrame(
    const db::eye::Recognition& recognition) const
{
    return frameIdToFrame_.at(recognition.frameId());
}

db::FeaturePrivacy TolokaRecognitionContext::getFramePrivacy(
    const db::eye::Recognition& recognition) const
{
    return frameIdToPrivacy_.at(recognition.frameId());
}

db::toloka::Tasks TolokaRecognitionContext::getTolokaTasks(
    const db::eye::Recognition& recognition) const
{
    db::toloka::Tasks tolokaTasks;
    for (const auto& recognitionTask : getRecognitionTasks(recognition)) {
        tolokaTasks.push_back(getTolokaTask(recognitionTask));
    }
    return tolokaTasks;
}

db::eye::RecognitionTasks TolokaRecognitionContext::getRecognitionTasks(
    const db::eye::Recognition& recognition) const
{
    auto recognitionTasksIt = recognitionIdToRecognitionTasks_.find(recognition.id());
    if (recognitionTasksIt != recognitionIdToRecognitionTasks_.end()) {
        return recognitionTasksIt->second;
    } else {
        return {};
    }
}

const db::toloka::Task& TolokaRecognitionContext::getTolokaTask(
    const db::eye::RecognitionTask& recognitionTask) const
{
    return recognitionTaskIdToTolokaTask_.at(recognitionTask.id());
}

} // namespace maps::mrc::eye
