#include <maps/wikimap/mapspro/services/mrc/tools/sign-positioning-dataset/common/constants.h>

#include <maps/wikimap/mapspro/libs/common/include/yandex/maps/wiki/common/pgpool3_helpers.h>
#include <maps/wikimap/mapspro/services/mrc/libs/config/include/config.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/feature_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/ride_recording_report_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/sign_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/track_point_gateway.h>
#include <yandex/maps/proto/offline-mrc/results.sproto.h>

#include <maps/libs/chrono/include/time_point.h>
#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/libs/json/include/builder.h>
#include <maps/libs/json/include/prettify.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/geolib/include/distance.h>
#include <maps/libs/geolib/include/point.h>
#include <maps/libs/log8/include/log8.h>
#include <yandex/maps/pb_stream2/reader.h>
#include <yandex/maps/proto/offline_recording/log_event.pb.h>
#include <yandex/maps/proto/offline_recording/mapkit2/location.pb.h>
#include <yandex/maps/proto/offline_recording/record.pb.h>
#include <yandex/maps/proto/offline_recording/report.pb.h>

#include <boost/lexical_cast.hpp>

#include <algorithm>
#include <sstream>
#include <unordered_map>

using namespace maps::mrc;
namespace spresults = yandex::maps::sproto::offline::mrc::results;
namespace precording = yandex::maps::proto::offline::recording;

using RecordingData = std::pair<db::rides::RideRecordingReport, std::string>;
using RecordingsData = std::vector<RecordingData>;

namespace {

const TString SENSORS_VERSION = "sensors3";
const TString SENSOR_TYPE_ACCELEROMETER = "accel";
const TString SENSOR_TYPE_GYROSCOPE = "gyro";
const maps::geolib3::Point2 MOSCOW_CENTER{37.622499, 55.753230};
const double MOSCOW_RADIUS_METERS = 50000;

bool sameRide(const db::rides::RideRecordingReport& prev,
              const db::rides::RideRecordingReport& next)
{
    return prev.sourceId() == next.sourceId()
           && next.startedAt() - prev.finishedAt()
                  <= std::chrono::seconds(60);
}

db::rides::RideRecordingReports
loadRecordingsMetadata(maps::wiki::common::PoolHolder& poolHolder,
                       maps::chrono::TimePoint beginDate,
                       maps::chrono::TimePoint endDate,
                       const std::string& sourceId)
{
    auto txn = poolHolder.pool().slaveTransaction();
    db::rides::RideRecordingReportGateway gtw{*txn};

    using Table = db::rides::table::RideRecordingReport;
    if (sourceId.empty()) {
        return gtw.load(beginDate <= Table::startedAt
                            && Table::finishedAt <= endDate,
                        maps::sql_chemistry::orderBy(Table::sourceId)
                            .orderBy(Table::startedAt));
    }

    return gtw.load(beginDate <= Table::startedAt
                        && Table::finishedAt <= endDate
                        && Table::sourceId == sourceId,
                    maps::sql_chemistry::orderBy(Table::sourceId)
                        .orderBy(Table::startedAt));
}

bool checkRecordingSuites(const RecordingsData& recordingsData)
{
    bool sensorsOK = false;
    bool locationOK = false;

    for (const auto& recordingData : recordingsData) {
        std::istringstream data(recordingData.second);

        maps::pb_stream2::Reader reader(&data);
        for (auto it = reader.begin(); it != reader.end(); ++it) {
            auto record = it->as<precording::record::Record>();
            if (record.HasExtension(precording::log_event::EVENT_RECORD)) {
                auto& eventRecord = record.GetExtension(
                    precording::log_event::EVENT_RECORD);
                if ((eventRecord.event() == SENSOR_TYPE_ACCELEROMETER
                     || eventRecord.event() == SENSOR_TYPE_GYROSCOPE)
                    && eventRecord.component() == SENSORS_VERSION) {
                    sensorsOK = true;
                }
            }
            else if (record.HasExtension(
                         precording::mapkit2::location::LOCATION_RECORD)) {
                const auto& locationRecord = record.GetExtension(
                    precording::mapkit2::location::LOCATION_RECORD);
                if (locationRecord.has_location()
                    && locationRecord.location().has_position()) {
                    const auto& pos = locationRecord.location().position();
                    const auto distance = maps::geolib3::fastGeoDistance(
                        {pos.lon(), pos.lat()}, MOSCOW_CENTER);
                    if (distance <= MOSCOW_RADIUS_METERS) {
                        locationOK = true;
                    }
                }
            }

            if (sensorsOK && locationOK) {
                return true;
            }
        }
    }
    return false;
}

db::Features loadImagesMetadata(maps::wiki::common::PoolHolder& poolHolder,
                                const RecordingsData& recordingsData)
{
    INFO() << "Loading features from MRC DB";

    REQUIRE(!recordingsData.empty(),
            "Stumbled over an empty recordings list");
    const auto sourceId = recordingsData.front().first.sourceId();
    const auto beginDate = recordingsData.front().first.startedAt();
    const auto endDate = recordingsData.back().first.finishedAt();

    auto txn = poolHolder.pool().slaveTransaction();
    db::FeatureGateway gtw{*txn};
    using Table = db::table::Feature;

    auto features
        = gtw.load(Table::sourceId == sourceId && beginDate <= Table::date
                       && Table::date <= endDate && Table::isPublished
                       && Table::pos != std::nullopt,
                   maps::sql_chemistry::orderBy(Table::date));

    // WORKAROUND: eliminate features with the same date to avoid having
    //             duplicates (it is a known bug atm).
    const auto nonUniqueBegin
        = std::unique(features.begin(), features.end(),
                      [](const db::Feature& lhs, const db::Feature& rhs) {
                          return lhs.timestamp() == rhs.timestamp();
                      });
    features.erase(nonUniqueBegin, features.end());

    return features;
}

std::string makeFileNamePrefix(const RecordingsData& recordingsData)
{
    const auto& recordingMetadata = recordingsData.front().first;
    return recordingMetadata.sourceId() + "-"
           + std::to_string(
                 recordingMetadata.startedAt().time_since_epoch().count());
}

struct BBox {
    std::string type;
    std::int32_t minX;
    std::int32_t minY;
    std::int32_t maxX;
    std::int32_t maxY;
};
using FeatureIdToBBoxes = std::unordered_multimap<db::TId, BBox>;

FeatureIdToBBoxes loadBBoxes(maps::wiki::common::PoolHolder& poolHolder,
                             db::Features& imagesMetadata)
{
    INFO() << "Obtaining the images bounding boxes";

    auto txn = poolHolder.pool().slaveTransaction();

    db::TIds featureIds;
    featureIds.reserve(imagesMetadata.size());
    for (const auto& feature : imagesMetadata) {
        featureIds.push_back(feature.id());
    }

    auto featureSigns = db::SignFeatureGateway{*txn}.load(
        db::table::SignFeature::featureId.in(featureIds));

    db::TIds signIds;
    for (const auto& featureSign : featureSigns) {
        signIds.push_back(featureSign.signId());
    }
    std::sort(signIds.begin(), signIds.end());
    signIds.erase(std::unique(signIds.begin(), signIds.end()), signIds.end());

    std::unordered_map<db::TId, std::string> signIdToSignType;
    auto signs = db::SignGateway{*txn}.load(db::table::Sign::id.in(signIds));
    for (const auto& sign : signs) {
        signIdToSignType[sign.id()]
            = boost::lexical_cast<std::string>(sign.type());
    }

    FeatureIdToBBoxes featureIdToBBoxes;
    for (auto&& featureSign : featureSigns) {
        featureIdToBBoxes.emplace(
            featureSign.featureId(),
            BBox{signIdToSignType[featureSign.signId()], featureSign.minX(),
                 featureSign.minY(), featureSign.maxX(), featureSign.maxY()});
    }

    return featureIdToBBoxes;
}

void filterFeaturesWithBBoxes(db::Features& imagesMetadata,
                              const FeatureIdToBBoxes& featureIdToBBoxes)
{
    INFO() << "Filtering images with bounding boxes";
    const auto noBBoxesBegin
        = std::remove_if(imagesMetadata.begin(), imagesMetadata.end(),
                         [&](const auto& feature) {
                             return !featureIdToBBoxes.count(feature.id());
                         });
    imagesMetadata.erase(noBBoxesBegin, imagesMetadata.end());
}

void saveBBoxes(db::Features& imagesMetadata,
                const FeatureIdToBBoxes& featureIdToBBoxes,
                const std::string& filenamePrefix)
{
    if (imagesMetadata.empty() || featureIdToBBoxes.empty()) {
        return;
    }

    INFO() << "Saving bboxes into " << filenamePrefix
           << dataset::IMAGES_META_FILENAME_POSTFIX;

    // Save bboxes into a file as JSON in the same order as features go.
    // [{
    //    "feature_id" : integer,
    //    "feature_ts" : timestamp,
    //    "exif_orientation" : integer,
    //    "quality" : float,
    //    "signs" : [
    //        {
    //          "sign_type" : "sign_type",
    //          "sign_id" : integer,
    //          "bbox" : [ x, y, widht, height ]
    //        },
    //        ...
    //    ]
    // },
    // ...]
    maps::json::Builder builder;
    builder << [&](maps::json::ArrayBuilder builder) {
        for (const auto& feature : imagesMetadata) {
            if (!feature.hasQuality()) {
                continue;
            }

            builder << [&](maps::json::ObjectBuilder builder) {

                builder["feature_id"] << feature.id();
                builder["feature_ts"]
                    << feature.timestamp().time_since_epoch().count();
                builder["quality"] << feature.quality();

                if (feature.hasOrientation()) {
                    builder["exif_orientation"]
                        = static_cast<int>(feature.orientation());
                }
                else {
                    // assuming an image is in normal orientation.
                    builder["exif_orientation"] = 1;
                }

                builder["signs"] << [&](maps::json::ArrayBuilder builder) {
                    auto[featureIdAndBBoxIt, end]
                        = featureIdToBBoxes.equal_range(feature.id());
                    for (; featureIdAndBBoxIt != end; ++featureIdAndBBoxIt) {
                        const auto& bbox = featureIdAndBBoxIt->second;
                        builder << [&](maps::json::ObjectBuilder builder) {
                            builder["sign_id"] << "unknown";
                            builder["sign_type"] << bbox.type;
                            builder["bbox"]
                                << [&](maps::json::ArrayBuilder builder) {
                                       builder << bbox.minX << bbox.minY
                                               << bbox.maxX << bbox.maxY;
                                   };
                        };
                    }
                };
            };
        }
    };

    const auto result = maps::json::prettifyJson(builder.str());
    std::ofstream out{filenamePrefix + dataset::IMAGES_META_FILENAME_POSTFIX};
    out << result << std::endl;
}

void putImages(spresults::Results& results,
               maps::mds::Mds& mdsClient,
               const db::Features& imagesMetadata)
{
    for (const auto& feature : imagesMetadata) {
        spresults::Image sprotoImage;

        sprotoImage.created()
            = feature.timestamp().time_since_epoch().count();
        sprotoImage.image() = mdsClient.get(feature.mdsKey());

        if (feature.hasPos()) {
            const auto geodeticPos = feature.geodeticPos();
            sprotoImage.estimatedPosition()->point().lon() = geodeticPos.x();
            sprotoImage.estimatedPosition()->point().lat() = geodeticPos.y();
        }

        if (feature.hasHeading()) {
            sprotoImage.estimatedPosition()->heading()
                = feature.heading().value();
        }

        results.images().push_back(std::move(sprotoImage));
    }
}

void putTrackPoints(spresults::Results& results,
                    maps::wiki::common::PoolHolder& poolHolder,
                    const RecordingsData& recordingsData)
{
    using Table = db::table::TrackPoint;

    const auto beginDate = recordingsData.front().first.startedAt();
    const auto endDate = recordingsData.back().first.finishedAt();
    const auto sourceId = recordingsData.back().first.sourceId();

    auto txn = poolHolder.pool().slaveTransaction();
    const auto trackPoints = db::TrackPointGateway{*txn}.load(
        beginDate <= Table::timestamp && Table::timestamp <= endDate
            && Table::sourceId == sourceId,
        maps::sql_chemistry::orderBy(Table::timestamp));

    for (const auto& trackPoint : trackPoints) {
        spresults::TrackPoint sprotoTrackPoint;

        const auto geodeticPos = trackPoint.geodeticPos();
        sprotoTrackPoint.location().point().lon() = geodeticPos.x();
        sprotoTrackPoint.location().point().lat() = geodeticPos.y();

        if (trackPoint.speedMetersPerSec()) {
            sprotoTrackPoint.location().speed()
                = *trackPoint.speedMetersPerSec();
        }

        if (trackPoint.heading()) {
            sprotoTrackPoint.location().heading()
                = trackPoint.heading()->value();
        }

        if (trackPoint.accuracyMeters()) {
            sprotoTrackPoint.location().accuracy()
                = *trackPoint.accuracyMeters();
        }

        sprotoTrackPoint.time()
            = trackPoint.timestamp().time_since_epoch().count();

        results.track().push_back(std::move(sprotoTrackPoint));
    }
}

void putRecordings(spresults::Results& results,
                   const RecordingsData& recordingsData)
{
    for (const auto& recordingData : recordingsData) {
        results.reports().push_back(recordingData.second);
    }
}

void saveResults(maps::wiki::common::PoolHolder& poolHolder,
                 maps::mds::Mds& mdsClient,
                 const RecordingsData& recordingsData,
                 const db::Features& imagesMetadata,
                 const std::string& filenamePrefix)
{
    if (imagesMetadata.empty()) {
        return;
    }

    INFO() << "Loading a ride data";

    spresults::Results results;

    putImages(results, mdsClient, imagesMetadata);
    putTrackPoints(results, poolHolder, recordingsData);
    putRecordings(results, recordingsData);

    INFO() << "Saving the ride data into " << filenamePrefix
           << dataset::RESULTS_FILENAME_POSTFIX;
    std::ofstream out{filenamePrefix + dataset::RESULTS_FILENAME_POSTFIX,
                      std::ios::binary};
    out << results;
    out.flush();
}

} // namespace

int main(int argc, char** argv)
{
    maps::cmdline::Parser parser(
        "A tool to load Moscow ride results with sensors");
    auto configPath = parser.string("config").help("path to configuration");
    auto beginDateOpt
        = parser.string("begin-date")
              .help("date in PG format e.g. \"2019-01-01 12:00:00\"")
              .required();
    auto endDateOpt
        = parser.string("end-date")
              .help("date in PG format e.g. \"2019-01-02 12:00:00\"");
    auto sourceIdOpt
        = parser.string("source-id").help("load data of this source ID");
    parser.parse(argc, argv);

    INFO() << "Loading ride datasets from time period" << beginDateOpt
           << " - " << endDateOpt;

    const auto beginDate = maps::chrono::parseSqlDateTime(beginDateOpt);
    const auto endDate = endDateOpt.defined()
                             ? maps::chrono::parseSqlDateTime(endDateOpt)
                             : maps::chrono::TimePoint::clock::now();

    auto cfg = maps::mrc::common::templateConfigFromCmdPath(configPath);
    auto mdsClient = cfg.makeMdsClient();
    auto poolHolder = cfg.makePoolHolder();

    auto recordingsMetadata = loadRecordingsMetadata(
        poolHolder, beginDate, endDate,
        sourceIdOpt.defined() ? static_cast<std::string>(sourceIdOpt)
                              : std::string{});

    if (recordingsMetadata.empty()) {
        INFO() << "No recordings for the specified date. Exiting.";
    }

    RecordingsData recordingsData;

    const auto collectRide = [&](const RecordingsData& recordingsData) {
        if (recordingsData.empty()) {
            return;
        }

        if (!checkRecordingSuites(recordingsData)) {
            INFO() << "No sensors data in recordings, skipping...";
            return;
        }

        INFO() << "Collect ride data for recordings "
               << recordingsData.front().first.sourceId()
               << " for time period "
               << maps::chrono::formatSqlDateTime(
                      recordingsData.front().first.startedAt())
               << " - " << maps::chrono::formatSqlDateTime(
                               recordingsData.back().first.finishedAt());

        const auto filenamePrefix = makeFileNamePrefix(recordingsData);

        auto imagesMetadata = loadImagesMetadata(poolHolder, recordingsData);
        auto featureIdToBBoxes = loadBBoxes(poolHolder, imagesMetadata);
        filterFeaturesWithBBoxes(imagesMetadata, featureIdToBBoxes);

        saveBBoxes(imagesMetadata, featureIdToBBoxes, filenamePrefix);
        saveResults(poolHolder, mdsClient, recordingsData, imagesMetadata,
                    filenamePrefix);
    };

    for (const auto& recordingMetadata : recordingsMetadata) {

        const auto recording = mdsClient.get(
            {recordingMetadata.mdsGroupId(), recordingMetadata.mdsPath()});

        if (!recordingsData.empty()
            && !sameRide(recordingsData.back().first, recordingMetadata)) {
            collectRide(recordingsData);
            recordingsData.clear();
        }

        INFO() << "Loaded recordings of " << recordingMetadata.sourceId()
               << " for time period " << maps::chrono::formatSqlDateTime(
                                             recordingMetadata.startedAt())
               << " - " << maps::chrono::formatSqlDateTime(
                               recordingMetadata.finishedAt());
        recordingsData.emplace_back(recordingMetadata, recording);
    }
    collectRide(recordingsData);

    INFO() << "Done with datasets loading";

    return EXIT_SUCCESS;
}
