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, INavCoverageReader
    {
        private readonly IChunkNavSourceReader chunkNavSourceReader;
        private readonly IWorldInterestReader worldInterestReader;
        private readonly INavCoverageHintReader navCoverageHintReader;
        private readonly ISubscriber<ChunkNavGeometryReadyMessage> chunkReadySubscriber;
        private readonly ISubscriber<ChunkNavGeometryRemovedMessage> chunkRemovedSubscriber;
        private readonly ISubscriber<WorldInterestChangedMessage> worldInterestChangedSubscriber;
        private readonly ISubscriber<NavCoverageHintChangedMessage> navCoverageHintChangedSubscriber;
        private readonly VoxelWorldNavMeshConfig config;
        private readonly Dictionary<int, NavCoverageWindowRuntime> coverageWindows = new Dictionary<int, NavCoverageWindowRuntime>();
        private readonly Queue<int> dirtyCoverageWindowIds = new Queue<int>();
        private readonly HashSet<int> queuedCoverageWindowIds = new HashSet<int>();
        private readonly List<int> dirtyCoverageWindowCandidates = new List<int>(16);
        private readonly List<WorldInterestPoint> interestPoints = new List<WorldInterestPoint>(8);
        private readonly List<Vector2Int> loadedChunkCoords = new List<Vector2Int>(128);
        private readonly List<NavMeshBuildSource> buildSources = new List<NavMeshBuildSource>(256);
        private readonly List<DesiredCoverageWindow> desiredCoverageWindows = new List<DesiredCoverageWindow>(8);
        private readonly List<ClusterAccumulator> clusterAccumulators = new List<ClusterAccumulator>(8);
        private readonly List<IDisposable> subscriptions = new List<IDisposable>(4);

        private int nextCoverageWindowId = 1;
        private int? activeBuildWindowId;

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

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

            RefreshInterestPoints();
            SyncCoverageWindows();
            MarkAllCoverageWindowsDirty();
        }

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

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

                int windowId = DequeueBestDirtyCoverageWindow();
                if (!TryStartCoverageBuild(windowId))
                {
                    startedBuilds++;
                    continue;
                }

                startedBuilds++;
            }
        }

        public bool IsPositionCovered(Vector3 worldPosition)
        {
            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                NavCoverageWindowRuntime window = pair.Value;
                if (window.State == NavCoverageState.Ready && NavMeshBoundsUtility.ContainsXZ(window.CoverageBounds, worldPosition))
                {
                    return true;
                }
            }

            return false;
        }

        public void GetCoverageWindows(List<NavCoverageWindowSnapshot> results)
        {
            if (results == null)
            {
                throw new ArgumentNullException(nameof(results));
            }

            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                NavCoverageWindowRuntime window = pair.Value;
                results.Add(new NavCoverageWindowSnapshot(window.Id, window.CoverageBounds, window.State, window.InterestCount));
            }
        }

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

            subscriptions.Clear();

            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                pair.Value.Dispose();
            }

            coverageWindows.Clear();
            queuedCoverageWindowIds.Clear();
            dirtyCoverageWindowIds.Clear();
            desiredCoverageWindows.Clear();
            clusterAccumulators.Clear();
            activeBuildWindowId = null;
        }

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

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

        private void OnWorldInterestChanged(WorldInterestChangedMessage message)
        {
            RefreshInterestPoints();
            SyncCoverageWindows();
            MarkAllCoverageWindowsDirty();
        }

        private void OnNavCoverageHintChanged(NavCoverageHintChangedMessage message)
        {
            RefreshInterestPoints();
            SyncCoverageWindows();
            MarkAllCoverageWindowsDirty();
        }

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

        private void SyncCoverageWindows()
        {
            NavCoveragePlanning.BuildDesiredCoverageWindows(
                interestPoints,
                config,
                Mathf.Max(1f, chunkNavSourceReader.ChunkWorldSize),
                desiredCoverageWindows,
                clusterAccumulators);

            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                pair.Value.MatchedThisFrame = false;
            }

            for (int i = 0; i < desiredCoverageWindows.Count; i++)
            {
                DesiredCoverageWindow desiredWindow = desiredCoverageWindows[i];
                NavCoverageWindowRuntime runtime = NavCoveragePlanning.FindBestMatchingCoverageWindow(
                    desiredWindow,
                    coverageWindows,
                    Mathf.Max(1f, chunkNavSourceReader.ChunkWorldSize),
                    config);
                if (runtime == null)
                {
                    runtime = new NavCoverageWindowRuntime(nextCoverageWindowId++, desiredWindow.CoverageBounds, desiredWindow.Priority, desiredWindow.InterestCount);
                    runtime.MatchedThisFrame = true;
                    coverageWindows.Add(runtime.Id, runtime);
                    EnqueueDirtyCoverageWindow(runtime.Id);
                    continue;
                }

                runtime.MatchedThisFrame = true;
                runtime.Priority = desiredWindow.Priority;
                runtime.InterestCount = desiredWindow.InterestCount;

                if (!NavMeshBoundsUtility.BoundsApproximatelyEqual(runtime.CoverageBounds, desiredWindow.CoverageBounds))
                {
                    runtime.CoverageBounds = desiredWindow.CoverageBounds;
                    runtime.State = NavCoverageState.Pending;
                    EnqueueDirtyCoverageWindow(runtime.Id);
                }
            }

            RemoveUnmatchedCoverageWindows();
        }

        private void RemoveUnmatchedCoverageWindows()
        {
            List<int> windowsToRemove = null;

            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                if (pair.Value.MatchedThisFrame)
                {
                    continue;
                }

                windowsToRemove ??= new List<int>();
                windowsToRemove.Add(pair.Key);
            }

            if (windowsToRemove == null)
            {
                return;
            }

            for (int i = 0; i < windowsToRemove.Count; i++)
            {
                RemoveCoverageWindow(windowsToRemove[i]);
            }
        }

        private void MarkAllCoverageWindowsDirty()
        {
            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                EnqueueDirtyCoverageWindow(pair.Key);
            }
        }

        private void MarkCoverageWindowsDirtyForChunk(Vector2Int chunkCoord)
        {
            Bounds chunkBounds = NavBuildSourceCollector.ExpandChunkBounds(
                chunkNavSourceReader,
                NavBuildSourceCollector.GetChunkWorldBounds(chunkNavSourceReader, chunkCoord),
                Mathf.Max(0, config.chunkCollectionMarginInChunks));

            foreach (KeyValuePair<int, NavCoverageWindowRuntime> pair in coverageWindows)
            {
                Bounds invalidationBounds = NavBuildSourceCollector.ExpandCoverageBounds(
                    chunkNavSourceReader,
                    pair.Value.CoverageBounds,
                    Mathf.Max(0, config.chunkCollectionMarginInChunks));

                if (NavMeshBoundsUtility.IntersectsXZ(invalidationBounds, chunkBounds))
                {
                    EnqueueDirtyCoverageWindow(pair.Key);
                }
            }
        }

        private void EnqueueDirtyCoverageWindow(int windowId)
        {
            if (!queuedCoverageWindowIds.Add(windowId))
            {
                if (activeBuildWindowId.HasValue && activeBuildWindowId.Value == windowId && coverageWindows.TryGetValue(windowId, out NavCoverageWindowRuntime activeWindow))
                {
                    activeWindow.BuildRequestedWhileRunning = true;
                }

                return;
            }

            dirtyCoverageWindowIds.Enqueue(windowId);
        }

        private int DequeueBestDirtyCoverageWindow()
        {
            dirtyCoverageWindowCandidates.Clear();
            while (dirtyCoverageWindowIds.Count > 0)
            {
                dirtyCoverageWindowCandidates.Add(dirtyCoverageWindowIds.Dequeue());
            }

            int bestIndex = 0;
            float bestScore = float.MaxValue;
            for (int i = 0; i < dirtyCoverageWindowCandidates.Count; i++)
            {
                int windowId = dirtyCoverageWindowCandidates[i];
                if (!coverageWindows.TryGetValue(windowId, out NavCoverageWindowRuntime window))
                {
                    continue;
                }

                float score = NavCoveragePlanning.GetCoveragePriorityScore(window, interestPoints);
                if (score < bestScore)
                {
                    bestScore = score;
                    bestIndex = i;
                }
            }

            int bestWindowId = dirtyCoverageWindowCandidates[bestIndex];
            queuedCoverageWindowIds.Remove(bestWindowId);

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

                dirtyCoverageWindowIds.Enqueue(dirtyCoverageWindowCandidates[i]);
            }

            dirtyCoverageWindowCandidates.Clear();
            return bestWindowId;
        }

        private bool TryStartCoverageBuild(int windowId)
        {
            if (!coverageWindows.TryGetValue(windowId, out NavCoverageWindowRuntime window))
            {
                return false;
            }

            buildSources.Clear();
            window.CollectionBounds = NavBuildSourceCollector.ExpandCoverageBounds(
                chunkNavSourceReader,
                window.CoverageBounds,
                Mathf.Max(0, config.chunkCollectionMarginInChunks));

            bool hasSources = NavBuildSourceCollector.CollectBuildSources(
                chunkNavSourceReader,
                window.CollectionBounds,
                loadedChunkCoords,
                buildSources);

            if (!hasSources || buildSources.Count == 0)
            {
                window.State = NavCoverageState.Pending;
                window.ResetCoverageData();
                return false;
            }

            Bounds buildBounds = NavMeshBoundsUtility.CalculateBounds(buildSources);
            NavMeshBoundsUtility.ExpandBounds(ref buildBounds, config.navBoundsHorizontalPadding, config.navBoundsVerticalPadding);
            window.BuildBounds = buildBounds;
            window.BuildRequestedWhileRunning = false;

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

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

            NavMeshBuildSettings buildSettings = UnityNavMesh.GetSettingsByID(config.agentTypeId);
            window.ActiveBuild = UnityNavMeshBuilder.UpdateNavMeshDataAsync(window.NavMeshData, buildSettings, buildSources, buildBounds);
            window.State = NavCoverageState.Building;
            activeBuildWindowId = windowId;
            return true;
        }

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

            if (!coverageWindows.TryGetValue(activeBuildWindowId.Value, out NavCoverageWindowRuntime window))
            {
                activeBuildWindowId = null;
                return;
            }

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

            window.ActiveBuild = null;
            window.State = NavCoverageState.Ready;
            int completedWindowId = activeBuildWindowId.Value;
            activeBuildWindowId = null;

            if (window.BuildRequestedWhileRunning)
            {
                window.BuildRequestedWhileRunning = false;
                EnqueueDirtyCoverageWindow(completedWindowId);
            }
        }

        private void RemoveCoverageWindow(int windowId)
        {
            if (!coverageWindows.TryGetValue(windowId, out NavCoverageWindowRuntime window))
            {
                return;
            }

            if (activeBuildWindowId.HasValue && activeBuildWindowId.Value == windowId)
            {
                activeBuildWindowId = null;
            }

            window.Dispose();
            coverageWindows.Remove(windowId);
            queuedCoverageWindowIds.Remove(windowId);
        }
    }
}