using System.Collections.Generic;
using NUnit.Framework;
using ProjectM.Simulation;
using Unity.Collections;
using Unity.Mathematics;

namespace ProjectM.Tests
{
    /// <summary>
    /// Pure, world-free tests for <see cref="ConveyorMath"/> — the DETERMINISTIC, ORDER-INDEPENDENT belt resolver
    /// that the server <c>ConveyorTransportSystem</c> applies. Drives the array-based <c>ResolveMoves</c> directly
    /// (no ECS world) so determinism is provable: a 4-cell line advances exactly one cell/tick, a Y-junction lets a
    /// single deterministic winner (lowest CellKey source) into a shared destination while the loser STALLS with no
    /// loss, a blocked (already-occupied) destination stalls its source, and the END-STATE is IDENTICAL under two
    /// shuffled input orders. Sink (machine-input) cells always accept. <c>DirOffset</c>/<c>CellKey</c> are pinned too.
    /// </summary>
    public class ConveyorMathTests
    {
        const byte DirPosX = 0;
        const byte DirNegX = 1;
        const byte DirPosZ = 2;
        const byte DirNegZ = 3;

        [Test]
        public void DirOffset_Maps_All_Four_Directions()
        {
            Assert.AreEqual(new int2(1, 0), ConveyorMath.DirOffset(DirPosX), "0 = +X");
            Assert.AreEqual(new int2(-1, 0), ConveyorMath.DirOffset(DirNegX), "1 = -X");
            Assert.AreEqual(new int2(0, 1), ConveyorMath.DirOffset(DirPosZ), "2 = +Z");
            Assert.AreEqual(new int2(0, -1), ConveyorMath.DirOffset(DirNegZ), "3 = -Z");
        }

        [Test]
        public void CellKey_Is_Stable_And_Order_Defining()
        {
            // The key must be a stable total order used to break move ties deterministically.
            long a = ConveyorMath.CellKey(new int2(0, 0));
            long b = ConveyorMath.CellKey(new int2(1, 0));
            long c = ConveyorMath.CellKey(new int2(0, 1));
            Assert.AreEqual(a, ConveyorMath.CellKey(new int2(0, 0)), "CellKey is deterministic for a given cell.");
            Assert.AreNotEqual(a, b);
            Assert.AreNotEqual(a, c);
            Assert.AreNotEqual(b, c, "Distinct cells map to distinct keys (no collision in-range).");
        }

        // ---- ResolveMoves harness ------------------------------------------------------------------------------

        /// <summary>One conveyor cell in a scenario: its grid cell, belt direction, and the item it currently holds.</summary>
        struct Belt
        {
            public int2 Cell;
            public byte Dir;
            public int Res;   // 0 = empty
            public int Cnt;   // 0 = empty
        }

        /// <summary>
        /// Drives <see cref="ConveyorMath.ResolveMoves"/> over a set of belts (+ optional sink cells), returning the
        /// resolved (srcIndex -&gt; destCell) move list as a dictionary keyed by source index for easy assertion.
        /// All native containers are Temp + disposed before returning.
        /// </summary>
        static Dictionary<int, int2> Resolve(Belt[] belts, int2[] sinks)
        {
            int n = belts.Length;
            var srcCells = new NativeArray<int2>(n, Allocator.Temp);
            var dirs = new NativeArray<byte>(n, Allocator.Temp);
            var itemRes = new NativeArray<int>(n, Allocator.Temp);
            var itemCnt = new NativeArray<int>(n, Allocator.Temp);
            var cellToIndex = new NativeHashMap<int2, int>(n, Allocator.Temp);
            var sinkCells = new NativeHashSet<int2>(math.max(1, sinks?.Length ?? 0), Allocator.Temp);
            var outMoveDst = new NativeArray<int2>(n, Allocator.Temp);
            var outMoveSrcIdx = new NativeArray<int>(n, Allocator.Temp);

            for (int i = 0; i < n; i++)
            {
                srcCells[i] = belts[i].Cell;
                dirs[i] = belts[i].Dir;
                itemRes[i] = belts[i].Res;
                itemCnt[i] = belts[i].Cnt;
                cellToIndex[belts[i].Cell] = i;
            }
            if (sinks != null)
                foreach (var s in sinks)
                    sinkCells.Add(s);

            ConveyorMath.ResolveMoves(srcCells, dirs, itemRes, itemCnt, cellToIndex, sinkCells,
                outMoveDst, outMoveSrcIdx, out int moveCount);

            var result = new Dictionary<int, int2>();
            for (int m = 0; m < moveCount; m++)
                result[outMoveSrcIdx[m]] = outMoveDst[m];

            srcCells.Dispose(); dirs.Dispose(); itemRes.Dispose(); itemCnt.Dispose();
            cellToIndex.Dispose(); sinkCells.Dispose(); outMoveDst.Dispose(); outMoveSrcIdx.Dispose();
            return result;
        }

        [Test]
        public void Line_Item_Advances_Exactly_One_Cell()
        {
            // 4-cell +X line at (0,0),(1,0),(2,0),(3,0). Only the head (index 0) holds an item; the cell ahead (1,0)
            // is EMPTY in the snapshot, so the item moves exactly one cell. No other belt holds an item -> no other move.
            var belts = new[]
            {
                new Belt { Cell = new int2(0, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = new int2(1, 0), Dir = DirPosX, Res = 0, Cnt = 0 },
                new Belt { Cell = new int2(2, 0), Dir = DirPosX, Res = 0, Cnt = 0 },
                new Belt { Cell = new int2(3, 0), Dir = DirPosX, Res = 0, Cnt = 0 },
            };

            var moves = Resolve(belts, sinks: null);

            Assert.AreEqual(1, moves.Count, "Exactly one item moves this tick.");
            Assert.IsTrue(moves.ContainsKey(0), "The head belt's item is the one that moves.");
            Assert.AreEqual(new int2(1, 0), moves[0], "The item advances exactly one cell along +X.");
        }

        [Test]
        public void Full_Line_All_Advance_One_Cell_From_PreMove_Snapshot()
        {
            // Every belt in a 4-cell line holds an item, EXCEPT the head cell ahead of the line is open. Because
            // occupancy is read from the PRE-MOVE snapshot (double-buffered), each item shifts forward one cell —
            // the cell ahead was occupied in the snapshot for the tail belts, so ONLY the lead item (whose target is
            // empty in the snapshot) may advance. This pins the snapshot (not live) occupancy rule.
            var belts = new[]
            {
                new Belt { Cell = new int2(0, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },     // -> (1,0) occupied in snapshot -> stall
                new Belt { Cell = new int2(1, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },     // -> (2,0) occupied in snapshot -> stall
                new Belt { Cell = new int2(2, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },     // -> (3,0) occupied in snapshot -> stall
                new Belt { Cell = new int2(3, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },     // -> (4,0) not a belt, not a sink -> no move
            };

            var moves = Resolve(belts, sinks: null);

            Assert.AreEqual(0, moves.Count,
                "A fully-packed belt line with a dead end produces no moves (snapshot occupancy blocks every step).");
        }

        [Test]
        public void Sink_Cell_Always_Accepts_The_Head_Item()
        {
            // A 2-cell +X line whose head's target cell is a SINK (machine input). Sinks always accept (deposit),
            // even though they are not conveyor cells in cellToIndex.
            var belts = new[]
            {
                new Belt { Cell = new int2(0, 0), Dir = DirPosX, Res = ResourceId.Biomass, Cnt = 2 },
            };
            var sinks = new[] { new int2(1, 0) };

            var moves = Resolve(belts, sinks);

            Assert.AreEqual(1, moves.Count, "The item moves into the adjacent sink.");
            Assert.AreEqual(new int2(1, 0), moves[0], "The destination is the sink cell.");
        }

        [Test]
        public void Y_Junction_Deterministic_Winner_By_Lowest_CellKey_Loser_Stalls()
        {
            // Two source belts both feed the SAME destination (1,1):
            //   A at (0,1) facing +X  -> (1,1)
            //   B at (1,0) facing +Z  -> (1,1)
            // The destination is an EMPTY conveyor cell -> it accepts AT MOST ONE; the tie breaks to the lowest
            // CellKey source. The loser STALLS (no move) with no item loss.
            var dst = new int2(1, 1);
            var a = new int2(0, 1);
            var b = new int2(1, 0);

            var belts = new[]
            {
                new Belt { Cell = a, Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },     // index 0
                new Belt { Cell = b, Dir = DirPosZ, Res = ResourceId.Ore, Cnt = 1 },     // index 1
                new Belt { Cell = dst, Dir = DirPosX, Res = 0, Cnt = 0 },                // index 2 (empty target)
            };

            var moves = Resolve(belts, sinks: null);

            // Exactly one of the two contenders wins; it targets dst. The winner is the lower-CellKey source.
            int winnerIdx = ConveyorMath.CellKey(a) < ConveyorMath.CellKey(b) ? 0 : 1;
            int loserIdx = winnerIdx == 0 ? 1 : 0;

            Assert.AreEqual(1, moves.Count, "Only one item may enter the shared (empty) destination this tick.");
            Assert.IsTrue(moves.ContainsKey(winnerIdx), "The lowest-CellKey source wins the contended cell.");
            Assert.AreEqual(dst, moves[winnerIdx], "The winner moves into the shared destination.");
            Assert.IsFalse(moves.ContainsKey(loserIdx), "The loser stalls in place (no item loss).");
        }

        [Test]
        public void Blocked_Cell_Stalls_The_Source()
        {
            // Head item at (0,0) facing +X; (1,0) is OCCUPIED in the snapshot (and its own item can't move because
            // (2,0) is not a belt/sink). The head item must STALL, not overwrite or destroy the blocker.
            var belts = new[]
            {
                new Belt { Cell = new int2(0, 0), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = new int2(1, 0), Dir = DirPosX, Res = ResourceId.Biomass, Cnt = 1 }, // blocker; -> (2,0) dead end
            };

            var moves = Resolve(belts, sinks: null);

            Assert.IsFalse(moves.ContainsKey(0), "A source whose destination is occupied in the snapshot stalls (no loss).");
            Assert.IsFalse(moves.ContainsKey(1), "The blocker itself has a dead-end target and also stalls.");
            Assert.AreEqual(0, moves.Count, "Nothing moves when the only path is blocked.");
        }

        [Test]
        public void EndState_Is_Identical_Under_Two_Shuffled_Input_Orders()
        {
            // A Y-junction plus a 4-cell line, fed in two DIFFERENT array orders. Because ResolveMoves iterates
            // sources SORTED by CellKey (not hashmap/array order), the resolved set of (srcCell -> destCell) moves
            // must be byte-for-byte identical regardless of input ordering. We compare keyed by CELL (order-stable),
            // not by array index (which differs between the two orderings).
            var dst = new int2(1, 1);
            var aCell = new int2(0, 1);
            var bCell = new int2(1, 0);

            var orderOne = new[]
            {
                new Belt { Cell = aCell, Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = bCell, Dir = DirPosZ, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = dst, Dir = DirPosX, Res = 0, Cnt = 0 },
                new Belt { Cell = new int2(5, 5), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 }, // lone item -> (6,5) sink
            };
            // Same scenario, shuffled: reverse + interleave the array order.
            var orderTwo = new[]
            {
                new Belt { Cell = new int2(5, 5), Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = dst, Dir = DirPosX, Res = 0, Cnt = 0 },
                new Belt { Cell = bCell, Dir = DirPosZ, Res = ResourceId.Ore, Cnt = 1 },
                new Belt { Cell = aCell, Dir = DirPosX, Res = ResourceId.Ore, Cnt = 1 },
            };
            var sinks = new[] { new int2(6, 5) };

            var movesOne = Resolve(orderOne, sinks);
            var movesTwo = Resolve(orderTwo, sinks);

            // Re-key both result sets by the SOURCE CELL (stable across orderings) -> destination cell.
            var byCellOne = ReKeyBySourceCell(orderOne, movesOne);
            var byCellTwo = ReKeyBySourceCell(orderTwo, movesTwo);

            Assert.AreEqual(byCellOne.Count, byCellTwo.Count, "Both orderings resolve the same number of moves.");
            foreach (var kv in byCellOne)
            {
                Assert.IsTrue(byCellTwo.ContainsKey(kv.Key), $"Source cell {kv.Key} moved in order-1 but not order-2.");
                Assert.AreEqual(kv.Value, byCellTwo[kv.Key], $"Source cell {kv.Key} resolved to a different destination across orderings (NON-deterministic).");
            }
        }

        static Dictionary<int2, int2> ReKeyBySourceCell(Belt[] belts, Dictionary<int, int2> movesBySrcIdx)
        {
            var byCell = new Dictionary<int2, int2>();
            foreach (var kv in movesBySrcIdx)
                byCell[belts[kv.Key].Cell] = kv.Value;
            return byCell;
        }
    }
}