TheDeclineOfWarriors/Assets/Features/VoxelWorldNavMesh/Runtime/VoxelWorldNavMeshService.cs

using System;
using System.Collections.Generic;
using InfiniteWorld.VoxelWorld.Contracts;
using MessagePipe;
using UnityEngine;
using UnityEngine.AI;
using VContainer.Unity;
using UnityNavMesh = UnityEngine.AI.NavMesh;
using UnityNavMeshBuilder = UnityEngine.AI.NavMeshBuilder;

namespace InfiniteWorld.VoxelWorld.NavMesh
{
    public sealed class VoxelWorldNavMeshService : IStartable, ITickable, IDisposable
    {
        private readonly IChunkNavSourceReader chunkNavSourceReader;
        private readonly IWorldInterestReader worldInterestReader;
        private readonly ISubscriber<ChunkNavGeometryReadyMessage> chunkReadySubscriber;
        private readonly ISubscriber<ChunkNavGeometryRemovedMessage> chunkRemovedSubscriber;
        private readonly ISubscriber<WorldInterestChangedMessage> worldInterestChangedSubscriber;
        private readonly VoxelWorldNavMeshConfig config;
        private readonly Dictionary<Vector2Int, NavRegionRuntime> navRegions = new Dictionary<Vector2Int, NavRegionRuntime>();
        private readonly Queue<Vector2Int> dirtyNavRegions = new Queue<Vector2Int>();
        private readonly HashSet<Vector2Int> queuedNavRegions = new HashSet<Vector2Int>();
        private readonly List<Vector2Int> loadedChunkCoords = new List<Vector2Int>(64);
        private readonly List<WorldInterestPoint> interestPoints = new List<WorldInterestPoint>(4);
        private readonly List<Vector2Int> dirtyRegionCandidates = new List<Vector2Int>(16);
        private readonly List<NavMeshBuildSource> buildSources = new List<NavMeshBuildSource>(64);
        private readonly HashSet<Vector2Int> currentInterestRegions = new HashSet<Vector2Int>();
        private readonly HashSet<Vector2Int> previousInterestRegions = new HashSet<Vector2Int>();
        private readonly List<IDisposable> subscriptions = new List<IDisposable>(3);

        private Vector2Int? activeBuildRegion;

        public VoxelWorldNavMeshService(
            IChunkNavSourceReader chunkNavSourceReader,
            IWorldInterestReader worldInterestReader,
            ISubscriber<ChunkNavGeometryReadyMessage> chunkReadySubscriber,
            ISubscriber<ChunkNavGeometryRemovedMessage> chunkRemovedSubscriber,
            ISubscriber<WorldInterestChangedMessage> worldInterestChangedSubscriber,
            VoxelWorldNavMeshConfig config)
        {
            this.chunkNavSourceReader = chunkNavSourceReader;
            this.worldInterestReader = worldInterestReader;
            this.chunkReadySubscriber = chunkReadySubscriber;
            this.chunkRemovedSubscriber = chunkRemovedSubscriber;
            this.worldInterestChangedSubscriber = worldInterestChangedSubscriber;
            this.config = config ?? new VoxelWorldNavMeshConfig();
        }

        public void Start()
        {
            subscriptions.Add(chunkReadySubscriber.Subscribe(OnChunkNavGeometryReady));
            subscriptions.Add(chunkRemovedSubscriber.Subscribe(OnChunkNavGeometryRemoved));
            subscriptions.Add(worldInterestChangedSubscriber.Subscribe(OnWorldInterestChanged));

            RefreshInterestPoints();

            loadedChunkCoords.Clear();
            chunkNavSourceReader.GetLoadedChunkCoords(loadedChunkCoords);
            for (int i = 0; i < loadedChunkCoords.Count; i++)
            {
                MarkDirtyForChunk(loadedChunkCoords[i]);
            }

            MarkWarmupRegionsDirty();
        }

        public void Tick()
        {
            RefreshInterestPoints();
            CompleteFinishedBuild();

            int startedBuilds = 0;
            int maxBuilds = Mathf.Max(1, config.maxNavMeshBuildsPerFrame);
            while (startedBuilds < maxBuilds)
            {
                if (activeBuildRegion.HasValue || dirtyNavRegions.Count == 0)
                {
                    break;
                }

                Vector2Int regionCoord = DequeueBestDirtyRegion();
                if (!TryStartRegionBuild(regionCoord))
                {
                    startedBuilds++;
                    continue;
                }

                startedBuilds++;
            }
        }

        public void Dispose()
        {
            for (int i = 0; i < subscriptions.Count; i++)
            {
                subscriptions[i]?.Dispose();
            }

            subscriptions.Clear();

            foreach (KeyValuePair<Vector2Int, NavRegionRuntime> pair in navRegions)
            {
                pair.Value.Dispose();
            }

            navRegions.Clear();
            queuedNavRegions.Clear();
            dirtyNavRegions.Clear();
            currentInterestRegions.Clear();
            previousInterestRegions.Clear();
            activeBuildRegion = null;
        }

        private void OnChunkNavGeometryReady(ChunkNavGeometryReadyMessage message)
        {
            MarkDirtyForChunk(message.Coord);
        }

        private void OnChunkNavGeometryRemoved(ChunkNavGeometryRemovedMessage message)
        {
            MarkDirtyForChunk(message.Coord);
        }

        private void OnWorldInterestChanged(WorldInterestChangedMessage message)
        {
            RefreshInterestPoints();
            MarkWarmupRegionsDirty();
        }

        private void RefreshInterestPoints()
        {
            interestPoints.Clear();
            worldInterestReader.GetInterestPoints(interestPoints);

            previousInterestRegions.Clear();
            foreach (Vector2Int region in currentInterestRegions)
            {
                previousInterestRegions.Add(region);
            }

            currentInterestRegions.Clear();
            for (int i = 0; i < interestPoints.Count; i++)
            {
                currentInterestRegions.Add(ChunkToRegion(WorldToChunk(interestPoints[i].Position)));
            }

            if (!AreSetsEqual(previousInterestRegions, currentInterestRegions))
            {
                MarkWarmupRegionsDirty();
            }
        }

        private void MarkWarmupRegionsDirty()
        {
            int radius = Mathf.Max(0, config.navWarmupRadiusInRegions);
            foreach (Vector2Int region in currentInterestRegions)
            {
                for (int y = -radius; y <= radius; y++)
                {
                    for (int x = -radius; x <= radius; x++)
                    {
                        EnqueueDirtyRegion(new Vector2Int(region.x + x, region.y + y));
                    }
                }
            }
        }

        private void MarkDirtyForChunk(Vector2Int chunkCoord)
        {
            int regionSize = Mathf.Max(1, config.navRegionSizeInChunks);
            Vector2Int regionCoord = ChunkToRegion(chunkCoord);
            EnqueueDirtyRegion(regionCoord);

            int localX = PositiveModulo(chunkCoord.x, regionSize);
            int localY = PositiveModulo(chunkCoord.y, regionSize);
            if (localX == 0)
            {
                EnqueueDirtyRegion(regionCoord + Vector2Int.left);
            }

            if (localX == regionSize - 1)
            {
                EnqueueDirtyRegion(regionCoord + Vector2Int.right);
            }

            if (localY == 0)
            {
                EnqueueDirtyRegion(regionCoord + Vector2Int.down);
            }

            if (localY == regionSize - 1)
            {
                EnqueueDirtyRegion(regionCoord + Vector2Int.up);
            }
        }

        private void EnqueueDirtyRegion(Vector2Int regionCoord)
        {
            if (!queuedNavRegions.Add(regionCoord))
            {
                if (activeBuildRegion.HasValue && activeBuildRegion.Value == regionCoord && navRegions.TryGetValue(regionCoord, out NavRegionRuntime activeRegion))
                {
                    activeRegion.BuildRequestedWhileRunning = true;
                }

                return;
            }

            dirtyNavRegions.Enqueue(regionCoord);
        }

        private Vector2Int DequeueBestDirtyRegion()
        {
            dirtyRegionCandidates.Clear();
            while (dirtyNavRegions.Count > 0)
            {
                dirtyRegionCandidates.Add(dirtyNavRegions.Dequeue());
            }

            int bestIndex = 0;
            float bestScore = float.MaxValue;
            for (int i = 0; i < dirtyRegionCandidates.Count; i++)
            {
                float score = GetRegionPriorityScore(dirtyRegionCandidates[i]);
                if (score < bestScore)
                {
                    bestScore = score;
                    bestIndex = i;
                }
            }

            Vector2Int best = dirtyRegionCandidates[bestIndex];
            queuedNavRegions.Remove(best);

            for (int i = 0; i < dirtyRegionCandidates.Count; i++)
            {
                if (i == bestIndex)
                {
                    continue;
                }

                dirtyNavRegions.Enqueue(dirtyRegionCandidates[i]);
            }

            dirtyRegionCandidates.Clear();
            return best;
        }

        private float GetRegionPriorityScore(Vector2Int regionCoord)
        {
            if (interestPoints.Count == 0)
            {
                return 0f;
            }

            Vector3 regionCenter = GetRegionCenter(regionCoord);
            float bestDistance = float.MaxValue;
            for (int i = 0; i < interestPoints.Count; i++)
            {
                float priority = Mathf.Max(0.01f, interestPoints[i].Priority);
                float distance = Vector3.SqrMagnitude(regionCenter - interestPoints[i].Position) / priority;
                if (distance < bestDistance)
                {
                    bestDistance = distance;
                }
            }

            return bestDistance;
        }

        private bool TryStartRegionBuild(Vector2Int regionCoord)
        {
            buildSources.Clear();
            bool hasCoreChunk = CollectBuildSources(regionCoord, buildSources);
            if (!hasCoreChunk || buildSources.Count == 0)
            {
                RemoveRegion(regionCoord);
                return false;
            }

            Bounds buildBounds = CalculateBounds(buildSources);
            ExpandBounds(ref buildBounds, config.navBoundsHorizontalPadding, config.navBoundsVerticalPadding);

            NavRegionRuntime region = GetOrCreateRegion(regionCoord);
            region.BuildRequestedWhileRunning = false;
            region.BuildBounds = buildBounds;

            if (region.NavMeshData == null)
            {
                region.NavMeshData = new NavMeshData(config.agentTypeId);
            }

            if (!region.Instance.valid)
            {
                region.Instance = UnityNavMesh.AddNavMeshData(region.NavMeshData);
            }

            NavMeshBuildSettings buildSettings = UnityNavMesh.GetSettingsByID(config.agentTypeId);
            region.ActiveBuild = UnityNavMeshBuilder.UpdateNavMeshDataAsync(region.NavMeshData, buildSettings, buildSources, buildBounds);
            activeBuildRegion = regionCoord;
            return true;
        }

        private bool CollectBuildSources(Vector2Int regionCoord, List<NavMeshBuildSource> results)
        {
            int regionSize = Mathf.Max(1, config.navRegionSizeInChunks);
            int baseChunkX = regionCoord.x * regionSize;
            int baseChunkY = regionCoord.y * regionSize;
            bool hasCoreChunk = false;

            for (int y = -1; y <= regionSize; y++)
            {
                for (int x = -1; x <= regionSize; x++)
                {
                    Vector2Int chunkCoord = new Vector2Int(baseChunkX + x, baseChunkY + y);
                    if (!chunkNavSourceReader.TryGetChunkNavSourceSnapshot(chunkCoord, out ChunkNavSourceSnapshot snapshot) || snapshot.Sources == null || snapshot.Sources.Length == 0)
                    {
                        continue;
                    }

                    if (x >= 0 && x < regionSize && y >= 0 && y < regionSize)
                    {
                        hasCoreChunk = true;
                    }

                    AppendBuildSources(snapshot.Sources, results);
                }
            }

            return hasCoreChunk;
        }

        private static void AppendBuildSources(ChunkNavBuildSourceDescriptor[] descriptors, List<NavMeshBuildSource> results)
        {
            for (int i = 0; i < descriptors.Length; i++)
            {
                ChunkNavBuildSourceDescriptor descriptor = descriptors[i];
                if (descriptor.Shape == NavMeshBuildSourceShape.Mesh && descriptor.Mesh == null)
                {
                    continue;
                }

                NavMeshBuildSource source = new NavMeshBuildSource
                {
                    area = descriptor.Area,
                    shape = descriptor.Shape,
                    transform = descriptor.Transform,
                    size = descriptor.Size,
                    sourceObject = descriptor.Shape == NavMeshBuildSourceShape.Mesh ? descriptor.Mesh : null
                };

                results.Add(source);
            }
        }

        private void CompleteFinishedBuild()
        {
            if (!activeBuildRegion.HasValue)
            {
                return;
            }

            if (!navRegions.TryGetValue(activeBuildRegion.Value, out NavRegionRuntime region))
            {
                activeBuildRegion = null;
                return;
            }

            if (region.ActiveBuild != null && !region.ActiveBuild.isDone)
            {
                return;
            }

            region.ActiveBuild = null;
            Vector2Int completedRegion = activeBuildRegion.Value;
            activeBuildRegion = null;

            if (region.BuildRequestedWhileRunning)
            {
                region.BuildRequestedWhileRunning = false;
                EnqueueDirtyRegion(completedRegion);
            }
        }

        private NavRegionRuntime GetOrCreateRegion(Vector2Int regionCoord)
        {
            if (!navRegions.TryGetValue(regionCoord, out NavRegionRuntime region))
            {
                region = new NavRegionRuntime();
                navRegions.Add(regionCoord, region);
            }

            return region;
        }

        private void RemoveRegion(Vector2Int regionCoord)
        {
            if (!navRegions.TryGetValue(regionCoord, out NavRegionRuntime region))
            {
                return;
            }

            if (activeBuildRegion.HasValue && activeBuildRegion.Value == regionCoord)
            {
                activeBuildRegion = null;
            }

            region.Dispose();
            navRegions.Remove(regionCoord);
        }

        private static Bounds CalculateBounds(List<NavMeshBuildSource> sources)
        {
            Bounds bounds = GetSourceBounds(sources[0]);
            for (int i = 1; i < sources.Count; i++)
            {
                bounds.Encapsulate(GetSourceBounds(sources[i]));
            }

            return bounds;
        }

        private static Bounds GetSourceBounds(NavMeshBuildSource source)
        {
            if (source.shape == NavMeshBuildSourceShape.Box)
            {
                return TransformBounds(source.transform, new Bounds(Vector3.zero, source.size));
            }

            Mesh mesh = source.sourceObject as Mesh;
            if (mesh != null)
            {
                return TransformBounds(source.transform, mesh.bounds);
            }

            return new Bounds(source.transform.GetColumn(3), Vector3.zero);
        }

        private static Bounds TransformBounds(Matrix4x4 matrix, Bounds localBounds)
        {
            Vector3 center = localBounds.center;
            Vector3 extents = localBounds.extents;

            Vector3[] corners =
            {
                new Vector3(center.x - extents.x, center.y - extents.y, center.z - extents.z),
                new Vector3(center.x - extents.x, center.y - extents.y, center.z + extents.z),
                new Vector3(center.x - extents.x, center.y + extents.y, center.z - extents.z),
                new Vector3(center.x - extents.x, center.y + extents.y, center.z + extents.z),
                new Vector3(center.x + extents.x, center.y - extents.y, center.z - extents.z),
                new Vector3(center.x + extents.x, center.y - extents.y, center.z + extents.z),
                new Vector3(center.x + extents.x, center.y + extents.y, center.z - extents.z),
                new Vector3(center.x + extents.x, center.y + extents.y, center.z + extents.z)
            };

            Bounds worldBounds = new Bounds(matrix.MultiplyPoint3x4(corners[0]), Vector3.zero);
            for (int i = 1; i < corners.Length; i++)
            {
                worldBounds.Encapsulate(matrix.MultiplyPoint3x4(corners[i]));
            }

            return worldBounds;
        }

        private static void ExpandBounds(ref Bounds bounds, float horizontalPadding, float verticalPadding)
        {
            Vector3 size = bounds.size;
            size.x = Mathf.Max(size.x + horizontalPadding * 2f, 0.1f);
            size.z = Mathf.Max(size.z + horizontalPadding * 2f, 0.1f);
            size.y = Mathf.Max(size.y + verticalPadding * 2f, 0.1f);
            bounds.size = size;
        }

        private Vector2Int WorldToChunk(Vector3 position)
        {
            float chunkSize = Mathf.Max(1f, chunkNavSourceReader.ChunkWorldSize);
            return new Vector2Int(
                Mathf.FloorToInt(position.x / chunkSize),
                Mathf.FloorToInt(position.z / chunkSize));
        }

        private Vector2Int ChunkToRegion(Vector2Int chunkCoord)
        {
            int size = Mathf.Max(1, config.navRegionSizeInChunks);
            return new Vector2Int(
                Mathf.FloorToInt(chunkCoord.x / (float)size),
                Mathf.FloorToInt(chunkCoord.y / (float)size));
        }

        private Vector3 GetRegionCenter(Vector2Int regionCoord)
        {
            float chunkSize = Mathf.Max(1f, chunkNavSourceReader.ChunkWorldSize);
            float regionSize = Mathf.Max(1, config.navRegionSizeInChunks) * chunkSize;
            return new Vector3(
                (regionCoord.x + 0.5f) * regionSize,
                0f,
                (regionCoord.y + 0.5f) * regionSize);
        }

        private static bool AreSetsEqual(HashSet<Vector2Int> left, HashSet<Vector2Int> right)
        {
            if (left.Count != right.Count)
            {
                return false;
            }

            foreach (Vector2Int value in left)
            {
                if (!right.Contains(value))
                {
                    return false;
                }
            }

            return true;
        }

        private static int PositiveModulo(int value, int modulus)
        {
            int result = value % modulus;
            return result < 0 ? result + modulus : result;
        }

        private sealed class NavRegionRuntime : IDisposable
        {
            public NavMeshData NavMeshData;
            public NavMeshDataInstance Instance;
            public AsyncOperation ActiveBuild;
            public bool BuildRequestedWhileRunning;
            public Bounds BuildBounds;

            public void Dispose()
            {
                if (ActiveBuild != null && !ActiveBuild.isDone && NavMeshData != null)
                {
                    UnityNavMeshBuilder.Cancel(NavMeshData);
                }

                if (Instance.valid)
                {
                    UnityNavMesh.RemoveNavMeshData(Instance);
                    Instance = default;
                }

                ActiveBuild = null;
                NavMeshData = null;
            }
        }
    }
}