#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/id.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detection/include/store.h>

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

#include <maps/libs/geolib/include/conversion.h>
#include <maps/libs/geolib/include/distance.h>

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

#include <maps/libs/common/include/make_batches.h>

namespace maps::mrc::eye {

DetectionStore::DetectionStore(
            db::IdTo<db::eye::DetectionGroup> groupById,
            db::IdTo<db::eye::Detection> detectionById,
            db::IdTo<db::eye::Frame> frameById,
            db::IdTo<db::eye::FrameLocation> frameLocationByFrameId,
            db::IdTo<db::eye::FramePrivacy> framePrivacyByFrameId,
            db::IdTo<db::eye::Device> deviceById)
    : groupById_(std::move(groupById))
    , detectionById_(std::move(detectionById))
    , frameById_(std::move(frameById))
    , frameLocationByFrameId_(std::move(frameLocationByFrameId))
    , framePrivacyByFrameId_(std::move(framePrivacyByFrameId))
    , deviceById_(std::move(deviceById))
{}

const db::eye::Detection& DetectionStore::detectionById(db::TId detectionId) const
{
    const auto it = detectionById_.find(detectionId);
    REQUIRE(it != detectionById_.end(), "Invalid detection id " << detectionId);
    return it->second;
}

db::TIdSet DetectionStore::detectionIds() const
{
    auto range = IterateKeys(detectionById_);
    return {range.begin(), range.end()};
}

const db::eye::DetectionGroup& DetectionStore::groupById(db::TId groupId) const
{
    const auto it = groupById_.find(groupId);
    REQUIRE(it != groupById_.end(), "Invalid group id " << groupId);
    return it->second;
}

const db::eye::DetectionGroup& DetectionStore::groupByDetectionId(db::TId detectionId) const
{
    return groupById(detectionById(detectionId).groupId());
}

db::TId DetectionStore::frameId(db::TId detectionId) const
{
    return groupByDetectionId(detectionId).frameId();
}

const db::eye::Frame& DetectionStore::frameById(db::TId frameId) const
{
    const auto it = frameById_.find(frameId);
    REQUIRE(it != frameById_.end(), "Invalid frame id " << frameId);
    return it->second;
}

const db::eye::Frame& DetectionStore::frameByDetectionId(db::TId detectionId) const
{
    return frameById(frameId(detectionId));
}

const db::eye::FrameLocation& DetectionStore::locationByFrameId(db::TId frameId) const
{
    const auto it = frameLocationByFrameId_.find(frameId);
    REQUIRE(it != frameLocationByFrameId_.end(), "Invalid frame id " << frameId);
    return it->second;
}

const db::eye::FrameLocation& DetectionStore::locationByDetectionId(db::TId detectionId) const
{
    return locationByFrameId(frameId(detectionId));
}

const db::eye::FramePrivacy& DetectionStore::privacyByFrameId(db::TId frameId) const
{
    const auto it = framePrivacyByFrameId_.find(frameId);
    REQUIRE(it != framePrivacyByFrameId_.end(), "Invalid frame id " << frameId);
    return it->second;
}

const db::eye::FramePrivacy& DetectionStore::privacyByDetectionId(db::TId detectionId) const
{
    return privacyByFrameId(frameId(detectionId));
}

const db::eye::Device& DetectionStore::deviceById(db::TId deviceId) const
{
    const auto it = deviceById_.find(deviceId);
    REQUIRE(it != deviceById_.end(), "Invalid device id " << deviceId);
    return it->second;
}

const db::eye::Device& DetectionStore::deviceByFrameId(db::TId frameId) const
{
    return deviceById(frameById(frameId).deviceId());
}

const db::eye::Device& DetectionStore::deviceByDetectionId(db::TId detectionId) const
{
    return deviceById(frameByDetectionId(detectionId).deviceId());
}

DetectionStore::Slice DetectionStore::slice(const db::TIdSet& detectionIds) const
{
    db::eye::DetectionGroups groups;
    db::eye::Detections detections;
    db::eye::Frames frames;
    db::eye::FrameLocations locations;
    db::eye::Devices devices;

    for (auto detectionId: detectionIds) {
        const auto& detection = detectionById(detectionId);

        detections.push_back(detection);

        const auto& group = groupById(detection.groupId());
        groups.push_back(group);

        const auto& frame = frameById(group.frameId());
        frames.push_back(frame);
        locations.push_back(locationByFrameId(frame.id()));

        devices.push_back(deviceById(frame.deviceId()));
    }

    return {
        std::move(groups),
        std::move(detections),
        std::move(frames),
        std::move(locations),
        std::move(devices)
    };
}


db::eye::DetectionType DetectionStore::getDetectionType(db::TId detectionId) const
{
    return groupByDetectionId(detectionId).type();
}

void DetectionStore::extendByDetections(
    pqxx::transaction_base& txn,
    db::eye::Detections detections)
{
    static const size_t BATCH_SIZE = 10000;

    auto detectionById = byId(std::move(detections));

    db::TIds groupIds;
    for (const auto& [_, detection]: detectionById) {
        groupIds.push_back(detection.groupId());
    }

    db::TIds frameIds;
    for (const auto& batch : maps::common::makeBatches(groupIds, BATCH_SIZE)) {
        db::TIds batchGroupIds{batch.begin(), batch.end()};
        auto groupById = byId(db::eye::DetectionGroupGateway(txn).loadByIds(batchGroupIds));

        for (const auto& [_, group]: groupById) {
            frameIds.push_back(group.frameId());
        }

        groupById_.merge(groupById);
    }

    db::TIds deviceIds;
    for (const auto& batch : maps::common::makeBatches(frameIds, BATCH_SIZE)) {
        db::TIds batchFrameIds{batch.begin(), batch.end()};

        auto frameById = byId(db::eye::FrameGateway(txn).loadByIds(batchFrameIds));
        auto locationByFrameId = byId(db::eye::FrameLocationGateway(txn).loadByIds(batchFrameIds));
        auto privacyByFrameId = byId(db::eye::FramePrivacyGateway(txn).loadByIds(batchFrameIds));

        for (const auto& [_, frame]: frameById) {
            deviceIds.push_back(frame.deviceId());
        }

        frameById_.merge(frameById);
        frameLocationByFrameId_.merge(locationByFrameId);
        framePrivacyByFrameId_.merge(privacyByFrameId);
    }


    for (const auto& batch : maps::common::makeBatches(deviceIds, BATCH_SIZE)) {
        db::TIds batchDeviceIds{batch.begin(), batch.end()};
        auto deviceById = byId(db::eye::DeviceGateway(txn).loadByIds(batchDeviceIds));

        deviceById_.merge(deviceById);
    }

    detectionById_.merge(detectionById);
}

void DetectionStore::extendByDetectionIds(
    pqxx::transaction_base& txn,
    const db::TIds& detectionIds)
{
    namespace table = db::eye::table;

    static const size_t BATCH_SIZE = 10000;

    for (const auto& batch : maps::common::makeBatches(detectionIds, BATCH_SIZE)) {
        auto detections = db::eye::DetectionGateway(txn).load(
            table::Detection::id.in({batch.begin(), batch.end()})
        );

        extendByDetections(txn, std::move(detections));
    }
}

void DetectionStore::extendByDetectionIds(
    pqxx::transaction_base& txn,
    const db::TIdSet& detectionIds)
{
    extendByDetectionIds(txn, db::TIds{detectionIds.begin(), detectionIds.end()});
}

void DetectionStore::extendByObjectRelations(
    pqxx::transaction_base& txn,
    const db::IdTo<db::TIdSet>& relationMap)
{
    db::TIds detectionIds;
    for (const auto& [primaryId, relatedDetectionIds] : relationMap) {
        detectionIds.push_back(primaryId);
        detectionIds.insert(detectionIds.end(),
            relatedDetectionIds.begin(), relatedDetectionIds.end()
        );
    }
    extendByDetectionIds(txn, detectionIds);
}

} // namespace maps::mrc::eye
