namespace StellaOps.ElkSharp;

internal static class ElkEdgeChannels
{
    internal static Dictionary<string, EdgeChannel> ComputeEdgeChannels(
        IReadOnlyCollection<ElkEdge> edges,
        IReadOnlyDictionary<string, ElkPositionedNode> positionedNodes,
        ElkLayoutDirection direction,
        IReadOnlyDictionary<string, LayerBoundary> layerBoundariesByNodeId)
    {
        var channels = new Dictionary<string, EdgeChannel>(edges.Count, StringComparer.Ordinal);
        var backwardEdges = new List<ElkEdge>();
        var forwardEdgesBySource = new Dictionary<string, List<ElkEdge>>(StringComparer.Ordinal);
        var outgoingCounts = edges
            .GroupBy(edge => edge.SourceNodeId, StringComparer.Ordinal)
            .ToDictionary(group => group.Key, group => group.Count(), StringComparer.Ordinal);

        foreach (var edge in edges)
        {
            var source = positionedNodes[edge.SourceNodeId];
            var target = positionedNodes[edge.TargetNodeId];
            var isBackward = direction == ElkLayoutDirection.LeftToRight
                ? (target.X + (target.Width / 2d)) < (source.X + (source.Width / 2d))
                : (target.Y + (target.Height / 2d)) < (source.Y + (source.Height / 2d));

            if (isBackward)
            {
                backwardEdges.Add(edge);
            }
            else
            {
                if (!forwardEdgesBySource.TryGetValue(edge.SourceNodeId, out var list))
                {
                    list = [];
                    forwardEdgesBySource[edge.SourceNodeId] = list;
                }

                list.Add(edge);
            }
        }

        var backwardGroups = backwardEdges
            .GroupBy(edge => edge.TargetNodeId, StringComparer.Ordinal)
            .SelectMany(targetGroup =>
            {
                var families = targetGroup
                    .GroupBy(edge => ElkEdgeChannelBands.ResolveLaneFamilyKey(edge.Label), StringComparer.Ordinal)
                    .OrderBy(group => ElkEdgeChannelBands.ResolveLaneFamilyPriority(group.First().Label))
                    .ThenBy(group => group.Key, StringComparer.Ordinal)
                    .ToArray();
                return families.Select((familyGroup, familyIndex) => new
                {
                    TargetNodeId = targetGroup.Key,
                    FamilyKey = familyGroup.Key,
                    FamilyIndex = familyIndex,
                    FamilyCount = families.Length,
                    Edges = familyGroup.ToArray(),
                    SharedOuterX = familyGroup.Max(edge =>
                    {
                        var s = positionedNodes[edge.SourceNodeId];
                        return s.X + s.Width;
                    }) + 56d,
                    Priority = ElkEdgeChannelBands.ResolveLaneFamilyPriority(familyGroup.First().Label),
                    Span = familyGroup.Max(edge =>
                    {
                        var s = positionedNodes[edge.SourceNodeId];
                        var t = positionedNodes[edge.TargetNodeId];
                        return direction == ElkLayoutDirection.LeftToRight
                            ? Math.Abs((s.X + (s.Width / 2d)) - (t.X + (t.Width / 2d)))
                            : Math.Abs((s.Y + (s.Height / 2d)) - (t.Y + (t.Height / 2d)));
                    }),
                });
            })
            .OrderByDescending(group => group.Span)
            .ThenBy(group => group.Priority)
            .ThenBy(group => group.TargetNodeId, StringComparer.Ordinal)
            .ToArray();

        var reservedHorizontalBands = new List<double>();
        for (var laneIndex = 0; laneIndex < backwardGroups.Length; laneIndex++)
        {
            var group = backwardGroups[laneIndex];
            var useSourceCollector = string.Equals(group.FamilyKey, "repeat", StringComparison.Ordinal);
            var preferredOuterY = double.NaN;
            if (direction == ElkLayoutDirection.LeftToRight && useSourceCollector)
            {
                var lowerCorridorY = ElkEdgeChannelCorridors.ResolveBackwardLowerCorridorY(group.Edges, positionedNodes);
                preferredOuterY = !double.IsNaN(lowerCorridorY)
                    ? lowerCorridorY
                    : ElkEdgeChannelCorridors.ResolveBackwardCorridorY(group.Edges, positionedNodes);
            }

            if (!double.IsNaN(preferredOuterY))
            {
                reservedHorizontalBands.Add(preferredOuterY);
            }

            foreach (var edge in group.Edges)
            {
                channels[edge.Id] = new EdgeChannel(
                    EdgeRouteMode.BackwardOuter,
                    laneIndex,
                    group.FamilyIndex,
                    group.FamilyCount,
                    0,
                    1,
                    0,
                    1,
                    -1,
                    0,
                    group.SharedOuterX,
                    preferredOuterY,
                    useSourceCollector,
                    double.NaN);
            }
        }

        var forwardEdgesByTarget = new Dictionary<string, List<ElkEdge>>(StringComparer.Ordinal);
        foreach (var sourceEdges in forwardEdgesBySource.Values)
        {
            foreach (var edge in sourceEdges)
            {
                if (!forwardEdgesByTarget.TryGetValue(edge.TargetNodeId, out var list))
                {
                    list = [];
                    forwardEdgesByTarget[edge.TargetNodeId] = list;
                }

                list.Add(edge);
            }
        }

        var sinkBandsByEdgeId = new Dictionary<string, (int BandIndex, int BandCount, double SharedOuterX, double PreferredOuterY)>(StringComparer.Ordinal);
        if (direction == ElkLayoutDirection.LeftToRight)
        {
            var reservedSinkBands = new List<double>(reservedHorizontalBands);
            foreach (var targetEdges in forwardEdgesByTarget)
            {
                var targetNode = positionedNodes[targetEdges.Key];
                var isSinkTarget = string.Equals(targetNode.Kind, "End", StringComparison.OrdinalIgnoreCase)
                    || !outgoingCounts.ContainsKey(targetEdges.Key);
                if (!isSinkTarget || targetEdges.Value.Count < 2)
                {
                    continue;
                }

                var sinkBands = targetEdges.Value
                    .GroupBy(edge => ElkEdgeChannelBands.ResolveLaneFamilyKey(edge.Label), StringComparer.Ordinal)
                    .OrderBy(group => ElkEdgeChannelBands.ResolveSinkLanePriority(group.First().Label))
                    .ThenBy(group => group.Key, StringComparer.Ordinal)
                    .ToArray();
                for (var bandIndex = 0; bandIndex < sinkBands.Length; bandIndex++)
                {
                    var sinkBandEdges = sinkBands[bandIndex].ToArray();
                    var sharedOuterX = sinkBands[bandIndex].Max(edge =>
                    {
                        var s = positionedNodes[edge.SourceNodeId];
                        return s.X + s.Width;
                    }) + 56d;
                    var familyKey = ElkEdgeChannelBands.ResolveLaneFamilyKey(sinkBands[bandIndex].First().Label);
                    var preferredOuterY = familyKey is "failure" or "timeout"
                        ? ElkEdgeChannelSinkCorridors.ResolveSinkCorridorY(sinkBandEdges, positionedNodes, reservedSinkBands)
                        : double.NaN;
                    if (!double.IsNaN(preferredOuterY))
                    {
                        reservedSinkBands.Add(preferredOuterY + (bandIndex * 24d));
                    }

                    foreach (var edge in sinkBandEdges)
                    {
                        sinkBandsByEdgeId[edge.Id] = (bandIndex, sinkBands.Length, sharedOuterX, preferredOuterY);
                    }
                }
            }
        }

        foreach (var sourceEdges in forwardEdgesBySource.Values)
        {
            var sorted = sourceEdges
                .OrderBy(e =>
                {
                    var t = positionedNodes[e.TargetNodeId];
                    return direction == ElkLayoutDirection.LeftToRight
                        ? t.Y + (t.Height / 2d)
                        : t.X + (t.Width / 2d);
                })
                .ToArray();

            for (var index = 0; index < sorted.Length; index++)
            {
                var targetEdges = forwardEdgesByTarget.GetValueOrDefault(sorted[index].TargetNodeId);
                var targetIncomingIndex = 0;
                var targetIncomingCount = 1;
                if (targetEdges is not null && targetEdges.Count > 1)
                {
                    var sortedBySourceY = targetEdges
                        .OrderBy(e =>
                        {
                            var s = positionedNodes[e.SourceNodeId];
                            return direction == ElkLayoutDirection.LeftToRight
                                ? s.Y + (s.Height / 2d)
                                : s.X + (s.Width / 2d);
                        })
                        .ToList();
                    targetIncomingIndex = sortedBySourceY.FindIndex(e => string.Equals(e.Id, sorted[index].Id, StringComparison.Ordinal));
                    targetIncomingCount = sortedBySourceY.Count;
                }

                var sinkBand = sinkBandsByEdgeId.GetValueOrDefault(sorted[index].Id, (-1, 0, 0d, double.NaN));
                var routeMode = EdgeRouteMode.Direct;
                if (sinkBandsByEdgeId.ContainsKey(sorted[index].Id))
                {
                    var sourceNode = positionedNodes[sorted[index].SourceNodeId];
                    var isGatewaySource = string.Equals(sourceNode.Kind, "Decision", StringComparison.OrdinalIgnoreCase)
                        || string.Equals(sourceNode.Kind, "Fork", StringComparison.OrdinalIgnoreCase);
                    if (isGatewaySource)
                    {
                        sinkBand = (-1, 0, 0d, double.NaN);
                    }
                    else
                    {
                        var familyKey = ElkEdgeChannelBands.ResolveLaneFamilyKey(sorted[index].Label);
                        if (familyKey is "failure" or "timeout")
                        {
                            routeMode = EdgeRouteMode.SinkOuterTop;
                        }
                        else
                        {
                            routeMode = EdgeRouteMode.SinkOuter;
                        }
                    }
                }

                channels[sorted[index].Id] = new EdgeChannel(
                    routeMode,
                    -1,
                    0,
                    1,
                    index,
                    sorted.Length,
                    targetIncomingIndex,
                    targetIncomingCount,
                    sinkBand.Item1,
                    sinkBand.Item2,
                    sinkBand.Item3,
                    sinkBand.Item4,
                    false,
                    double.NaN);
            }
        }

        ElkEdgeChannelBands.AllocateDirectForwardChannelBands(edges, positionedNodes, layerBoundariesByNodeId, channels, direction);

        return channels;
    }
}