namespace StellaOps.ElkSharp;

internal static partial class ElkCompoundLayout
{
    private static ElkNode[][] OptimizeLayerOrderingForHierarchy(
        ElkNode[][] initialLayers,
        IReadOnlyDictionary<string, List<string>> incomingNodeIds,
        IReadOnlyDictionary<string, List<string>> outgoingNodeIds,
        IReadOnlyDictionary<string, int> inputOrder,
        int iterationCount,
        ElkCompoundHierarchy hierarchy,
        IReadOnlySet<string> dummyNodeIds)
    {
        var layers = initialLayers
            .Select(layer => layer.ToList())
            .ToArray();
        var effectiveIterations = Math.Max(1, iterationCount);

        for (var iteration = 0; iteration < effectiveIterations; iteration++)
        {
            for (var layerIndex = 1; layerIndex < layers.Length; layerIndex++)
            {
                OrderLayerWithHierarchy(layers, layerIndex, incomingNodeIds, inputOrder, hierarchy, dummyNodeIds);
            }

            for (var layerIndex = layers.Length - 2; layerIndex >= 0; layerIndex--)
            {
                OrderLayerWithHierarchy(layers, layerIndex, outgoingNodeIds, inputOrder, hierarchy, dummyNodeIds);
            }
        }

        return layers
            .Select(layer => layer.ToArray())
            .ToArray();
    }

    private static void OrderLayerWithHierarchy(
        IReadOnlyList<List<ElkNode>> layers,
        int layerIndex,
        IReadOnlyDictionary<string, List<string>> adjacentNodeIds,
        IReadOnlyDictionary<string, int> inputOrder,
        ElkCompoundHierarchy hierarchy,
        IReadOnlySet<string> dummyNodeIds)
    {
        var currentLayer = layers[layerIndex];
        if (currentLayer.Count == 0)
        {
            return;
        }

        var positions = ElkNodeOrdering.BuildNodeOrderPositions(layers);
        var layerNodesById = currentLayer.ToDictionary(node => node.Id, StringComparer.Ordinal);
        var realNodeIdsInLayer = currentLayer
            .Where(node => !dummyNodeIds.Contains(node.Id))
            .Select(node => node.Id)
            .ToHashSet(StringComparer.Ordinal);

        var orderedRealIds = OrderNodesForSubtree(
            null,
            hierarchy,
            realNodeIdsInLayer,
            adjacentNodeIds,
            positions,
            inputOrder);
        var orderedDummies = currentLayer
            .Where(node => dummyNodeIds.Contains(node.Id))
            .OrderBy(node => ElkNodeOrdering.ResolveOrderingRank(node.Id, adjacentNodeIds, positions))
            .ThenBy(node => positions[node.Id])
            .ThenBy(node => inputOrder[node.Id])
            .ToArray();

        currentLayer.Clear();
        foreach (var nodeId in orderedRealIds)
        {
            currentLayer.Add(layerNodesById[nodeId]);
        }

        currentLayer.AddRange(orderedDummies);
    }

    private static List<string> OrderNodesForSubtree(
        string? parentNodeId,
        ElkCompoundHierarchy hierarchy,
        IReadOnlySet<string> realNodeIdsInLayer,
        IReadOnlyDictionary<string, List<string>> adjacentNodeIds,
        IReadOnlyDictionary<string, int> positions,
        IReadOnlyDictionary<string, int> inputOrder)
    {
        var blocks = new List<HierarchyOrderBlock>();
        foreach (var childNodeId in hierarchy.GetChildIds(parentNodeId))
        {
            if (realNodeIdsInLayer.Contains(childNodeId))
            {
                blocks.Add(BuildBlock([childNodeId]));
                continue;
            }

            if (!hierarchy.IsCompoundNode(childNodeId))
            {
                continue;
            }

            var descendantNodeIds = OrderNodesForSubtree(
                childNodeId,
                hierarchy,
                realNodeIdsInLayer,
                adjacentNodeIds,
                positions,
                inputOrder);
            if (descendantNodeIds.Count == 0)
            {
                continue;
            }

            blocks.Add(BuildBlock(descendantNodeIds));
        }

        return blocks
            .OrderBy(block => block.Rank)
            .ThenBy(block => block.MinCurrentPosition)
            .ThenBy(block => block.MinInputOrder)
            .SelectMany(block => block.NodeIds)
            .ToList();

        HierarchyOrderBlock BuildBlock(IReadOnlyList<string> nodeIds)
        {
            var ranks = new List<double>();
            foreach (var nodeId in nodeIds)
            {
                var nodeRank = ElkNodeOrdering.ResolveOrderingRank(nodeId, adjacentNodeIds, positions);
                if (!double.IsInfinity(nodeRank))
                {
                    ranks.Add(nodeRank);
                }
            }

            ranks.Sort();
            var rank = ranks.Count switch
            {
                0 => double.PositiveInfinity,
                _ when ranks.Count % 2 == 1 => ranks[ranks.Count / 2],
                _ => (ranks[(ranks.Count / 2) - 1] + ranks[ranks.Count / 2]) / 2d,
            };
            return new HierarchyOrderBlock(
                nodeIds,
                rank,
                nodeIds.Min(nodeId => positions.GetValueOrDefault(nodeId, int.MaxValue)),
                nodeIds.Min(nodeId => inputOrder.GetValueOrDefault(nodeId, int.MaxValue)));
        }
    }

    private readonly record struct HierarchyOrderBlock(
        IReadOnlyList<string> NodeIds,
        double Rank,
        int MinCurrentPosition,
        int MinInputOrder);
}