#include "graph_coverage_explorer.h"

#include <maps/wikimap/mapspro/services/mrc/browser/lib/tools.h>
#include <maps/wikimap/mapspro/services/mrc/libs/fb/include/common.h>
#include <maps/libs/geolib/include/closest_point.h>
#include <maps/libs/geolib/include/distance.h>

#include <set>

namespace maps::mrc::browser {
namespace {

const Meters COVERAGE_EXPLORER_SNAP_DISTANCE{export_gen::COVERAGE_METERS_SNAP_THRESHOLD};

constexpr geolib3::Radians COVERAGE_EXPLORER_SNAP_ANGLE{geolib3::PI / 12};

bool areConnected(
    const common::geometry::PolylinePosition& lhs,
    const common::geometry::PolylinePosition& rhs)
{
    return lhs == rhs;
}

bool areConnected(
    const geolib3::Point2& lhs,
    const geolib3::Point2& rhs,
    Meters maxDistance)
{
    return geolib3::fastGeoDistance(lhs, rhs) <= maxDistance.value();
}

bool areConnected(const GraphCoverageExplorer::Node& lhs,
                  const GraphCoverageExplorer::Node& rhs,
                  Meters maxDistance)
{
    return areConnected(lhs.ray.pos, rhs.ray.pos, maxDistance);
}

double relativePosition(
    const geolib3::Segment2& segment, const geolib3::Point2& point)
{
    return std::sqrt(
        geolib3::squaredLength(point - segment.start()) /
        geolib3::squaredLength(segment.vector()));
}

// To convert PolylinePosition back and forth between common::geometry:: and fb:: ones
template<class To, class From>
To convertTo(const From& from)
{
    using IdxT = decltype(To{0, 0}.segmentIdx());
    using PosT = decltype(To{0, 0}.segmentRelPosition());
    return {
        static_cast<IdxT>(from.segmentIdx()),
        static_cast<PosT>(from.segmentRelPosition())};
}

template<>
common::geometry::SubPolyline convertTo<
    common::geometry::SubPolyline,
    fb::CoveredSubPolyline>(const fb::CoveredSubPolyline& from)
{
    return {
        convertTo<common::geometry::PolylinePosition>(from.begin()),
        convertTo<common::geometry::PolylinePosition>(from.end())};
}

std::optional<fb::CoveredSubPolyline> findNearestBackwardCoverage(
    const std::vector<fb::CoveredSubPolyline>& coveredSubpolylines,
    const common::geometry::PolylinePosition& position)
{
    const auto it = std::lower_bound(
        coveredSubpolylines.crbegin(),
        coveredSubpolylines.crend(),
        position,
        [](const auto& coveredSubpolyline, const auto& position) {
            return coveredSubpolyline.begin() >
                   convertTo<fb::PolylinePosition>(position);
        });
    if (it == coveredSubpolylines.crend()) {
        return std::nullopt;
    }
    return *it;
}

std::optional<fb::CoveredSubPolyline> findNearestForwardCoverage(
    const std::vector<fb::CoveredSubPolyline>& coveredSubpolylines,
    const common::geometry::PolylinePosition& position)
{
    const auto it = std::lower_bound(
        coveredSubpolylines.cbegin(),
        coveredSubpolylines.cend(),
        position,
        [](const auto& coveredSubpolyline, const auto& position) {
            return coveredSubpolyline.begin() <
                   convertTo<fb::PolylinePosition>(position);
        });
    if (it == coveredSubpolylines.cend()) {
        return std::nullopt;
    }
    return *it;
}

using EdgeIdSet = std::set<road_graph::EdgeId>;
using EdgeIdVector = std::vector<road_graph::EdgeId>;

bool isReverseEdge(
    const road_graph::Graph& graph,
    road_graph::EdgeId edgeId,
    road_graph::EdgeId maybeReverseEdgeId)
{
    if (const auto reverseEdgeId = graph.reverse(edgeId)) {
        return graph.base(*reverseEdgeId) == graph.base(maybeReverseEdgeId);
    }
    return false;
}

EdgeIdVector getInEdgeIds(const road_graph::Graph& graph,
                          road_graph::EdgeId edgeId)
{
    EdgeIdSet result;
    const auto inEdgeIds = graph.inEdgeIds(graph.edge(edgeId).source);

    for (auto inEdgeId: inEdgeIds) {
        if (!isReverseEdge(graph, edgeId, inEdgeId)) {
            result.insert(inEdgeId);
        }
    }

    return {result.begin(), result.end()};
}

EdgeIdVector getOutEdgeIds(const road_graph::Graph& graph,
                           road_graph::EdgeId edgeId)
{
    EdgeIdSet result;
    const auto outEdgeIds = graph.outEdgeIds(graph.edge(edgeId).target);

    for (auto outEdgeId: outEdgeIds) {
        if (!isReverseEdge(graph, edgeId, outEdgeId)) {
            result.insert(outEdgeId);
        }
    }

    return {result.begin(), result.end()};
}

} // namespace

struct GraphCoverageExplorer::SnappedPoint {
    road_graph::EdgeId edgeId;
    common::geometry::PolylinePosition position;
    Meters distance;
};

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::find(const geolib3::Point2& geoPoint,
                            int maxFc,
                            Meters searchRadius) const
{
    const auto snappedPoint = snapPointToCoveredEdge(geoPoint, maxFc, searchRadius);

    if (!snappedPoint) {
        return std::nullopt;
    }

    return findNearestNodeOnEdge(
        geoPoint, snappedPoint->edgeId, snappedPoint->position);
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::find(const db::Ray& ray, Meters searchRadius) const
{
    const auto snappedPoint = snapRayToCoveredEdge(ray, searchRadius);

    if (!snappedPoint) {
        return std::nullopt;
    }

    return findNearestNodeOnEdge(
        ray.pos, snappedPoint->edgeId, snappedPoint->position);
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::find(const db::Feature& feature) const
{
    for (auto edgeId: getCoveredEdgeIds(feature.geodeticPos())) {
        const auto coveredSubpolylines = getCoveredSubpolylines(edgeId);
        const auto it = std::find_if(
            coveredSubpolylines.cbegin(),
            coveredSubpolylines.cend(),
            [featureId = feature.id()](const auto& coveredSubpolyline) {
                return coveredSubpolyline.featureId() == featureId;
            });

        if (it != coveredSubpolylines.cend()) {
            return resolveNode(toNode(edgeId, *it));
        }
    }

    return std::nullopt;
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::uturn(const GraphCoverageExplorer::Node& node) const
{
    OptNode reverseNode = find(reverseRay(node.ray));
    const auto reverseEdgeId = reverse(node.edgeId);

    // If there is a reversed edge in the topology and reverse find didnt find
    // a node on it then uturn is considered absent.
    if (reverseEdgeId && reverseNode &&
        reverseNode->edgeId != graph_.graph().base(*reverseEdgeId)) {
        return std::nullopt;
    }
    // FIXME: it could be possible that there is no coverage on a small
    //        reverse edge but it exist on its continuation. But for now we ignore
    //        this fact for the sake of simpllicity.
    return reverseNode;
}

GraphCoverageExplorer::Nodes
GraphCoverageExplorer::backward(const GraphCoverageExplorer::Node& node,
                                Meters maxDistance) const
{
    OptNode backwardEdgeNode = nextEdgeBackwardNode(node);

    if (backwardEdgeNode && isResolved(*backwardEdgeNode)) {
        // If there is a backward node which starts (resolved to) on the same
        // edge then there is a single backward connection.
        return areConnected(node, *backwardEdgeNode, maxDistance)
                   ? Nodes{*backwardEdgeNode} : Nodes{};
    }

    // Check incoming edges for backward nodes
    EdgeIdVector backwardEdgeIds;

    const auto populateBackwardEdgeIds = [&](road_graph::EdgeId edgeId) {
        const auto edgeIds = getInEdgeIds(graph_.graph(), edgeId);
        backwardEdgeIds.insert(backwardEdgeIds.end(), edgeIds.begin(), edgeIds.end());
    };

    std::set<road_graph::EdgeId> visited{node.edgeId};
    std::vector<Node> connections;

    populateBackwardEdgeIds(node.edgeId);

    while (!backwardEdgeIds.empty()) {
        const auto backwardEdgeId = backwardEdgeIds.back();
        backwardEdgeIds.pop_back();

        if (!visited.insert(backwardEdgeId).second) {
            continue;
        }

        OptNode backwardNode;

        const auto coveredSupbolylines = getCoveredSubpolylines(backwardEdgeId);
        if (!coveredSupbolylines.empty()) {
            backwardNode = toNode(backwardEdgeId, coveredSupbolylines.back());
        }

        if (backwardNode && isResolved(*backwardNode)) {
            if (areConnected(node, *backwardNode, maxDistance)) {
                connections.push_back(*backwardNode);
            }
            continue;
        }

        const geolib3::Point2 backwardStartPoint = pointAtEdgeStart(backwardEdgeId);
        if (areConnected(node.ray.pos, backwardStartPoint, maxDistance)) {
            populateBackwardEdgeIds(backwardEdgeId);
        }
    }

    return connections;
}

GraphCoverageExplorer::Nodes
GraphCoverageExplorer::forward(const GraphCoverageExplorer::Node& node,
                               Meters maxDistance) const
{
    OptNode forwardEdgeNode = nextEdgeForwardNode(node);

    if (forwardEdgeNode) {
        // If there is a forward node which on the same edge then there is a
        // single forward connection
        return areConnected(node, *forwardEdgeNode, maxDistance)
                   ? Nodes{*forwardEdgeNode}
                   : Nodes{};
    }

    // Check outgoing edges for forward nodes
    EdgeIdVector forwardEdgeIds;

    const auto populateForwardEdgeIds = [&](road_graph::EdgeId edgeId) {
        const auto edgeIds = getOutEdgeIds(graph_.graph(), edgeId);
        forwardEdgeIds.insert(forwardEdgeIds.end(), edgeIds.begin(), edgeIds.end());
    };

    std::set<road_graph::EdgeId> visited{node.edgeId};
    std::vector<Node> connections;

    populateForwardEdgeIds(node.edgeId);

    while (!forwardEdgeIds.empty()) {
        const auto forwardEdgeId = forwardEdgeIds.back();
        forwardEdgeIds.pop_back();

        if (!visited.insert(forwardEdgeId).second) {
            continue;
        }

        OptNode forwardNode;

        // Coverage subpolylines in the beginning of an edge are most
        // likely just continuation of nodes from its incoming edges. So
        // skip them.
        for (const auto& coverage: getCoveredSubpolylines(forwardEdgeId)) {
            Node nextNode = toNode(forwardEdgeId, coverage);
            if (isResolved(nextNode)) {
                forwardNode = std::move(nextNode);
                break;
            }
        }

        if (forwardNode) {
            if (areConnected(node, *forwardNode, maxDistance)) {
                connections.push_back(*forwardNode);
            }
            continue;
        }

        const geolib3::Point2 forwardEndPoint = pointAtEdgeEnd(forwardEdgeId);
        if (areConnected(node.ray.pos, forwardEndPoint, maxDistance)) {
            populateForwardEdgeIds(forwardEdgeId);
        }
    }

    return connections;
}

bool GraphCoverageExplorer::isResolved(const Node& node) const {
    return resolveNode(node).edgeId == node.edgeId;
}

geolib3::Segment2 GraphCoverageExplorer::getSegment(
    road_graph::EdgeId edgeId, uint16_t segmentIdx) const
{
    return graph_.graph().edgeData(edgeId).geometry().segmentAt(segmentIdx);
}

geolib3::Point2 GraphCoverageExplorer::pointAtEdgeEnd(road_graph::EdgeId edgeId) const
{
    const auto edgeEnd = completeEdgeSubpolyline(edgeId).end();
    const auto segment = getSegment(edgeId, edgeEnd.segmentIdx());
    return segment.pointByPosition(edgeEnd.segmentRelPosition());
}

geolib3::Point2 GraphCoverageExplorer::pointAtEdgeStart(road_graph::EdgeId edgeId) const
{
    const auto segment = getSegment(edgeId, 0u);
    return segment.pointByPosition(0.0);
}

GraphCoverageExplorer::Node
GraphCoverageExplorer::toNode(road_graph::EdgeId edgeId,
                              const fb::CoveredSubPolyline& from) const
{
    const auto segment = getSegment(edgeId, from.begin().segmentIdx());

    const db::Ray ray{
        segment.pointByPosition(from.begin().segmentRelPosition()),
        geolib3::fastGeoDirection(segment)};

    return
    {
        .edgeId = graph_.graph().base(edgeId),
        .coverage = convertTo<common::geometry::SubPolyline>(from),
        .featureId = from.featureId(),
        .ray = ray
    };
}

bool GraphCoverageExplorer::edgeHasCoverage(road_graph::EdgeId edgeId) const
{
    const auto coveredEdge =
        graph_.getCoveredEdgeById(graph_.graph().base(edgeId));

    if (!coveredEdge) {
        return false;
    }

    return coverageSelector_(*coveredEdge) != nullptr;
};

std::optional<GraphCoverageExplorer::SnappedPoint>
GraphCoverageExplorer::snapPointToEdge(road_graph::EdgeId edgeId,
                                       const geolib3::Point2& point,
                                       const std::optional<Meters>& maxDistance) const
{
    const geolib3::Polyline2 edgePolyline{
        graph_.graph().edgeData(edgeId).geometry()};

    const auto closestPoint = geolib3::closestPoint(edgePolyline, point);
    const Meters distance{geolib3::fastGeoDistance(point, closestPoint)};
    if(maxDistance && *maxDistance < distance) {
        return std::nullopt;
    }

    const auto closestSegmentIdx = edgePolyline.segmentIndex(closestPoint);
    ASSERT(closestSegmentIdx);

    const auto segment = edgePolyline.segmentAt(*closestSegmentIdx);

    return SnappedPoint{
        .edgeId = edgeId,
        .position =
            common::geometry::PolylinePosition{
                *closestSegmentIdx,
                relativePosition(segment, closestPoint)},
        .distance = distance};
}

std::optional<road_graph::EdgeId>
GraphCoverageExplorer::reverseForTrafficDirection(const geolib3::Point2& point,
                                                  const SnappedPoint& snappedPoint) const
{
    const auto reverseEdgeId = reverse(snappedPoint.edgeId);
    if (!reverseEdgeId) {
        return std::nullopt;
    }

    if (!graph_.getCoveredEdgeById(*reverseEdgeId)) {
        return std::nullopt;
    }

    // If there is a covered reverse edge then compute on which side of an
    // original edge the point is on and pick either the original edge or
    // its reverse edge depending on local traffic side.
    const auto edgeData = graph_.graph().edgeData(snappedPoint.edgeId);
    const auto trafficSide = edgeData.trafficSide();
    const geolib3::Polyline2 edgePolyline{edgeData.geometry()};

    const auto segment =
        edgePolyline.segmentAt(snappedPoint.position.segmentIdx());

    auto ccwRadians = orientedAngleBetween(
        geolib3::fastGeoDirection(segment),
        geolib3::fastGeoDirection(geolib3::Segment2{segment.start(), point}));

    if (ccwRadians.value() < 0.0) {
        ccwRadians = PI * 2 + ccwRadians;
    }

    const bool snapToReverseEdge =
        (trafficSide == road_graph::EdgeTrafficSide::Right &&
         ccwRadians < PI) ||
        (trafficSide == road_graph::EdgeTrafficSide::Left && PI < ccwRadians);

    if (snapToReverseEdge) {
        return *reverseEdgeId;
    }

    return std::nullopt;
}

std::vector<road_graph::EdgeId>
GraphCoverageExplorer::getCoveredEdgeIds(
    const geolib3::Point2& point, int maxFc, Meters bboxSize) const
{
    std::vector<road_graph::EdgeId> edgeIds;

    for (auto edgeId: graph_.getCoveredEdgeIdsByBbox(
             db::expandBbox(point.boundingBox(), bboxSize.value()), maxFc)) {

        if (edgeHasCoverage(edgeId)) {
            edgeIds.push_back(edgeId);
        }
    }
    return edgeIds;
}

std::optional<GraphCoverageExplorer::SnappedPoint>
GraphCoverageExplorer::snapPointToCoveredEdge(const geolib3::Point2& point,
                                              int maxFc,
                                              Meters searchRadius) const
{
    std::optional<SnappedPoint> snappedPoint;

    for (auto edgeId: getCoveredEdgeIds(point, maxFc, searchRadius)) {
        if (const auto candidate =
                snapPointToEdge(edgeId, point, searchRadius)) {

            if (!snappedPoint ||
                candidate->distance < snappedPoint->distance) {
                snappedPoint = candidate;
            }
        }
    }

    if (!snappedPoint) {
        return std::nullopt;
    }

    if (const auto reverseEdgeId =
            reverseForTrafficDirection(point, *snappedPoint)) {
        return snapPointToEdge(*reverseEdgeId, point, searchRadius);
    }

    return snappedPoint;
}

std::optional<GraphCoverageExplorer::SnappedPoint>
GraphCoverageExplorer::snapRayToCoveredEdge(db::Ray ray,
                                            Meters searchRadius) const
{
    db::Ray snappedRay = ray;
    const auto edgeId = snapToGraph(
        graph_,
        snappedRay,
        makeSnapLossFunc(
            COVERAGE_EXPLORER_SNAP_DISTANCE, COVERAGE_EXPLORER_SNAP_ANGLE));

    if (!edgeId) {
        return std::nullopt;
    }

    return snapPointToEdge(*edgeId, snappedRay.pos, searchRadius);
}

std::vector<fb::CoveredSubPolyline>
GraphCoverageExplorer::getCoveredSubpolylines(road_graph::EdgeId edgeId) const
{
    auto coveredEdge = graph_.getCoveredEdgeById(edgeId);
    if (!coveredEdge) {
        return {};
    }

    if (const fb::TEdgeCoverage* coverage = coverageSelector_(*coveredEdge)) {
        return coverage->coveredSubpolylines;
    }

    return {};
}

GraphCoverageExplorer::Node
GraphCoverageExplorer::resolveNode(const Node& node) const
{
    if (!isBackwardTipNode(node)) {
        // This node starts from somewhere in the middle of an edge so there
        // is nothing to resolve.
        return node;
    }

    OptNode resolved;
    const auto resolveTo = [&](const Node& backwardNode) {
        if (!resolved ||
            geolib3::fastGeoDistance(resolved->ray.pos, node.ray.pos) <
            geolib3::fastGeoDistance(resolved->ray.pos, backwardNode.ray.pos)) {
            resolved = backwardNode;
        }
    };

    for (auto backwardEdgeId: getInEdgeIds(graph_.graph(), node.edgeId)) {
        const auto coveredSupbolylines = getCoveredSubpolylines(backwardEdgeId);
        if (coveredSupbolylines.empty()) {
            continue;
        }

        const auto backwardNode = toNode(backwardEdgeId, coveredSupbolylines.back());

        if (isForwardTipNode(backwardNode) && node.featureId == backwardNode.featureId) {
            resolveTo(resolveNode(backwardNode));
        }
    }

    if (resolved) {
        return *resolved;
    }
    return node;
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::findNearestNodeOnEdge(
    const geolib3::Point2& position,
    const road_graph::EdgeId edgeId,
    const common::geometry::PolylinePosition& snappedPosition) const
{
    const auto coveredSubPolylines = getCoveredSubpolylines(edgeId);

    const auto backwardCov =
        findNearestBackwardCoverage(coveredSubPolylines, snappedPosition);
    const auto forwardCov =
        findNearestForwardCoverage(coveredSubPolylines, snappedPosition);

    if (!backwardCov && !forwardCov) {
        return std::nullopt;

    } else if (!backwardCov && forwardCov) {
        // No need to resolve forward coverage on the same edge because it
        // starts from it.
        return toNode(edgeId, *forwardCov);

    } else if (!forwardCov && backwardCov) {
        return resolveNode(toNode(edgeId, *backwardCov));
    }

    const auto backwardNode = resolveNode(toNode(edgeId, *backwardCov));
    // No need to resolve forward coverage on the same edge because it
    // starts from it.
    const auto forwardNode = toNode(edgeId, *forwardCov);

    if (fastGeoDistance(position, backwardNode.ray.pos) <
        fastGeoDistance(position, forwardNode.ray.pos)) {
        return backwardNode;

    } else {
        return forwardNode;
    }
}

common::geometry::SubPolyline
GraphCoverageExplorer::completeEdgeSubpolyline(road_graph::EdgeId edgeId) const
{
    const std::size_t segmentsNum = graph_.getEdgeSegmentsNum(edgeId);

    return {
        common::geometry::PolylinePosition{0, 0.0},
        common::geometry::PolylinePosition{segmentsNum - 1, 1.0}};
}

std::optional<road_graph::EdgeId> GraphCoverageExplorer::reverse(
    road_graph::EdgeId edgeId) const
{
    edgeId = graph_.graph().base(edgeId);
    if (const auto reversedEdgeId = graph_.graph().reverse(edgeId)) {
        return graph_.graph().base(*reversedEdgeId);
    }
    return std::nullopt;
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::nextEdgeBackwardNode(const Node& node) const {
    const auto coveredSupbolylines = getCoveredSubpolylines(node.edgeId);
    auto it = std::find_if(
        coveredSupbolylines.crbegin(),
        coveredSupbolylines.crend(),
        [&](const auto& coveredSubpolyline) {
            return coveredSubpolyline.featureId() == node.featureId;
        });
    ASSERT(it != coveredSupbolylines.crend());
    const auto backwardIt = std::next(it);

    if (backwardIt == coveredSupbolylines.crend()) {
        return std::nullopt;
    }

    return toNode(node.edgeId, *backwardIt);
}

GraphCoverageExplorer::OptNode
GraphCoverageExplorer::nextEdgeForwardNode(const Node& node) const {
    const auto coveredSupbolylines = getCoveredSubpolylines(node.edgeId);
    const auto it = std::find_if(
        coveredSupbolylines.cbegin(),
        coveredSupbolylines.cend(),
        [&](const auto& coveredSubpolyline) {
            return coveredSubpolyline.featureId() == node.featureId;
        });
    ASSERT(it != coveredSupbolylines.cend());
    const auto forwardIt = std::next(it);

    if (forwardIt == coveredSupbolylines.cend()) {
        return std::nullopt;
    }

    return toNode(node.edgeId, *forwardIt);
}

bool GraphCoverageExplorer::isBackwardTipNode(const Node& node) const
{
    return areConnected(
        node.coverage.begin(),
        completeEdgeSubpolyline(node.edgeId).begin());
}

bool GraphCoverageExplorer::isForwardTipNode(const Node& node) const
{
    return areConnected(
        node.coverage.end(),
        completeEdgeSubpolyline(node.edgeId).end());
}

} // namespace maps::mrc::browser
