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

#include <maps/wikimap/mapspro/services/mrc/libs/toloka_manager/include/toloka_manager.h>

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

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

#include <library/cpp/iterator/zip.h>

namespace maps::mrc::eye {

namespace {

db::eye::DetectedHouseNumbers makeDetectedHouseNumbers(
    const common::ImageOrientation& orientation,
    const common::Size& originalSize,
    const db::toloka::Tasks& tasks)
{
    db::eye::DetectedHouseNumbers houseNumbers;
    for (const auto& task : tasks) {
        auto input = toloka::parseJson<toloka::HouseNumberRecognitionInput>(task.inputValues());
        auto output = toloka::parseJson<toloka::HouseNumberRecognitionOutput>(*task.outputValues());

        if (output.state() == toloka::HouseNumberRecognitionState::Ok) {
            houseNumbers.push_back({
                common::revertByImageOrientation(
                    input.bbox(),
                    originalSize,
                    orientation
                ),
                1.,
                output.houseNumber()
            });
        }
    }

    return houseNumbers;
}

struct DetectionRequest {
    using TaskType = toloka::HouseNumberDetectionTask;

    DetectionRequest(
        db::TId recognitionId,
        const toloka::DetectionInput& input,
        db::eye::RecognitionSource source)
        : recognitionId(recognitionId)
        , input(input)
        , source(source)
    {}

    db::TId recognitionId;
    toloka::DetectionInput input;
    db::eye::RecognitionSource source;

    db::eye::RecognitionTask createRecognitionTask(db::TId taskId) const {
        return {recognitionId, taskId};
    }
};

struct RecognitionRequest {
    using TaskType = toloka::HouseNumberRecognitionTask;

    RecognitionRequest(
        db::TId recognitionId,
        const toloka::HouseNumberRecognitionInput& input,
        db::TId parentRecognitionTaskId,
        db::eye::RecognitionSource source)
        : recognitionId(recognitionId)
        , input(input)
        , parentRecognitionTaskId(parentRecognitionTaskId)
        , source(source)
    {}

    db::TId recognitionId;
    toloka::HouseNumberRecognitionInput input;
    db::TId parentRecognitionTaskId;
    db::eye::RecognitionSource source;

    db::eye::RecognitionTask createRecognitionTask(db::TId taskId) const {
        return {recognitionId, taskId, parentRecognitionTaskId};
    }
};

template <typename TaskRequest>
void createNewTasks(
    pqxx::transaction_base& txn,
    const std::vector<TaskRequest>& requests,
    db::toloka::Platform platform)
{
    std::vector<typename TaskRequest::TaskType::InputType> inputs;
    for (const auto& request : requests) {
        inputs.push_back(request.input);
    }

    auto tasks = toloka::createTasks<typename TaskRequest::TaskType>(
        txn, platform, inputs);

    db::eye::RecognitionTasks newRecognitionTasks;
    for (const auto& [request, task] : Zip(requests, tasks)) {
        newRecognitionTasks.emplace_back(request.createRecognitionTask(task.id()));
    }

    db::eye::RecognitionTaskGateway(txn).insertx(newRecognitionTasks);
}

} // namespace

size_t handleHouseNumberRecognitions(
    pqxx::transaction_base& txn,
    const FrameUrlResolver& makeFrameUrl,
    db::eye::Recognitions recognitions)
{
    TolokaRecognitionContext context(txn, recognitions);

    db::eye::Recognitions updatedRecognitions;
    std::vector<DetectionRequest> detectionTolokaRequests;
    std::vector<DetectionRequest> detectionYangRequests;
    std::vector<RecognitionRequest> recognitionTolokaRequests;
    std::vector<RecognitionRequest> recognitionYangRequests;

    for (auto&& recognition : recognitions) {
        const db::eye::Frame& frame = context.getFrame(recognition);
        const db::FeaturePrivacy& privacy = context.getFramePrivacy(recognition);
        const std::string imageURL = makeFrameUrl.image(frame, recognition.orientation(), privacy);

        db::eye::RecognitionTasks recognitionTasks = context.getRecognitionTasks(recognition);

        if (recognitionTasks.empty()) {
            DetectionRequest request{
                recognition.id(),
                toloka::DetectionInput(imageURL),
                recognition.source()
            };

            if (recognition.source() == db::eye::RecognitionSource::Toloka) {
                detectionTolokaRequests.push_back(std::move(request));
            } else {
                detectionYangRequests.push_back(std::move(request));
            }
        } else if (recognitionTasks.size() == 1u) {
            db::toloka::Task tolokaTask = context.getTolokaTask(recognitionTasks[0]);

            if (isCompleted(tolokaTask)) {
                auto output = toloka::parseJson<toloka::DetectionOutput>(*tolokaTask.outputValues());

                if (output.bboxes().empty()) {
                    recognition.setValue(
                        makeDetectedHouseNumbers(recognition.orientation(), frame.originalSize(), {})
                    );
                    updatedRecognitions.push_back(recognition);
                } else {
                    for (const auto& outputBBox : output.bboxes()) {
                        auto size = common::transformByImageOrientation(
                            frame.originalSize(),
                            recognition.orientation()
                        );
                        common::ImageBox bbox(
                            common::round<size_t>(outputBBox.minX() * size.width),
                            common::round<size_t>(outputBBox.minY() * size.height),
                            common::round<size_t>(outputBBox.maxX() * size.width),
                            common::round<size_t>(outputBBox.maxY() * size.height)
                        );

                        RecognitionRequest request{
                            recognition.id(),
                            toloka::HouseNumberRecognitionInput(imageURL, bbox),
                            recognitionTasks[0].id(),
                            recognition.source()
                        };

                        if (recognition.source() == db::eye::RecognitionSource::Toloka) {
                            recognitionTolokaRequests.push_back(std::move(request));
                        } else {
                            recognitionYangRequests.push_back(std::move(request));
                        }
                    }
                }
            }
        } else {
            db::toloka::Tasks hnRecognitionTasks;
            for (const auto& recognitionTask : recognitionTasks) {
                db::toloka::Task task = context.getTolokaTask(recognitionTask);
                if (task.type() == toloka::HouseNumberRecognitionTask::DB_TASK_TYPE) {
                    hnRecognitionTasks.push_back(task);
                }
            }
            if (std::all_of(hnRecognitionTasks.begin(), hnRecognitionTasks.end(), isCompleted)) {
                db::eye::Recognition updatedRecognition = recognition;
                updatedRecognition.setValue(
                    makeDetectedHouseNumbers(recognition.orientation(), frame.originalSize(), hnRecognitionTasks)
                );
                updatedRecognitions.push_back(updatedRecognition);
            }
        }
    }

    INFO() << "House number recognitions to update: " << updatedRecognitions.size();
    db::eye::RecognitionGateway(txn).updatex(updatedRecognitions);

    INFO() << "New house number detection toloka tasks: " << detectionTolokaRequests.size();
    createNewTasks(txn, detectionTolokaRequests, db::toloka::Platform::Toloka);
    INFO() << "New house number detection yang tasks: " << detectionYangRequests.size();
    createNewTasks(txn, detectionYangRequests, db::toloka::Platform::Yang);

    INFO() << "New house number recognition toloka tasks: " << recognitionTolokaRequests.size();
    createNewTasks(txn, recognitionTolokaRequests, db::toloka::Platform::Toloka);
    INFO() << "New house number recognition yang tasks: " << recognitionYangRequests.size();
    createNewTasks(txn, recognitionYangRequests, db::toloka::Platform::Yang);

    return updatedRecognitions.size()
           + detectionTolokaRequests.size()
           + detectionYangRequests.size()
           + recognitionTolokaRequests.size()
           + recognitionYangRequests.size();
}

} // namespace maps::mrc::eye
