// imports...


namespace ECS.Systems
{
   // Simple struct to hold the direction and coordinates of a diggable neighbor
   public struct DiggableNeighbor
   {
       public int2 Direction;
       public int2 Target;
   }


   [UpdateInGroup(typeof(SimulationSystemGroup))]
   [UpdateBefore(typeof(AntMoveSystem))]
   [BurstCompile]
   public partial struct DiggerAntSystem : ISystem
   {
       /// <summary>
       /// Find Digger ants that don't have a target and aren't digging.
       /// Calculate a new target and enable
       /// </summary>
       private EntityQuery _diggerAntQuery;


	 /// This isn't ideal, Systems should be stateless 
       private NativeArray<Random> _threadRandoms;


       public void OnCreate(ref SystemState state)
       {
           state.RequireForUpdate<EndSimulationEntityCommandBufferSystem.Singleton>();
           state.RequireForUpdate<FlowFields>();
           state.RequireForUpdate<Tilemap>();
           _threadRandoms = new NativeArray<Random>(
               JobsUtility.MaxJobThreadCount,
               Allocator.Persistent);


           RandomUtils.InitializeThreadRandoms(_threadRandoms, 652763, 17);
           // Ants assigned to digging that aren't digging, hauling, moving or resting
           _diggerAntQuery = new EntityQueryBuilder(Allocator.Temp)
               .WithAll<AnimationIndexOverride, AntTag, Assignment, LocalTransform>()
               .WithDisabled<Digging, Hauling, Heading, Resting>()
               .Build(state.EntityManager);
           _diggerAntQuery.SetSharedComponentFilter(new Assignment
           {
               Value = AssignmentType.Dig
           });
       }


       public void OnUpdate(ref SystemState state)
       {
           // The tilemap is a grid but we store it as flattened array
           var tilemap = SystemAPI.GetSingleton<Tilemap>();
           var flowFields = SystemAPI.GetSingleton<FlowFields>();
           var neighbors = new NativeArray<int2>(8, Allocator.TempJob);
           TilemapUtils.PopulateNeighbors(neighbors);


           // Get the built-in system that handles batching at the end of the simulation
           var ecbSingleton = SystemAPI.GetSingleton<EndSimulationEntityCommandBufferSystem.Singleton>();
           // Create the thread-safe writer
           var ecb = ecbSingleton.CreateCommandBuffer(state.WorldUnmanaged).AsParallelWriter();


           var handle = new DiggerAntJob
           {
               DiggerEntityMap = tilemap.DiggerEntityMap.AsParallelWriter(),
               DiggingFlowField = flowFields.DiggingVectorField,
               DiggingIntegrationField = flowFields.DiggingIntegrationField,
               ECB = ecb,
               MapHeight = TilemapConstants.Height,
               MapWidth = TilemapConstants.Width,
               Neighbors = neighbors,
               TileStates = tilemap.TileStates,
               ThreadRandoms = _threadRandoms,
           }.ScheduleParallel(_diggerAntQuery, state.Dependency);
           // In the absence of a Complete() handle ECS will resolve the dependencies automatically and run in parallel.
           // Dependent systems wait until completion to run
           state.Dependency = handle;


       }


       [BurstCompile]
       public partial struct DiggerAntJob : IJobEntity
       {
           public NativeParallelMultiHashMap<int, Entity>.ParallelWriter DiggerEntityMap;
           [ReadOnly] public NativeArray<int2> DiggingFlowField;
           [ReadOnly] public NativeArray<ushort> DiggingIntegrationField;
           public EntityCommandBuffer.ParallelWriter ECB;
           public int MapHeight;
           public int MapWidth;
           [DeallocateOnJobCompletion] [ReadOnly] public NativeArray<int2> Neighbors;
           [ReadOnly] public NativeArray<byte> TileStates;
           [NativeSetThreadIndex] private int _threadIndex;


           [NativeDisableContainerSafetyRestriction]
           public NativeArray<Random> ThreadRandoms;


           public void Execute(
               ref AnimationIndexOverride animationIndexOverride,
               [ChunkIndexInQuery] int chunkIndex,
               ref Digging digging,
               ref Heading heading,
               EnabledRefRW<Digging> isDiggingEnabled,
               EnabledRefRW<Heading> isHeadingEnabled,
               EnabledRefRW<Resting> isRestingEnabled,
               in LocalTransform localTransform,
               ref Resting resting,
               Entity entity)
           {
               var random = ThreadRandoms[_threadIndex];


               // 20% chance to just rest regardless of where the ant is standing
               if (random.NextInt(0, 10) > 1)
               {
                   // Where is the ant currently standing in grid coordinates?
                   var currentX = (int)math.round(localTransform.Position.x);
                   var currentY = (int)math.round(-localTransform.Position.y);
                   var currentIndex = TilemapUtils.GetFlattenedIndex(currentX, currentY);
                   var currentTileState = TileStates[currentIndex];


                   // If standing on a diggable tile, dig for a random duration
                   // and finish the iteration(return)
                   if (currentTileState == (byte)TileState.Filled)
                   {
                       animationIndexOverride.Value = (int)AntAnimationIndex.Dig;
                       var duration = random.NextFloat(1.5f, 3.5f);
                       digging.Duration = duration;
                       isDiggingEnabled.ValueRW = true;
                       DiggerEntityMap.Add(currentIndex, entity);
                       return;
                   }


                   // If not standing on a diggable tile, then check the flow field
                   // Read the Flow Field arrow for this tile
                   var flowDirection = DiggingFlowField[currentIndex];
                   var targetX = currentX + flowDirection.x;
                   var targetY = currentY + flowDirection.y;
                   // Debug.Log($"Flow direction: {flowDirection}");


                   // If the flow field is not 0,0, then there is still a path towards a diggable tile
                   // set the heading with the flow field direction
                   if (!flowDirection.Equals(int2.zero))
                   {
                       animationIndexOverride.Value = (int)AntAnimationIndex.Walk;
                       // Convert grid direction to a world direction (invert Y)
                       heading.Direction = new int2(flowDirection.x, -flowDirection.y);
                       heading.Target = new int2(targetX, -targetY);
                       isHeadingEnabled.ValueRW = true;
                   }
                   else
                   {
                       // If the flow field is 0,0, we are standing next to a diggable tile or there is no path
                       // to a diggable tile.
                       // Debug.Log("Flow field is 0");


                       // WE REACHED THE END OF THE FLOW FIELD.
                       // Check the 8 adjacent neighbors for a dirt tile. We wouldn't need to check
                       // for the tiles that we just walked through, but it's easier to check all neighbors
                       // than to hold the backtracking state
                       var diggableNeighbors = new NativeList<DiggableNeighbor>(7, Allocator.TempJob);


                       for (var i = 0; i < Neighbors.Length; i++)
                       {
                           var nx = currentX + Neighbors[i].x;
                           var ny = currentY + Neighbors[i].y;


                           // Prevent out-of-bounds
                           if (nx < 0 || nx >= MapWidth || ny < 0 || ny >= MapHeight) continue;


                           var nIndex = TilemapUtils.GetFlattenedIndex(nx, ny);
                           if (TileStates[nIndex] != (byte)TileState.Filled) continue;
                           diggableNeighbors.Add(new DiggableNeighbor
                           {
                               Direction = new int2(Neighbors[i].x, Neighbors[i].y),
                               Target = new int2(nx, -ny)
                           });


                       }


                       // If there is a diggable neighbor, move to the dirt and start digging.
                       if (!diggableNeighbors.IsEmpty)
                       {
                           // Set animation index to walking
                           animationIndexOverride.Value = (int)AntAnimationIndex.Walk;


                           // Select a diggable neighbor at random
                           var diggableNeighbor = diggableNeighbors[random.NextInt(0, diggableNeighbors.Length)];
                           // Convert grid direction to a world direction (invert Y)
                           heading.Direction = diggableNeighbor.Direction;
                           heading.Target = diggableNeighbor.Target;
                           isHeadingEnabled.ValueRW = true;
                       }
                       else
                       {
                           // There is no path to a diggable tile. The flow field is 0, and there is no dirt in any adjacent tiles.
                           // Use ECB to revert the ant to Idle assignment.
                           // We should avoid ECB's when possible, in this case it's preferable
                           // since we can run this in parallel and also change the shared Assignment component
                           // value. The alternative is to convert the assignment to multiple enableable components


                           if (DiggingIntegrationField[currentIndex] == ushort.MaxValue)
                           {
                               // DEAD END / NO DIRT LEFT
                               // The BFS algorithm confirmed there is absolutely no path.
                               ECB.SetSharedComponent(chunkIndex, entity,
                                   new Assignment { Value = AssignmentType.Idle });
                               animationIndexOverride.Value = (int)AntAnimationIndex.Idle;
                               resting.Duration = random.NextFloat(1.5f, 3.5f);
                               isRestingEnabled.ValueRW = true;
                           }
                           else
                           {
                               // STALE MAP
                               // Rest for a brief moment and wait for the Flow Field to update!
                               animationIndexOverride.Value = (int)AntAnimationIndex.Idle;
                               resting.Duration = random.NextFloat(0.5f, 1.5f); // Short rest
                               isRestingEnabled.ValueRW = true;
                           }
                       }


                       diggableNeighbors.Dispose();
                   }
               }
               else
               {
                   // Randomly resting
                   animationIndexOverride.Value = (int)AntAnimationIndex.Idle;
                   resting.Duration = random.NextFloat(1.5f, 3.5f);
                   isRestingEnabled.ValueRW = true;
               }


               ThreadRandoms[_threadIndex] = random;
           }
       }


       public void OnDestroy(ref SystemState state)
       {
           _diggerAntQuery.Dispose();
           _threadRandoms.Dispose();
       }
   }
}

// imports...


namespace ECS.Systems
{
   /// <summary>
   /// Apply linear speed to ants so they will move in their <see cref="Heading"/> direction
   /// </summary>
   /// <remarks>
   ///
   /// </remarks>
   [UpdateInGroup(typeof(SimulationSystemGroup))]
   [BurstCompile]
   public partial struct AntMoveSystem : ISystem
   {
       /// <summary>
       /// Query focusing on ants that need to move to another tile
       /// </summary>
       private EntityQuery _antQuery;


       [BurstCompile]
       public void OnCreate(ref SystemState state)
       {
           _antQuery = new EntityQueryBuilder(Allocator.Temp)
               .WithAll<
                   AntTag,
                   FlipXOverride,
                   Heading,
                   LocalTransform,
                   MovementSpeed>()
               .Build(state.EntityManager);
       }


       [BurstCompile]
       public void OnUpdate(ref SystemState state)
       {
           var deltaTime = SystemAPI.Time.DeltaTime;


           var handle = new AntMoveJob
           {
               DeltaTime = deltaTime
           }.ScheduleParallel(_antQuery, state.Dependency);
           state.Dependency = handle;
       }


       [BurstCompile]
       public partial struct AntMoveJob : IJobEntity
       {
           public float DeltaTime;


           public void Execute(
               ref FlipXOverride flipXOverride,
               ref Heading heading,
               EnabledRefRW<Heading> isHeadingEnabled,
               in MovementSpeed movementSpeed,
               ref LocalTransform localTransform)
           {
               // Get the offset vector and the exact distance to the target
               var offset = heading.Target - localTransform.Position.xy;
               var distanceToTarget = math.length(offset);


               // Calculate wow far will the ant walk this frame
               var distanceToTravelThisFrame = movementSpeed.Value * DeltaTime;
               if (distanceToTarget <= distanceToTravelThisFrame)
               {
                   // Snap to the target since the ant will arrive or overshoot
                   localTransform.Position.xy = heading.Target;
                   isHeadingEnabled.ValueRW = false;
               }
               else
               {
                   // Normalize for constant speed
                   var normalizedDirection = offset / distanceToTarget;


                   // Face the correct direction
                   flipXOverride.Value = normalizedDirection.x > 0 ? 1 : -1;


                   // Walk forward
                   localTransform.Position.xy += normalizedDirection * distanceToTravelThisFrame;
               }
           }
       }
   }
}