#include <maps/wikimap/mapspro/services/mrc/eye/lib/feedback/include/worker.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/feedback/include/metadata.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/feedback/include/push_feedback.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/txn.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/id.h>

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

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

namespace maps::mrc::eye {

namespace {

db::eye::Hypotheses merge(
    db::eye::Hypotheses hypotheses,
    const db::eye::Hypotheses& newHypotheses)
{
    hypotheses.insert(
        hypotheses.end(),
        newHypotheses.begin(), newHypotheses.end()
    );
    std::sort(hypotheses.begin(), hypotheses.end(),
        [](const auto& lhs, const auto& rhs) {
            return lhs.id() < rhs.id();
        }
    );
    hypotheses.erase(
        std::unique(hypotheses.begin(), hypotheses.end(),
            [](const auto& lhs, const auto& rhs) {
                return lhs.id() == rhs.id();
            }
        ),
        hypotheses.end()
    );
    return hypotheses;
}

std::pair<db::TId, db::TId> getInterval(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    const db::eye::HypothesisTypes& types,
    size_t limit)
{
    db::TId endTxnId = beginTxnId;

    const auto hypothesisTxnIds = db::eye::HypothesisGateway(txn).loadTxnIds(
        db::eye::table::Hypothesis::txnId >= beginTxnId
            and db::eye::table::Hypothesis::type.in(types)
            and not db::eye::table::Hypothesis::deleted,
        sql_chemistry::limit(limit).orderBy(db::eye::table::Hypothesis::txnId)
    );

    const auto groupTxnIds = db::eye::DetectionGroupGateway(txn).loadTxnIds(
        db::eye::table::DetectionGroup::txnId >= beginTxnId
            and db::eye::table::DetectionGroup::approved,
        sql_chemistry::limit(limit).orderBy(db::eye::table::DetectionGroup::txnId)
    );

    if (!hypothesisTxnIds.empty() && !groupTxnIds.empty()) {
        endTxnId = std::min(hypothesisTxnIds.back(), groupTxnIds.back()) + 1;
    } else if (!hypothesisTxnIds.empty()) {
        endTxnId = hypothesisTxnIds.back() + 1;
    } else if (!groupTxnIds.empty()) {
        endTxnId = groupTxnIds.back() + 1;
    }

    return {beginTxnId, endTxnId};
}

db::eye::Hypotheses loadBatch(
    pqxx::transaction_base& txn,
    const db::TIds& hypothesisIds,
    const db::eye::HypothesisTypes& hypothesisTypes)
{
    return db::eye::HypothesisGateway(txn).load(
        db::eye::table::Hypothesis::id.in(hypothesisIds)
        and db::eye::table::Hypothesis::type.in(hypothesisTypes)
        and not db::eye::table::Hypothesis::deleted
    );
}

db::eye::Hypotheses loadBatchForLoopMode(
    pqxx::transaction_base& txn,
    db::TId beginTxnId, db::TId endTxnId,
    const db::eye::HypothesisTypes& hypothesisTypes)
{
    db::eye::Hypotheses hypotheses = db::eye::HypothesisGateway(txn).load(
        db::eye::table::Hypothesis::txnId >= beginTxnId
        and db::eye::table::Hypothesis::txnId < endTxnId
        and db::eye::table::Hypothesis::type.in(hypothesisTypes)
        and not db::eye::table::Hypothesis::deleted
    );

    db::eye::Hypotheses hypothesesFromGroups = db::eye::HypothesisGateway(txn).load(
        db::eye::table::Hypothesis::id == db::eye::table::HypothesisObject::hypothesisId
        and not db::eye::table::Hypothesis::deleted
        and db::eye::table::Object::id == db::eye::table::HypothesisObject::objectId
        and not db::eye::table::Object::deleted
        and db::eye::table::Object::primaryDetectionId == db::eye::table::Detection::id
        and not db::eye::table::Detection::deleted
        and db::eye::table::Detection::groupId == db::eye::table::DetectionGroup::id
        and db::eye::table::DetectionGroup::txnId >= beginTxnId
        and db::eye::table::DetectionGroup::txnId < endTxnId
        and db::eye::table::DetectionGroup::approved
    );

    return merge(hypotheses, hypothesesFromGroups);
}

db::eye::Hypotheses filterUnpublishedHypotheses(
    pqxx::transaction_base& txn,
    db::eye::Hypotheses hypotheses)
{
    db::TIds hypothesisIds = collectIds(hypotheses);

    auto hypothesisFeedbacks = db::eye::HypothesisFeedbackGateway(txn).load(
        db::eye::table::HypothesisFeedback::hypothesisId.in(hypothesisIds)
    );
    db::TIdSet publishedHypothesesIds;
    for (const auto& hypothesisFeedback : hypothesisFeedbacks) {
        publishedHypothesesIds.insert(hypothesisFeedback.hypothesisId());
    }
    auto isAlreadyPublished = [&](db::TId hypothesisId) {
        return publishedHypothesesIds.count(hypothesisId);
    };

    hypotheses.erase(
        std::remove_if(hypotheses.begin(), hypotheses.end(),
            [&](const auto& hypothesis) {
                return isAlreadyPublished(hypothesis.id());
            }
        ),
        hypotheses.end()
    );

    return hypotheses;
}

db::eye::RecognitionSource
selectRecognitionRequestSource(db::FeaturePrivacy privacy)
{
    if (privacy == db::FeaturePrivacy::Public) {
        return db::eye::RecognitionSource::Toloka;
    } else {
        return db::eye::RecognitionSource::Yang;
    }
}

db::eye::Recognitions createRecognitionRequests(const HypothesisContext::Items& items) {
    constexpr db::TId VERSION = 0u;

    db::eye::Recognitions requests;
    for (const auto& item : items) {
        requests.emplace_back(
            item.frame.id(),
            item.frame.orientation(),
            toRecognitionType(item.detectionGroup.type()),
            selectRecognitionRequestSource(item.framePrivacy.type()),
            VERSION
        );
    }
    return requests;
}

using RecognitionTuple
    = std::tuple<
        db::TId,
        common::ImageOrientation,
        db::eye::RecognitionType,
        db::eye::RecognitionSource,
        db::TId
    >;

RecognitionTuple makeRecognitionTuple(const db::eye::Recognition& recognition) {
    return std::tuple(
        recognition.frameId(),
        recognition.orientation(),
        recognition.type(),
        recognition.source(),
        recognition.version()
    );
}

bool basesOnObjectNewerThan(
    const HypothesisContext& context,
    chrono::TimePoint timeThreshold)
{
    for (const auto& item : context.items()) {
        if (item.frame.time() < timeThreshold) {
            return false;
        }
    }

    return true;
}

bool isInRussia(
    const db::eye::Hypothesis& hypothesis,
    const privacy::GeoIdProvider& geoIdProvider)
{
    static const db::TId RUSSIA_GEO_ID = 225;

    const auto geoIds = geoIdProvider.load(hypothesis.geodeticPos());
    return std::any_of(
        geoIds.begin(), geoIds.end(),
        [&](auto geoId) {
            return geoId == RUSSIA_GEO_ID;
        }
    );
}

bool shouldPublishHypothesisOnTrafficSign(
    const db::eye::Hypothesis& hypothesis,
    const HypothesisContext& context,
    const privacy::GeoIdProvider& geoIdProvider)
{
    // All hypotheses about trucks traffic signs that were seen before 2020-09-01
    // have been processed. In order not to generate duplicate hypotheses skip
    // all the signs that were seen before this date.
    static const chrono::TimePoint TRUCKS_TRAFFIC_SIGN_TIME_THRESHOLD
        = chrono::parseIsoDateTime("2020-09-01T00:00:00+03");
    static const chrono::TimePoint TRAFFIC_SIGN_TIME_THRESHOLD
        = chrono::parseIsoDateTime("2021-09-01T00:00:00+03");

    auto attrs = hypothesis.attrs<db::eye::TrafficSignAttrs>();
    switch (attrs.feedbackType) {
        case wiki::social::feedback::Type::TrucksProhibitedSign:
        case wiki::social::feedback::Type::TrucksSpeedLimitSign:
        case wiki::social::feedback::Type::TrucksManeuverRestrictionSign:
        case wiki::social::feedback::Type::WeightLimitingSign:
        case wiki::social::feedback::Type::DimensionsLimitingSign:
            return basesOnObjectNewerThan(context, TRUCKS_TRAFFIC_SIGN_TIME_THRESHOLD)
                && isInRussia(hypothesis, geoIdProvider);
        default:
            return basesOnObjectNewerThan(context, TRAFFIC_SIGN_TIME_THRESHOLD);
    }
}

bool shouldPublishHypothesisOnParkingSign(
    const db::eye::Hypothesis& hypothesis,
    const HypothesisContext& context)
{
    static const chrono::TimePoint PARKING_TRAFFIC_SIGN_TIME_THRESHOLD
        = chrono::parseIsoDateTime("2021-09-01T00:00:00+03");

    if (hypothesis.type() == db::eye::HypothesisType::AbsentParking) {
        if (hypothesis.attrs<db::eye::AbsentParkingAttrs>().isToll) {
            return true;
        }
    }

    return basesOnObjectNewerThan(context, PARKING_TRAFFIC_SIGN_TIME_THRESHOLD);
}

bool shouldPublishHypothesisOnSpeedBump(
    const db::eye::Hypothesis& /*hypothesis*/,
    const HypothesisContext& /*context*/)
{
    // Публикуем всегда, так как это новый тип гипотез
    return true;
}

bool shouldPublishHypothesis(
    const db::eye::Hypothesis& hypothesis,
    const HypothesisContext& context,
    const privacy::GeoIdProvider& geoIdProvider)
{
    // Не публикуем гипотезы, если они основаны на объектах, которые
    // существовали до 1 сентября 2021 года, если для данного типа гипотез
    // не задано уникального правила. Ограничение добавлено, чтобы избежать
    // дублей с гипотезами, которые мог сгенерить signs_detector
    static const chrono::TimePoint COMMON_TIME_THRESHOLD
        = chrono::parseIsoDateTime("2021-09-01T00:00:00+03");

    switch (hypothesis.type()) {
        case db::eye::HypothesisType::TrafficSign:
            return shouldPublishHypothesisOnTrafficSign(hypothesis, context, geoIdProvider);
        case db::eye::HypothesisType::AbsentParking:
        case db::eye::HypothesisType::WrongParkingFtType:
            return shouldPublishHypothesisOnParkingSign(hypothesis, context);
        case db::eye::HypothesisType::SpeedBump:
            return shouldPublishHypothesisOnSpeedBump(hypothesis, context);
        default:
            return basesOnObjectNewerThan(context, COMMON_TIME_THRESHOLD);
    }
}

} // namespace

void PushFeedbackWorker::processBatch(const db::TIds& hypothesisIds) {
    HypothesisStore store;

    {
        auto readTxn = getSlaveTxn();

        db::eye::Hypotheses hypotheses = loadBatch(
            *readTxn, hypothesisIds, config_.hypothesisTypes
        );

        store = HypothesisStore::fromHypotheses(
            *readTxn,
            filterUnpublishedHypotheses(*readTxn, std::move(hypotheses))
        );
    }

    auto [hypothesesToPublish, tolokaRequests] = process(store);

    pushHypotheses(hypothesesToPublish);

    auto readTxn = getSlaveTxn();
    auto writeTxn = getMasterWriteTxn(*(config_.mrc.pool));
    saveTolokaRequests(*readTxn, *writeTxn, tolokaRequests);
    commitIfNeed(*writeTxn);
}

bool PushFeedbackWorker::processBatchInLoopMode(size_t batchSize) {
    HypothesisStore store;
    db::TId beginTxnId;
    db::TId endTxnId;

    {
        auto readTxn = getSlaveTxn();

        std::tie(beginTxnId, endTxnId) = getInterval(
            *readTxn, pushFeedbackMetadata(*readTxn).getTxnId(),
            config_.hypothesisTypes, batchSize
        );

        db::eye::Hypotheses hypotheses = loadBatchForLoopMode(
            *readTxn, beginTxnId, endTxnId, config_.hypothesisTypes
        );

        store = HypothesisStore::fromHypotheses(
            *readTxn,
            filterUnpublishedHypotheses(*readTxn, std::move(hypotheses))
        );
    }

    auto [hypothesesToPublish, tolokaRequests] = process(store);

    pushHypotheses(hypothesesToPublish);

    auto readTxn = getSlaveTxn();
    auto writeTxn = getMasterWriteTxn(*(config_.mrc.pool));
    saveTolokaRequests(*readTxn, *writeTxn, tolokaRequests);

    auto metadata = pushFeedbackMetadata(*writeTxn);
    metadata.updateTxnId(endTxnId);
    metadata.updateTime();

    commitIfNeed(*writeTxn);

    return beginTxnId != endTxnId;
}

void PushFeedbackWorker::pushHypotheses(
    const std::vector<std::pair<db::eye::Hypothesis, HypothesisContext>>& hypothesesToPublish)
{
    if (!config_.mrc.commit || !config_.pushToSocial) {
        INFO() << "Skip push hypotheses to social";
        return;
    }

    for (const auto& [hypothesis, context] : hypothesesToPublish) {
        INFO() << "Push hypothesis " << hypothesis.id() << " to social";
        db::TId feedbackId = pushFeedback(
            hypothesis, context,
            *(config_.frameUrlResolver), *(config_.feedbackUrlResolver)
        );

        db::eye::HypothesisFeedback hypothesisFeedback(hypothesis.id(), feedbackId);
        auto writeTxn = getMasterWriteTxn(*(config_.mrc.pool));
        INFO() << "Save feedback with " << feedbackId
               << " for hypothesis " << hypothesis.id();
        db::eye::HypothesisFeedbackGateway(*writeTxn).insertx(hypothesisFeedback);
        commitIfNeed(*writeTxn);
    }
}

PushFeedbackWorkerResult PushFeedbackWorker::process(const HypothesisStore& store) const {
    PushFeedbackWorkerResult result;

    for (const auto& [hypothesis, context] : store) {
        if (context.empty()) {
            WARN() << "Hypothesis " << hypothesis.id() << " has empty context";
            continue;
        }

        if (!shouldPublishHypothesis(hypothesis, context, *config_.geoIdProvider)) {
            WARN() << "Ignore hypothesis " << hypothesis.id();
            continue;
        }

        auto unapprovedItems = context.getItemsToApprove();

        if (!unapprovedItems.empty()) {
            db::eye::Recognitions recognitionRequests = createRecognitionRequests(unapprovedItems);
            result.recognitionRequests.insert(
                result.recognitionRequests.end(),
                recognitionRequests.begin(), recognitionRequests.end()
            );
            INFO() << "Created " << recognitionRequests.size()
                   << " recognition requests for hypothesis " << hypothesis.id();
        } else {
            result.hypothesesToPublish.emplace_back(hypothesis, context);
            INFO() << "Need publish hypothesis " << hypothesis.id();
        }
    }

    return result;
}

db::eye::Recognitions PushFeedbackWorker::removeDuplicate(
    pqxx::transaction_base& txn,
    db::eye::Recognitions requests)
{
    std::sort(requests.begin(), requests.end(),
        [](const auto& lhs, const auto& rhs) {
            return makeRecognitionTuple(lhs) < makeRecognitionTuple(rhs);
        }
    );
    requests.erase(
        std::unique(requests.begin(), requests.end(),
            [](const auto& lhs, const auto& rhs) {
                return makeRecognitionTuple(lhs) == makeRecognitionTuple(rhs);
            }
        ),
        requests.end()
    );

    db::TIds frameIds;
    for (const auto& request : requests) {
        frameIds.push_back(request.frameId());
    }

    auto dbRequests = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::frameId.in(frameIds)
            and db::eye::table::Recognition::source.in({
                db::eye::RecognitionSource::Toloka,
                db::eye::RecognitionSource::Yang
            })
    );

    std::set<RecognitionTuple> tupleSet;
    for (const auto& dbRequest : dbRequests) {
        tupleSet.insert(makeRecognitionTuple(dbRequest));
    }

    requests.erase(
        std::remove_if(requests.begin(), requests.end(),
            [&](const auto& request) {
                return tupleSet.count(makeRecognitionTuple(request));
            }
        ),
        requests.end()
    );

    return requests;
}

void PushFeedbackWorker::saveTolokaRequests(
    pqxx::transaction_base& readTxn,
    pqxx::transaction_base& writeTxn,
    db::eye::Recognitions& tolokaRequests)
{
    if (!config_.approveInToloka) {
        INFO() << "Skip save toloka requests";
        return;
    }

    tolokaRequests = removeDuplicate(readTxn, std::move(tolokaRequests));
    db::eye::RecognitionGateway(writeTxn).insertx(tolokaRequests);
    INFO() << "Save " <<  tolokaRequests.size() << " new requests to toloka";
}

} // namespace maps::mrc::eye
