namespace ReactorMaintenance.Simulation;

internal static class NetworkPropagationSystem
{
    public static LevelState Propagate(LevelState level)
    {
        var fuel = ClearTransient(level.Fuel);
        var coolant = ClearTransient(level.Coolant);
        var electricity = ClearTransient(level.Electricity);
        var next = level.WithRuntimeArrays(fuel, coolant, electricity, level.Surface.ToArray(), level.Props.ToArray());

        foreach (var carrier in Enum.GetValues<ECarrierType>())
            next = PropagateCarrier(next, carrier);

        return next;
    }

    private static UndergroundCell[] ClearTransient(IReadOnlyList<UndergroundCell> layer)
    {
        return layer.Select(cell => cell with { Amount = 0, Intensity = 0 }).ToArray();
    }

    private static LevelState PropagateCarrier(LevelState level, ECarrierType carrier)
    {
        var layer = level.Layer(carrier).ToArray();
        var sources = LevelTraversal.AllPositions(level).Where(position => level.GetProp(position) is { Type: EPropType.Flow, SwitchState: EPropSwitchState.Enabled, Carrier: var sourceCarrier } && sourceCarrier == carrier && level.GetUnderground(position, carrier).CarriesFlow).ToArray();
        var junctions = JunctionFlowAnalyzer.Analyze(level).Where(junction => junction.IsValid && junction.Carrier == carrier).ToDictionary(junction => junction.Position);

        foreach (var source in sources)
            ApplySourceFlow(level, layer, source, carrier, junctions);

        return carrier switch {
            ECarrierType.Fuel => level with { Fuel = layer },
            ECarrierType.Coolant => level with { Coolant = layer },
            ECarrierType.Electricity => level with { Electricity = layer },
            _ => throw new ArgumentOutOfRangeException(nameof(carrier), carrier, "Unsupported carrier.")
        };
    }

    private static void ApplySourceFlow(LevelState level, UndergroundCell[] layer, GridPosition source, ECarrierType carrier, IReadOnlyDictionary<GridPosition, JunctionFlow> junctions)
    {
        var open = new Queue<(GridPosition Position, int Distance, float AmountFactor, float IntensityFactor)>();
        var best = new Dictionary<GridPosition, float>();
        open.Enqueue((source, 0, 1, 1));
        best[source] = 1;

        while (open.Count > 0)
        {
            var current = open.Dequeue();
            var amount = Balancing.Current.ClampValue((Balancing.Current.SourceAmount * current.AmountFactor) - (current.Distance * Balancing.Current.DistanceAmountFalloff));
            var intensity = Balancing.Current.ClampValue((Balancing.Current.SourceIntensity * current.IntensityFactor) - (current.Distance * Balancing.Current.DistanceIntensityFalloff));
            var index = level.Index(current.Position);
            layer[index] = layer[index] with {
                Amount = Math.Max(layer[index].Amount, amount),
                Intensity = Math.Max(layer[index].Intensity, intensity)
            };

            foreach (var next in current.Position.Neighbors().Where(level.InBounds))
            {
                if (!level.GetUnderground(next, carrier).CarriesFlow)
                    continue;

                var weights = BranchWeights(current.Position, next, junctions);
                var amountFactor = current.AmountFactor * weights.Amount;
                var intensityFactor = current.IntensityFactor * weights.Intensity;
                if (amountFactor <= 0 || intensityFactor <= 0)
                    continue;

                if (best.TryGetValue(next, out var oldBest) && oldBest >= amountFactor)
                    continue;

                best[next] = amountFactor;
                open.Enqueue((next, current.Distance + 1, amountFactor, intensityFactor));
            }
        }
    }

    private static (float Amount, float Intensity) BranchWeights(GridPosition from, GridPosition to, IReadOnlyDictionary<GridPosition, JunctionFlow> junctions)
    {
        if (!junctions.TryGetValue(from, out var junction))
            return (1, 1);

        var weight = junction.WeightFor(to);
        return (weight, weight);
    }
}