#include "company_data_bundle_queue.h"

#include <yandex/maps/wiki/common/pg_utils.h>
#include <maps/libs/enum_io/include/enum_io.h>
#include <maps/libs/geolib/include/serialization.h>
#include <util/string/hex.h>
#include <pqxx/pqxx>

namespace maps::wiki::walkers_export_downloader {

namespace {

constexpr enum_io::Representations<CompanyDataBundle::State> COMPANY_DATA_BUNDLE_STATE_ENUM_REPRESENTATION {
    {CompanyDataBundle::State::Draft, "draft"},
    {CompanyDataBundle::State::Waiting, "waiting"},
    {CompanyDataBundle::State::Submitted, "submitted"},
    {CompanyDataBundle::State::Canceled, "canceled"},
    {CompanyDataBundle::State::Failed, "failed"},
};

std::string
quote(pqxx::transaction_base& work, std::string_view str)
{
    return work.quote(std::string(std::move(str)));
}

std::string
toWKB(const geolib3::Point2& poiCoordinate)
{
    return geolib3::WKB::toString(poiCoordinate);
}

geolib3::Point2
fromWKB(const std::string& wkb)
{
    return geolib3::EWKB::read<geolib3::SpatialReference::Epsg3395, geolib3::Point2>(HexDecode(wkb));
}

PhotoData readPhotoData(const pqxx::row& row)
{
    return PhotoData {
        .photoId = row["photo_id"].as<PhotoId>(),
        .subject = enum_io::fromString<PhotoSubject>(row["subject"].c_str()),
        .shootingPoint = fromWKB(row["shooting_point"].c_str()),
        .shootingTarget = fromWKB(row["shooting_target"].c_str())
    };
}

std::map<PhotoId, PhotoData>
readPhotosData(const pqxx::result& rows)
{
    std::map<PhotoId, PhotoData> photoDatas;
    for (const auto& row : rows) {
        auto data = readPhotoData(row);
        photoDatas.emplace(data.photoId, std::move(data));
    }
    return photoDatas;
}

std::map<PhotoId, PhotoData>
readPhotosData(pqxx::transaction_base& work, const std::set<PhotoId>& ids)
{
    const auto rows = work.exec(
        "SELECT * FROM sprav.photo_data "
        "WHERE photo_id IN ('" + common::join(ids, "','") + "')");
    return readPhotosData(rows);
}

CompanyDataBundle
readBundle(const pqxx::row& row)
{
    CompanyDataBundle bundle {
        .id = row["bundle_id"].as<BundleId>(),
        .permalinkId = row["permalink_id"].as<PermalinkId>(),
        .createdAt = maps::chrono::parseSqlDateTime(row["created_at"].c_str()),
        .state = enum_io::fromString<CompanyDataBundle::State>(row["bundle_state"].c_str()),
        .coordinate = fromWKB(row["altay_poi_position"].as<std::string>())
    };
    const auto& rowSubmittedAt = row["submitted_at"];
    if (!rowSubmittedAt.is_null()) {
        bundle.submittedAt = maps::chrono::parseSqlDateTime(rowSubmittedAt.c_str());
    }
    const auto& rowFeedbackTaskId = row["feedback_task_id"];
    if (!rowFeedbackTaskId.is_null()) {
        bundle.feedbackTaskId = rowFeedbackTaskId.as<FeedbackTaskId>();
    }
    return bundle;
}

std::vector<CompanyDataBundle>
readBundles(pqxx::result&& rows, pqxx::transaction_base& work)
{
    std::vector<CompanyDataBundle> bundles;
    std::set<BundleId> bundleIds;
    for (const auto& row : rows) {
        auto bundle = readBundle(row);
        bundleIds.insert(bundle.id);
        bundles.emplace_back(std::move(bundle));
    }
    const auto photoRows = work.exec(
        "SELECT bp.bundle_id, p.* FROM sprav.company_bundle_photo bp JOIN sprav.photo_data p USING(photo_id) "
        "WHERE " + common::whereClause("bundle_id", bundleIds));
    std::map<BundleId, std::set<PhotoId>> bundleToPhotos;
    for (const auto& row : photoRows) {
        bundleToPhotos[row["bundle_id"].as<BundleId>()].insert(row["photo_id"].as<PhotoId>());
    }
    const auto photos = readPhotosData(photoRows);
    for (auto& bundle : bundles) {
        const auto it = bundleToPhotos.find(bundle.id);
        if (it == bundleToPhotos.end()) {
            continue;
        }
        const auto& photoIds = it->second;
        for (const auto& photoId : photoIds) {
            const auto photoIt = photos.find(photoId);
            if (photoIt != photos.end()) {
                bundle.photos.emplace_back(photoIt->second);
            }
        }
    }
    return bundles;
}

} // namespace

CompanyDataBundleQueue::CompanyDataBundleQueue(pgpool3::Pool& socialPool)
    : socialPool_(socialPool)
{
}

BundleId
CompanyDataBundleQueue::createBundle(
    PermalinkId permalinkId,
    const geolib3::Point2& poiCoordinate)
{
    auto work = socialPool_.masterWriteableTransaction();
    const auto rows = work->exec(
        "INSERT INTO sprav.company_bundle(permalink_id, altay_poi_position) "
        "VALUES(" + std::to_string(permalinkId) +
        ", ST_GeomFromWKB(" + work->quote_raw(toWKB(poiCoordinate)) + ", 3395)) "
        "RETURNING bundle_id;");
    work->commit();
    return rows[0][0].as<BundleId>();
}

std::vector<BundleId>
CompanyDataBundleQueue::findBundles(
    PermalinkId permalinkId,
    CompanyDataBundle::State state) const
{
    auto work = socialPool_.masterReadOnlyTransaction();
    const auto rows = work->exec(
        "SELECT bundle_id FROM sprav.company_bundle "
        "WHERE permalink_id = " + std::to_string(permalinkId) +
        " AND bundle_state = " + quote(*work, toString(state)));
    std::vector<BundleId> bundleIds;
    bundleIds.reserve(rows.size());
    for (const auto& row : rows) {
        bundleIds.push_back(row[0].as<BundleId>());
    }
    return bundleIds;
}

void
CompanyDataBundleQueue::assignPhotoToBundle(BundleId id, const PhotoId& photoId)
{
    auto work = socialPool_.masterWriteableTransaction();
    work->exec(
        "INSERT INTO sprav.company_bundle_photo (bundle_id, photo_id) "
        "VALUES (" + std::to_string(id) + ", '" + photoId + "')");
    work->commit();
}

std::optional<PhotoData>
CompanyDataBundleQueue::loadPhotoData(const PhotoId& photoId) const
{
    auto work = socialPool_.masterReadOnlyTransaction();
    const auto datas = readPhotosData(*work, {photoId});
    const auto it = datas.find(photoId);
    return
        it == datas.end()
        ? std::optional<PhotoData>()
        : it->second;
}

void
CompanyDataBundleQueue::storePhotoData(const PhotoData& data)
{
    auto work = socialPool_.masterWriteableTransaction();
    work->exec(
        "INSERT INTO sprav.photo_data (photo_id, shooting_point, shooting_target, subject) "
        "VALUES ('" + data.photoId + "',"
        " ST_GeomFromWKB(" + work->quote_raw(toWKB(data.shootingPoint)) + ", 3395), "
        " ST_GeomFromWKB(" + work->quote_raw(toWKB(data.shootingTarget)) + ", 3395), " +
        quote(*work, toString(data.subject)) + ")");
    work->commit();
}

void
CompanyDataBundleQueue::setBundleState(BundleId id, CompanyDataBundle::State state)
{
    auto work = socialPool_.masterWriteableTransaction();
    work->exec(
        "UPDATE sprav.company_bundle SET bundle_state = " +
        quote(*work, toString(state)) + " WHERE "
        "bundle_id = " + std::to_string(id));
    work->commit();
}

void
CompanyDataBundleQueue::markAsSubmitted(BundleId id, FeedbackTaskId taskId)
{
    auto work = socialPool_.masterWriteableTransaction();
    work->exec(
        "UPDATE sprav.company_bundle SET bundle_state = " +
        quote(*work, toString(CompanyDataBundle::State::Submitted)) + ", "
        " submitted_at = now(), feedback_task_id = " +
        std::to_string(taskId) + " WHERE "
        "bundle_id = " + std::to_string(id));
    work->commit();
}

std::vector<CompanyDataBundle>
CompanyDataBundleQueue::load(const std::vector<BundleId>& ids) const
{
    if (ids.empty()) {
        return {};
    }
    auto work = socialPool_.masterReadOnlyTransaction();
    return readBundles(
        work->exec(
            "SELECT * FROM sprav.company_bundle "
            "WHERE " + common::whereClause("bundle_id", ids)),
        *work);
}

std::vector<CompanyDataBundle>
CompanyDataBundleQueue::loadWaitingBatch(BundleId minId, size_t count) const
{
    std::string query;
    query =
        "SELECT * FROM sprav.company_bundle WHERE bundle_state = 'waiting' ";
    if (minId) {
        query += "AND bundle_id > " + std::to_string(minId) + " ";
    }
    query +=
        "ORDER BY created_at ASC ";
    if (count) {
        query += "LIMIT " + std::to_string(count);
    }
    auto work = socialPool_.masterReadOnlyTransaction();
    return readBundles(work->exec(query), *work);
}

size_t
CompanyDataBundleQueue::waitingCount() const
{
    auto work = socialPool_.masterReadOnlyTransaction();
    const auto rows = work->exec(
        "SELECT count(*) FROM sprav.company_bundle "
        "WHERE bundle_state = 'waiting'");
    if (rows.empty() || rows[0][0].is_null()) {
        return 0;
    }
    return rows[0][0].as<size_t>();
}

std::optional<maps::chrono::TimePoint>
CompanyDataBundleQueue::oldestWaitingCreatedAt() const
{
    auto work = socialPool_.masterReadOnlyTransaction();
    const auto rows = work->exec(
        "SELECT max(created_at) FROM sprav.company_bundle "
        "WHERE bundle_state = 'waiting'");
    if (rows.empty() || rows[0][0].is_null()) {
        return {};
    }
    return maps::chrono::parseSqlDateTime(rows[0][0].c_str());
}

DEFINE_ENUM_IO(CompanyDataBundle::State, COMPANY_DATA_BUNDLE_STATE_ENUM_REPRESENTATION);

} // namespace maps::wiki::walkers_export_downloader
