#include "feature_export.h"

#include <maps/libs/sql_chemistry/include/batch_load.h>
#include <maps/wikimap/mapspro/libs/common/include/yandex/maps/wiki/common/pgpool3_helpers.h>
#include <maps/libs/common/include/make_batches.h>
#include <maps/libs/log8/include/log8.h>

#include <maps/wikimap/mapspro/services/mrc/libs/common/include/opencv.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/common.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/feature.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/feature_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/object_in_photo_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/ugc/assignment_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/ugc/task_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/io.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/operation.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/serialization.h>

#include <optional>

namespace maps::mrc::yt_export {
namespace {

constexpr std::size_t MAX_PARALLEL_UPLOADING_JOBS = 20;
constexpr std::size_t BATCH_SIZE = 1000;
constexpr std::size_t YT_FEATURE_TABLE_INDEX = 0;
constexpr std::size_t YT_PRIVATE_OBJECT_TABLE_INDEX = 1;

db::Features loadPhotos(maps::pgpool3::Pool& pool, db::TId tasksGroupId)
{
    auto txn = pool.slaveTransaction();

    const auto taskIds = db::ugc::TaskGateway{*txn}.loadIds(
        db::ugc::table::Task::tasksGroupId == tasksGroupId);

    const auto assignmentIds = db::ugc::AssignmentGateway{*txn}.loadIds(
        db::ugc::table::Assignment::taskId.in(taskIds));

    sql_chemistry::BatchLoad<db::table::Feature> batch{
        BATCH_SIZE,
        db::table::Feature::isPublished
            && db::table::Feature::assignmentId.in(assignmentIds)};

    db::Features photos;
    while (batch.next(*txn)) {
        photos.insert(photos.end(), batch.begin(), batch.end());
    }

    return photos;
}

db::ObjectsInPhoto loadPrivacyObjects(maps::pgpool3::Pool& pool,
                                      const db::Features& photos)
{
    db::ObjectsInPhoto objects;
    auto txn = pool.slaveTransaction();

    for (auto photosBatch :
         maps::common::makeBatches(photos, BATCH_SIZE)) {

        db::TIds featureIds;
        featureIds.reserve(BATCH_SIZE);
        for (const auto& photo : photosBatch) {
            featureIds.push_back(photo.id());
        }

        db::ObjectsInPhoto objectsBatch = db::ObjectInPhotoGateway{*txn}.load(
            db::table::ObjectInPhotoTable::featureId.in(featureIds)
            && db::table::ObjectInPhotoTable::type.in(
                   {db::ObjectInPhotoType::Face,
                    db::ObjectInPhotoType::LicensePlate}));

        for (auto&& object : objectsBatch) {
            objects.push_back(std::move(object));
        }
    }

    return objects;
}

class YTTablesHolder {
public:
    YTTablesHolder(const NYT::IClientPtr& ytClient,
                   const TString& basePath,
                   const TString& tasksGroupIdStr)
        : photosTablePath(
              basePath + "/feature_" + tasksGroupIdStr)
        , privacyObjectsTablePath(
              basePath + "/feature_object_" + tasksGroupIdStr)
        , ytClient_(ytClient)
    {
        if (!ytClient->Exists(basePath)) {
            INFO() << "Create YT node " << basePath;
            ytClient->Create(basePath, NYT::NT_MAP,
                             NYT::TCreateOptions().Recursive(true));
        }
    }

    ~YTTablesHolder()
    {
        for (const auto& path : {photosTablePath, privacyObjectsTablePath}) {
            try {
                ytClient_->Remove(path);
            }
            catch (std::exception& e) {
                ERROR() << "Error while removing " << path << "\n"
                        << e.what();
            }
            catch (...) {
                WARN() << "Unknown error while removing " << path;
            }
        }
    }

    const TString photosTablePath;
    const TString privacyObjectsTablePath;

private:
    NYT::IClientPtr ytClient_;
};

YTTablesHolder uploadDataToTmpTables(const common::Config& mrcConfig,
                                     const NYT::IClientPtr& ytClient,
                                     db::TId tasksGroupId)
{
    const TString baseTmpPath{};

    YTTablesHolder ytTablesHolder{
        ytClient,
        TString{mrcConfig.externals().yt().path() + "/yt_features_export"},
        TString{std::to_string(tasksGroupId)}};

    auto poolHolder =
        mrcConfig.makePoolHolder(maps::mrc::common::LONG_READ_DB_ID,
                                 maps::mrc::common::LONG_READ_POOL_ID);
    const auto photos = loadPhotos(poolHolder.pool(), tasksGroupId);

    const auto privacyObjects
        = loadPrivacyObjects(poolHolder.pool(), photos);

    INFO() << "YT: Store " << photos.size() << " features to temporary table "
           << ytTablesHolder.photosTablePath;
    yt::saveToTable(*ytClient, ytTablesHolder.photosTablePath, photos);

    INFO() << "YT: Store " << privacyObjects.size()
           << " privacy objects to temporary table "
           << ytTablesHolder.privacyObjectsTablePath;
    yt::saveToTable(*ytClient, ytTablesHolder.privacyObjectsTablePath,
                    privacyObjects);

    INFO() << "YT: Upload feature images to " << ytTablesHolder.photosTablePath;
    yt::uploadFeatureImage(
        mrcConfig, *ytClient, ytTablesHolder.photosTablePath,
        ytTablesHolder.photosTablePath, MAX_PARALLEL_UPLOADING_JOBS);

    return ytTablesHolder;
}

using namespace NYT;
using FeatureWithImage = yt::FeatureWithImage<db::Feature>;

class FeaturesExporter
    : public IReducer<TTableReader<TNode>, TTableWriter<TNode>> {
public:
    void Do(TReader* reader, TWriter* writer) override
    {
        std::optional<FeatureWithImage> featureWithImage{std::nullopt};
        db::ObjectsInPhoto privacyObjects;

        for (; reader->IsValid(); reader->Next()) {
            auto row = reader->GetRow();

            if (reader->GetTableIndex() == YT_FEATURE_TABLE_INDEX) {

                featureWithImage = yt::deserialize<FeatureWithImage>(row);
            }
            else if (reader->GetTableIndex() == YT_PRIVATE_OBJECT_TABLE_INDEX) {
                privacyObjects.push_back(
                    yt::deserialize<db::ObjectInPhoto>(row));
            }
        }

        for (const auto& privacyObject : privacyObjects) {
            mrc::common::ellipticalBlur(featureWithImage->image,
                                        privacyObject.imageBox());
        }

        auto row = yt::serialize(*featureWithImage);
        row["pos"] = yt::serialize(
            featureWithImage->feature.geodeticPos());
        writer->AddRow(row);
    }
};
REGISTER_REDUCER(FeaturesExporter);

void exportFeatures(const NYT::IClientPtr& ytClient,
                    const TString& photosTablePath,
                    const TString& privacyObjectsTablePath,
                    const TString& ytExportTablePath)
{
    ytClient->Sort(TSortOperationSpec()
                       .AddInput(photosTablePath)
                       .Output(photosTablePath)
                       .SortBy({"feature_id"}));

    ytClient->Sort(TSortOperationSpec()
                       .AddInput(privacyObjectsTablePath)
                       .Output(privacyObjectsTablePath)
                       .SortBy({"feature_id"}));

    ytClient->Reduce(TReduceOperationSpec()
                         .ReduceBy({"feature_id"})
                         .AddInput<TNode>(photosTablePath)
                         .AddInput<TNode>(privacyObjectsTablePath)
                         .AddOutput<TNode>(ytExportTablePath),
                     new FeaturesExporter);
}

} // anonymous namespace

void exportFeatures(const common::Config& mrcConfig,
                    db::TId tasksGroupId,
                    const TString& ytExportTablePath)
{
    const auto ytClient = mrcConfig.externals().yt().makeClient();

    REQUIRE(!ytClient->Exists(ytExportTablePath),
            "Node exists: " << ytExportTablePath);

    const auto ytTablesHolder
        = uploadDataToTmpTables(mrcConfig, ytClient, tasksGroupId);

    exportFeatures(ytClient, ytTablesHolder.photosTablePath,
                   ytTablesHolder.privacyObjectsTablePath, ytExportTablePath);

    INFO() << "YT: Done with features export";
}

} // namespace maps::mrc::yt_export
