#include "generate_tasks.h"

#include <maps/libs/log8/include/log8.h>
#include <maps/libs/sql_chemistry/include/batch_load.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/algorithm/retry.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/ugc/gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/ugc/tasks_group_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/check_generator_result.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/data_model.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/deadends.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/districts.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/draw_optimization.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/geometry.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/static_graph_loader.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/tasks_generator.h>
#include <maps/libs/common/include/make_batches.h>
#include <maps/libs/geolib/include/bounding_box.h>
#include <maps/libs/geolib/include/const.h>
#include <maps/libs/geolib/include/conversion.h>
#include <maps/libs/geolib/include/convex_hull.h>
#include <maps/libs/geolib/include/prepared_polygon.h>
#include <maps/libs/geolib/include/serialization.h>
#include <maps/libs/geolib/include/spatial_relation.h>
#include <maps/libs/geolib/include/static_geometry_searcher.h>

#include <boost/lexical_cast.hpp>

#include <algorithm>
#include <iterator>
#include <vector>
#include <unordered_set>
#include <unordered_map>
#include <sstream>

namespace maps::mrc::gen_targets {

namespace {

geolib3::BoundingBox
expandedBoundingBox(const geolib3::MultiPolygon2& geodeticGeom)
{
    constexpr double EXPANSION_METERS = 5000;
    auto bbox = geodeticGeom.boundingBox();

    return geolib3::BoundingBox(
        fastGeoShift(bbox.lowerCorner(),
                     geolib3::Vector2(-EXPANSION_METERS, -EXPANSION_METERS)),
        fastGeoShift(bbox.upperCorner(),
                     geolib3::Vector2(EXPANSION_METERS, EXPANSION_METERS)));
}

std::pair<object::MrcRegions, object::MrcRegions>
loadResidentialAndRestrictedRegions(
    object::Loader& loader, const db::ugc::TasksGroup& tasksGroup)
{
    std::pair<object::MrcRegions, object::MrcRegions> result;

    for (auto&& region : loader.loadMrcRegions(tasksGroup.mercatorGeom().boundingBox())) {
        if (region.type() == object::MrcRegion::Type::Residential &&
            !tasksGroup.ignorePrivateArea()) {
            result.first.push_back(std::move(region));
        } else if (region.type() == object::MrcRegion::Type::Restricted) {
            result.second.push_back(std::move(region));
        }
    }
    return result;
}

RoadNetworkData
loadRoadGraphFragmentForTasksGroup(StaticGraphLoader& graphLoader,
                                   const db::ugc::TasksGroup& group)
{
    const size_t ALLOW_UTURNS_ON_CROSSROADS_WHERE_FC_GE = 7;

    return graphLoader.loadGraph(
            expandedBoundingBox(geolib3::convertMercatorToGeodetic(
                group.mercatorGeom()
            )),
            ALLOW_UTURNS_ON_CROSSROADS_WHERE_FC_GE /* allow uturns on crossroads where fc >= 7 */,
            group.useToll() == db::ugc::UseToll::Yes);
}


std::string makeTaskName(const std::string& baseName, size_t idx)
{
    return baseName + "_" + std::to_string(idx);
}

std::chrono::seconds estimateTaskDuration(double distance)
{
    constexpr double AVG_SPEED_MPS = 10.;
    return std::chrono::seconds(
        static_cast<long int>(distance / AVG_SPEED_MPS));
}

geolib3::Polygon2 makeConvexHull(const PathWithData& path)
{
    const double BUFFER_WIDTH_GEODETIC_GRAD =
        360. * 10. / geolib3::WGS84_EQUATOR_LENGTH; // 10 meters
    size_t pointsNumber = 0;
    for (const auto& pathEdge : path) {
        pointsNumber += pathEdge.edge.geom.pointsNumber();
    }

    geolib3::PointsVector allPoints;
    allPoints.reserve(pointsNumber);

    for (const auto& pathEdge : path) {
        const auto& points = pathEdge.edge.geom.points();
        std::copy(points.begin(), points.end(),
                  std::back_inserter(allPoints));
    }

    return geolib3::bufferedConvexHull(allPoints, BUFFER_WIDTH_GEODETIC_GRAD);
}

template <typename Enum>
auto toIntegral(const std::vector<Enum>& values)
    -> std::vector<std::underlying_type_t<Enum>>
{
    std::vector<std::underlying_type_t<Enum>> result;
    result.reserve(values.size());

    for (auto value : values) {
        result.push_back(static_cast<std::underlying_type_t<Enum>>(value));
    }
    return result;
}

db::ugc::Task makeTask(db::TId tasksGroupId,
                       const std::string& taskName,
                       const PathWithData& path,
                       bool withRouting,
                       const std::vector<db::CameraDeviation>& cameraDeviations)
{
    ASSERT(path.size() > 0);
    const Locale RUSSIAN = boost::lexical_cast<Locale>("ru");
    db::ugc::Task task;
    task.setTasksGroupId(tasksGroupId)
        .addName(RUSSIAN, taskName)
        .setStatus(db::ugc::TaskStatus::Draft)
        .setCameraDeviations(cameraDeviations);

    double distance = 0.;
    uint32_t edgeIdx = 0;

    for (const auto& pathEdge : path) {
        const Edge& edge = pathEdge.edge;
        distance += edge.length;

        task.addTarget(
            edge.geom,
            db::ugc::Oneway::No, // FIXME
            db::ugc::Traffic::RightHand, // TODO add option to tasksGroup
            db::ugc::Direction::Forward,
            withRouting ? std::make_optional(edgeIdx) : std::nullopt,
            std::nullopt /* backwardPos */);
        ++edgeIdx;
    }

    task.setDistanceInMeters(distance)
        .setDuration(estimateTaskDuration(distance))
        .setGeodeticHull(geolib3::MultiPolygon2({makeConvexHull(path)}));

    return task;
}

void saveTasks(pqxx::transaction_base& txn,
               const db::ugc::TasksGroup& group,
               const PathsWithData& paths)
{
    size_t taskIdx = 0;
    db::ugc::TaskGateway gtw(txn);
    for (const auto& path : paths) {
        if (path.empty()) {
            continue;
        }
        ++taskIdx;
        auto task = makeTask(group.id(),
                             makeTaskName(group.name(), taskIdx), path,
                             group.useRouting() == db::ugc::UseRouting::Yes,
                             group.cameraDeviations());
        gtw.insert(task);
    }
}

void calculateTasksGroupStats(
        const RoadNetworkData& roadNetwork,
        const std::unordered_set<EdgeId>& targetEdgeIds,
        const PathsWithData& paths,
        db::ugc::TasksGroup& tasksGroup)
{
    std::unordered_set<EdgeId> countedTargetEdges;

    double totalTasksLength = 0.;
    double uniqueCoveredTargetsLength = 0.;
    double targetsLength = 0.;

    for (auto edgeId : targetEdgeIds) {
        targetsLength += roadNetwork.edge(edgeId).length;
    }

    for (const auto& path : paths) {
        for(auto pathEdge : path) {
            const auto& edge = pathEdge.edge;
            totalTasksLength += edge.length;
            if (pathEdge.isUsefulAsTarget && !countedTargetEdges.count(pathEdge.edge.id)) {
                uniqueCoveredTargetsLength += edge.length;
                countedTargetEdges.insert(pathEdge.edge.id);
            }
        }
    }

    double coverageRatio = targetsLength > 0. ? uniqueCoveredTargetsLength / targetsLength : 0.;

    tasksGroup.setTotalLengthMeters(totalTasksLength)
        .setUniqueLengthMeters(uniqueCoveredTargetsLength)
        .setGraphCoverageRatio(coverageRatio);
}


Path convertToPath(const std::unordered_set<EdgeId>& edgeIds)
{
    Path path;
    path.reserve(edgeIds.size());
    for(auto edgeId : edgeIds) {
        path.push_back({edgeId, true /*isUsefulAsTarget*/});
    }
    return path;
}

void buildStaticIndexOverPolygons(
    geolib3::StaticGeometrySearcher<geolib3::Polygon2, size_t>& searcher,
    const std::vector<geolib3::Polygon2>& polygons)
{

    for (size_t idx = 0; idx < polygons.size(); ++idx) {
        searcher.insert(&polygons[idx], idx);
    }

    searcher.build();
}

std::vector<geolib3::PreparedPolygon2>
makePreparedPolygons(const std::vector<geolib3::Polygon2>& polygons)
{
    std::vector<geolib3::PreparedPolygon2> result;
    result.reserve(polygons.size());

    for(const auto& polygon : polygons) {
        result.emplace_back(polygon);
    }
    return result;
}

std::vector<geolib3::Polygon2> convertToPolygons(const geolib3::MultiPolygon2& geom)
{
    std::vector<geolib3::Polygon2> result;
    result.reserve(geom.polygonsNumber());

    for(size_t idx = 0; idx < geom.polygonsNumber(); ++idx) {
        result.emplace_back(geom.polygonAt(idx));
    }
    return result;
}

Paths
distributeTargetsAmongPolygons(const RoadNetworkData& roadNetwork,
                               const std::unordered_set<EdgeId>& edgeIds,
                               const geolib3::MultiPolygon2& tasksGroupGeodeticGeom)
{
    ASSERT(tasksGroupGeodeticGeom.polygonsNumber() > 0);
    Paths result;

    if (tasksGroupGeodeticGeom.polygonsNumber() == 1) {
        result.push_back(convertToPath(edgeIds));
        return result;
    }

    result.resize(tasksGroupGeodeticGeom.polygonsNumber());
    auto polygonsVec = convertToPolygons(tasksGroupGeodeticGeom);
    auto preparedPolygons = makePreparedPolygons(polygonsVec);
    geolib3::StaticGeometrySearcher<geolib3::Polygon2, size_t> polygonsSearcher;
    buildStaticIndexOverPolygons(polygonsSearcher, polygonsVec);

    for (auto edgeId : edgeIds) {
        const geolib3::Polyline2& edgeGeom = roadNetwork.edge(edgeId).geom;
        auto searchResult = polygonsSearcher.find(edgeGeom.boundingBox());
        for (auto it = searchResult.first; it != searchResult.second; ++it)
        {
            size_t idx = it->value();

            if (polylineIsMostlyInsidePolygon(edgeGeom, preparedPolygons[idx],
                                              polygonsVec[idx].boundingBox()))
            {
                result[idx].push_back({edgeId, true /*isUsefulAsTarget*/});
                break;
            }
        }
    }
    return result;
}

PathWithData convertToPathWithData(const RoadNetworkData& roadNetwork,
                                   const Path& path)
{
    PathWithData result;
    result.reserve(path.size());
    for(const auto& pathEdge : path) {
        result.push_back({roadNetwork.edge(pathEdge.edgeId), pathEdge.isUsefulAsTarget});
    }
    return result;
}

} // namespace

void setFailedGenerationStatus(pgpool3::Pool& pool,
                               db::TId tasksGroupId,
                               const std::string& errorMsg)
{
    auto tasksGroup = db::ugc::TasksGroupGateway(*pool.masterReadOnlyTransaction())
                          .loadById(tasksGroupId);
    ERROR() << "Error in generating tasks for group " << tasksGroupId
            << ": " << errorMsg;
    if (tasksGroup.status() != db::ugc::TasksGroupStatus::Generating) {
        return;
    }
    tasksGroup.setStatus(db::ugc::TasksGroupStatus::Failed);
    db::ugc::TasksGroupGateway(*pool.masterWriteableTransaction()).update(tasksGroup);
}

bool generateTasks(pgpool3::Pool& pool,
                   object::Loader& objectLoader,
                   StaticGraphLoader& graphLoader,
                   db::TId tasksGroupId)
try {
    auto tasksGroup = db::ugc::TasksGroupGateway(*pool.masterReadOnlyTransaction())
                          .loadById(tasksGroupId);

    common::retryOnException<maps::Exception>(
        common::RetryPolicy()
            .setInitialTimeout(std::chrono::milliseconds(50))
            .setMaxAttempts(10)
            .setTimeoutBackoff(2),
        [&] {
            db::ugc::TasksGroupGateway(*pool.masterReadOnlyTransaction())
                .reload(tasksGroup);
            if (tasksGroup.status() != db::ugc::TasksGroupStatus::Generating) {
                throw maps::Exception{}
                    << "Tasks group object must be in Generating state "
                    << tasksGroupId;
            }
        });

    INFO() << "Loading road graph";
    auto roadNetworkData = loadRoadGraphFragmentForTasksGroup(graphLoader, tasksGroup);

    INFO() << "Loading exclusion regions";
    auto [residentialRegions, restrictedRegions] = loadResidentialAndRestrictedRegions(objectLoader, tasksGroup);

    INFO() << "There are " << restrictedRegions.size() << " restricted area";
    for (const auto& region : restrictedRegions) {
        auto edgeIds = roadNetworkData.getEdgesWithinPolygon(
            geolib3::convertMercatorToGeodetic(region.geom()),
            0 /* minFc */,
            10 /* maxFc */,
            {} /* prohibitedEdges */);
        INFO() << "Found " << edgeIds.size() << " intersected edges";
        for(EdgeId edgeId : edgeIds) {
            roadNetworkData.deleteEdge(edgeId);
        }
    }

    INFO() << "There are " << residentialRegions.size() << " residential area";
    std::unordered_set<EdgeId> unrecommendedEdges;
    for (const auto& region : residentialRegions) {
        auto edgeIds = roadNetworkData.getEdgesWithinPolygon(
            geolib3::convertMercatorToGeodetic(region.geom()),
            7 /* minFc */,
            10 /* maxFc */,
            {} /* prohibitedEdges */);
        INFO() << "Found " << edgeIds.size() << " intersected edges";
        for(EdgeId edgeId : edgeIds) {
            unrecommendedEdges.insert(edgeId);
        }
    }

    std::unordered_set<EdgeId> deadEndEdges;
    if (tasksGroup.excludeDeadends()) {
        deadEndEdges = getDeadEnds(roadNetworkData, tasksGroup.fcs());
    }

    INFO() << "Generating tasks";

    auto edgesFilter = [&](const Edge& edge) {
        return std::find(tasksGroup.fcs().begin(), tasksGroup.fcs().end(), edge.fc) != tasksGroup.fcs().end() &&
            unrecommendedEdges.count(edge.id) == 0 &&
            deadEndEdges.count(edge.id) == 0;
    };

    auto tasksGroupGeodeticGeom = geolib3::convertMercatorToGeodetic(tasksGroup.mercatorGeom());
    auto targetRoadGraphEdges =
        selectTargetEdges(roadNetworkData, tasksGroupGeodeticGeom, edgesFilter);

    graphLoader.removeActualizedEdges(tasksGroup, targetRoadGraphEdges);

    PathsWithData pathsWithData;
    if (tasksGroup.useRouting() == db::ugc::UseRouting::Yes) {
        auto paths = generateTasks(roadNetworkData, targetRoadGraphEdges,
                                   tasksGroup.recommendedTaskLengthMeters());

        std::vector<const Path*> pathPtrs;
        pathPtrs.reserve(paths.size());
        for(const auto& path : paths) {
            pathPtrs.push_back(&path);
        }

        checkResult(roadNetworkData, targetRoadGraphEdges, pathPtrs);

        pathsWithData.reserve(paths.size());
        for(const auto& path : paths) {
            pathsWithData.push_back(convertToPathWithData(roadNetworkData, path));
        }
    }
    else {
        auto paths = distributeTargetsAmongPolygons(roadNetworkData, targetRoadGraphEdges, tasksGroupGeodeticGeom);
        pathsWithData.reserve(paths.size());
        for (const auto& path : paths) {
            pathsWithData.push_back(optimizeTargertsForDrawing(path, roadNetworkData));
        }
    }

    INFO() << "Saving tasks";
    auto txn = pool.masterWriteableTransaction();
    saveTasks(*txn, tasksGroup, pathsWithData);
    calculateTasksGroupStats(roadNetworkData, targetRoadGraphEdges, pathsWithData, tasksGroup);
    tasksGroup.setStatus(db::ugc::TasksGroupStatus::Open);
    db::ugc::TasksGroupGateway(*txn).update(tasksGroup);

    INFO() << "Commiting to database";
    txn->commit();

    INFO() << "Done generating tasks for tasks group " << tasksGroupId;
    return true;
} catch (const maps::Exception& e) {
    setFailedGenerationStatus(pool, tasksGroupId, (std::stringstream() << e).str());
    return false;
} catch (const std::exception& e) {
    setFailedGenerationStatus(pool, tasksGroupId, e.what());
    return false;
} catch (...) {
    setFailedGenerationStatus(pool, tasksGroupId, "unknown error");
    return false;
}

} // namespace maps::mrc::gen_targets
