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

#include <maps/libs/geolib/include/bounding_box.h>

#include <maps/libs/json/include/builder.h>

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

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

#include <contrib/libs/opencv/include/opencv2/opencv.hpp>

#include <random>
#include <vector>
#include <fstream>

using namespace maps;
using namespace maps::http;
using namespace maps::geolib3;

namespace {

std::string getFeatureURL(int64_t featureId) {
    return "https://core-nmaps-mrc-browser.maps.yandex.ru/feature/" + std::to_string(featureId) + "/image";
}

void bboxToJson(const BoundingBox& bbox, json::ObjectBuilder& b) {
    b["type"] = "rectangle";
    b["data"] = [&](json::ObjectBuilder b) {
        b["p1"] = [&](json::ObjectBuilder b) {
            b["x"] = bbox.minX();
            b["y"] = bbox.minY();
        };
        b["p2"] = [&](json::ObjectBuilder b) {
            b["x"] = bbox.maxX();
            b["y"] = bbox.maxY();
        };
    };
}

std::vector<BoundingBox> readBBoxes(const NYT::TNode& node) {
    std::vector<BoundingBox> bboxes;
    for (int i = 0; i < node.AsList().ysize(); i++) {
        const NYT::TNode& coords = node[i]["bbox"];
        bboxes.push_back({
            Point2(coords[0][0].AsInt64(), coords[0][1].AsInt64()),
            Point2(coords[1][0].AsInt64(), coords[1][1].AsInt64())
        });
    }
    return bboxes;
}

cv::Mat loadImage(http::Client& client, const std::string& url) {
    return common::retry(
        [&]() {
            http::Request request(client, http::GET, http::URL(url));
            http::Response response = request.perform();
            if (response.status() == 200) {
                std::string data = response.readBody();
                return cv::imdecode(
                    cv::Mat(1, data.size(), CV_8UC1, const_cast<char*>(data.data())),
                    cv::IMREAD_COLOR
                );
            } else {
                throw maps::RuntimeError()
                    << "Failed to download image: " << url;
            }
        },
        maps::common::RetryPolicy()
            .setTryNumber(6)
            .setInitialCooldown(std::chrono::seconds(1))
            .setCooldownBackoff(2.)
    );
}

std::vector<BoundingBox> normalizeBBoxes(
    const std::vector<BoundingBox>& bboxes, const std::string& featureURL)
{
    std::vector<BoundingBox> normalizedBBoxes;

    http::Client httpClient;
    for (const BoundingBox& bbox : bboxes) {
        cv::Mat image = loadImage(httpClient, featureURL);
        normalizedBBoxes.push_back({
            {bbox.minX() / image.cols, bbox.minY() / image.rows},
            {bbox.maxX() / image.cols, bbox.maxY() / image.rows}
        });
    }

    return normalizedBBoxes;
}

class Task {
public:
    Task(const NYT::TNode& node)
        : featureId_(node["feature_id"].AsInt64())
        , featureURL_(getFeatureURL(featureId_))
        , bboxes_(normalizeBBoxes(readBBoxes(node["objects"]), featureURL_))
    {
    }

    void toJson(json::ObjectBuilder& builder, bool dumpAnnotation) const {
        builder["inputValues"] = [&](json::ObjectBuilder b) {
            b["feature_id"] = featureId_;
            b["image"] = featureURL_;
            if (dumpAnnotation && !bboxes_.empty()) {
                b["annotations"] << [&](json::ArrayBuilder b) {
                    for (const BoundingBox& bbox : bboxes_) {
                        b << [&](json::ObjectBuilder b) {
                            bboxToJson(bbox, b);
                        };
                    }
                };
            }
        };
    }

private:
    int64_t featureId_;
    std::string featureURL_;
    std::vector<BoundingBox> bboxes_;
};

} // namespace

int main(int argc, const char** argv)
try {
    NYT::Initialize(argc, argv);

    cmdline::Parser parser("Prepare toloka tasks in JSON");

    cmdline::Option<std::string> ytPath = parser.string("yt_path")
        .required()
        .help("Path to YT table with features");

    cmdline::Option<std::string> tasksPath = parser.string("tasks_path")
        .required()
        .help("Path to output json with tasks");

    cmdline::Option<size_t> tasksCount = parser.size_t("count")
        .required()
        .help("Tasks count in output json");

    cmdline::Option<bool> dumpAnnotation = parser.flag("dump_annotation")
        .help("Write annotation from YT table to tasks");

    cmdline::Option<bool> doShuffle = parser.flag("random_shuffle")
        .help("Shuffle rows in YT tables");

    parser.parse(argc, const_cast<char**>(argv));

    INFO() << "Create YT client: yt::hahn";
    NYT::IClientPtr client = NYT::CreateClient("hahn");

    std::vector<Task> tasks;

    INFO() << "Reading tasks from YT table: " << ytPath;
    NYT::TTableReaderPtr<NYT::TNode> reader
        = client->CreateTableReader<NYT::TNode>(NYT::TYPath(ytPath));
    if (!doShuffle) {
        INFO() << "Reading " << tasksCount << " rows from YT table";
        size_t addedCount = 0;
        for (; addedCount < tasksCount && reader->IsValid(); reader->Next()) {
            try {
                Task task(reader->GetRow());
                tasks.push_back(task);
                addedCount++;
            } catch (const maps::RuntimeError& e) {
                INFO() << "Failed to create task";
            }
        }
    } else {
        INFO() << "Reading all rows from YT table";
        std::vector<NYT::TNode> rows;
        for (; reader->IsValid(); reader->Next()) {
            rows.push_back(reader->GetRow());
        }
        INFO() << "Shuffling rows";
        std::mt19937 rndGen{std::random_device{}()};
        INFO() << "Reading " << tasksCount << " rows";
        std::shuffle(rows.begin(), rows.end(), rndGen);
        size_t addedCount = 0;
        for (auto it = rows.begin(); addedCount < tasksCount && it != rows.end() ; it++) {
            try {
                Task task(*it);
                tasks.push_back(task);
                addedCount++;
            } catch (const maps::Exception& e) {
                INFO() << "Failed to create task: " << e.what();
            }
        }
    }

    INFO() << "Dumping tasks to json: " << tasksPath;
    std::ofstream ofs(tasksPath);
    json::Builder builder(ofs);
    builder << [&](json::ArrayBuilder b) {
        for (const Task& task : tasks) {
            b << [&](json::ObjectBuilder b) {
                task.toJson(b, dumpAnnotation);
            };
        }
    };
    ofs.close();

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    INFO() << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    INFO() << e.what();
    return EXIT_FAILURE;
}
catch (...) {
    INFO() << "Caught unknown exception";
    return EXIT_FAILURE;
}
