VoxelEngine/World/Streaming/ChunkStreamer.cs

using System.Collections.Concurrent;
using System.Linq;
using OpenTK.Mathematics;
using Serilog;
using VoxelEngine.Core;
using VoxelEngine.World.Blocks;
using VoxelEngine.World.Chunk;
using VoxelEngine.World.WorldGen;
using VoxelChunk = VoxelEngine.World.Chunk.Chunk;

namespace VoxelEngine.World.Streaming;

public sealed class ChunkStreamer : IDisposable
{
    private static readonly ILogger Log = Serilog.Log.ForContext<ChunkStreamer>();

    // How many chunks below the surface the spawn area must load before the
    // loading screen clears (just enough to stand on and see; the rest streams in).
    private const int SpawnReadyDepthChunks = 2;

    private readonly Configuration _cfg;
    private readonly IWorld _world;
    private readonly Func<Vector3i, VoxelChunk> _generate;
    private readonly Func<int, int, int> _surfaceHeightAt;
    // Per chunk-column (cx,cz) surface-TOP height, memoized so each column's
    // samples are taken once and reused across ticks. Pruned to the horizontal
    // window each recompute so it stays bounded. Full-column loading only needs
    // the top; the bottom is always bedrock (cy 0).
    private readonly Dictionary<(int cx, int cz), int> _surfaceTopCache = new();
    private readonly IChunkMesher _mesher;
    private readonly BlockRegistry _registry;
    private readonly IJobExecutor _executor;
    private readonly Func<float[], IChunkRenderer> _rendererFactory;

    private readonly Dictionary<Vector3i, IChunkRenderer> _renderers = new();
    private readonly HashSet<Vector3i> _tracked = new();
    private readonly HashSet<Vector3i> _ready = new();
    private readonly HashSet<Vector3i> _dirty = new();
    private readonly PriorityQueue<Vector3i, int> _pendingLoad = new();
    private readonly ConcurrentQueue<ChunkJobResult> _completed = new();
    private readonly Queue<Vector3i> _pendingUnload = new();
    private readonly HashSet<Vector3i> _failed = new();

    private Vector3i _lastPlayerChunk = new(int.MinValue, int.MinValue, int.MinValue);
    private int _inFlight;
    private bool _spawnReady;
    private bool _spawnAnchorSet;
    private Vector3i _spawnAnchor;

    public int RendererCount => this._renderers.Count;
    public int TrackedCount => this._tracked.Count;
    public bool IsReady(Vector3i pos) => this._ready.Contains(pos);
    public bool SpawnReady => this._spawnReady;

    public void MarkDirty(Vector3i chunkPos)
    {
        if (!this._ready.Contains(chunkPos)) return;
        this._dirty.Add(chunkPos);
    }

    public ChunkStreamer(
        Configuration cfg,
        IWorld world,
        Func<Vector3i, VoxelChunk> generate,
        Func<int, int, int> surfaceHeightAt,
        IChunkMesher mesher,
        BlockRegistry registry,
        IJobExecutor executor,
        Func<float[], IChunkRenderer> rendererFactory)
    {
        this._cfg = cfg;
        this._world = world;
        this._generate = generate;
        this._surfaceHeightAt = surfaceHeightAt;
        this._mesher = mesher;
        this._registry = registry;
        this._executor = executor;
        this._rendererFactory = rendererFactory;
    }

    public void Tick(Vector3i playerChunkPos)
    {
        if (playerChunkPos != this._lastPlayerChunk)
        {
            this.RecomputeTargetSet(playerChunkPos);
            this._lastPlayerChunk = playerChunkPos;
        }

        // Unloads first to free in-flight slots, then dispatch BEFORE drain:
        // with a synchronous executor (tests), Submit() runs work to completion
        // before returning, so a single tick must drain its own dispatches.
        // With the production threadpool executor, jobs don't finish in this
        // tick -- Drain harvests results submitted on previous ticks. Either
        // way this ordering is correct.
        this.ProcessUnloads();
        this.DispatchJobs();
        this.DrainCompleted();
        this.UpdateSpawnReady();
    }

    // Draws every loaded chunk for which `isChunkVisible` returns true. Returns
    // the number of chunks actually drawn (passed the visibility test). The
    // visibility predicate lives with the caller so the streamer stays free of
    // any rendering/Frustum dependency.
    public int DrawVisible(Func<Vector3i, bool> isChunkVisible, Action<Vector3i, IChunkRenderer> draw)
    {
        int drawn = 0;
        foreach (KeyValuePair<Vector3i, IChunkRenderer> kvp in this._renderers)
        {
            if (!isChunkVisible(kvp.Key)) continue;
            draw(kvp.Key, kvp.Value);
            drawn++;
        }
        return drawn;
    }

    public void Dispose()
    {
        foreach (IChunkRenderer r in this._renderers.Values) r.Dispose();
        this._renderers.Clear();
    }

    private void RecomputeTargetSet(Vector3i p)
    {
        if (!this._spawnAnchorSet)
        {
            this._spawnAnchor = p;
            this._spawnAnchorSet = true;
        }

        HashSet<Vector3i> target = new();
        this.CollectTargets(p, this._cfg.RenderDistanceHorizontal, target);

        foreach (Vector3i pos in target)
        {
            if (this._failed.Contains(pos)) continue;
            if (this._tracked.Add(pos))
            {
                int dx = pos.X - p.X;
                int dy = pos.Y - p.Y;
                int dz = pos.Z - p.Z;
                int priority = dx * dx
                             + dz * dz
                             + dy * dy * this._cfg.VerticalPriorityWeight;
                this._pendingLoad.Enqueue(pos, priority);
            }
        }

        // Anything currently tracked but no longer in target → unload.
        List<Vector3i> toRemove = new();
        foreach (Vector3i tracked in this._tracked)
        {
            if (!target.Contains(tracked)) toRemove.Add(tracked);
        }
        foreach (Vector3i pos in toRemove)
        {
            this._tracked.Remove(pos);
            this._pendingUnload.Enqueue(pos);
        }

        this.PruneSurfaceCache(p);
    }

    // Drop cached surface tops for columns outside the current horizontal
    // window so the cache stays bounded as the player explores.
    private void PruneSurfaceCache(Vector3i p)
    {
        int keep = this._cfg.RenderDistanceHorizontal + 1;
        List<(int cx, int cz)>? stale = null;
        foreach ((int cx, int cz) key in this._surfaceTopCache.Keys)
        {
            if (Math.Abs(key.cx - p.X) > keep || Math.Abs(key.cz - p.Z) > keep)
                (stale ??= new()).Add(key);
        }
        if (stale != null)
            foreach ((int cx, int cz) key in stale) this._surfaceTopCache.Remove(key);
    }

    // Fills `into` with the chunk positions to load for a viewer at chunk
    // `center`, over a square horizontal region of Chebyshev radius `radiusH`.
    // Per column the full vertical stack from bedrock (cy 0) up to the
    // surface-top chunk (clamped up to the sea-level chunk so submerged columns
    // load their opaque-water surface) is loaded, plus any section above that
    // holding player edits (e.g. a built tower).
    private void CollectTargets(Vector3i center, int radiusH, HashSet<Vector3i> into)
    {
        for (int dx = -radiusH; dx <= radiusH; dx++)
        for (int dz = -radiusH; dz <= radiusH; dz++)
        {
            int cx = center.X + dx;
            int cz = center.Z + dz;

            int top = this.SurfaceTopChunk(cx, cz);
            AddColumn(into, cx, cz, 0, top);

            // Player-built structures can rise above the surface; load any edited
            // section above the column top so they remain visible.
            for (int cy = top + 1; cy < WorldConstants.WorldHeightChunks; cy++)
            {
                Vector3i pos = new(cx, cy, cz);
                if (this._world.ModificationStore.HasDiff(pos)) into.Add(pos);
            }
        }
    }

    private static void AddColumn(HashSet<Vector3i> into, int cx, int cz, int loCy, int hiCy)
    {
        int lo = Math.Max(0, loCy);
        int hi = Math.Min(WorldConstants.WorldHeightChunks - 1, hiCy);
        for (int cy = lo; cy <= hi; cy++)
            into.Add(new Vector3i(cx, cy, cz));
    }

    // Chunk Y of a column's loaded top: the terrain surface chunk, raised to the
    // sea-level (waterline) chunk for submerged columns, and clamped to the world
    // ceiling. The bottom of a loaded column is always bedrock (cy 0).
    private int SurfaceTopChunk(int cx, int cz)
        => Math.Min(
            WorldConstants.WorldHeightChunks - 1,
            Math.Max(
                ChunkMath.FloorDiv(this.SurfaceTopHeight(cx, cz), VoxelChunk.Size),
                ChunkMath.FloorDiv(WorldConstants.SeaLevel, VoxelChunk.Size)));

    // Max terrain surface height across a chunk column's 16x16 footprint,
    // sampled at the four corners and the centre (cheap bound on the column's
    // top). Cached per column. The bottom of the loaded stack is always bedrock,
    // so only the top is needed.
    private int SurfaceTopHeight(int cx, int cz)
    {
        if (this._surfaceTopCache.TryGetValue((cx, cz), out int cached))
            return cached;

        int x0 = cx * VoxelChunk.Size;
        int z0 = cz * VoxelChunk.Size;
        int x1 = x0 + VoxelChunk.Size - 1;
        int z1 = z0 + VoxelChunk.Size - 1;
        int xm = x0 + VoxelChunk.Size / 2;
        int zm = z0 + VoxelChunk.Size / 2;

        int max = Math.Max(
            Math.Max(this._surfaceHeightAt(x0, z0), this._surfaceHeightAt(x1, z0)),
            Math.Max(Math.Max(this._surfaceHeightAt(x0, z1), this._surfaceHeightAt(x1, z1)),
                     this._surfaceHeightAt(xm, zm)));

        this._surfaceTopCache[(cx, cz)] = max;
        return max;
    }

    private void DrainCompleted()
    {
        int budget = this._cfg.MaxMeshUploadsPerFrame;
        while (budget > 0 && this._completed.TryDequeue(out ChunkJobResult? result))
        {
            this._inFlight--;
            budget--;
            this.InstallResult(result);
        }
    }

    private void ProcessUnloads()
    {
        while (this._pendingUnload.TryDequeue(out Vector3i pos))
        {
            if (this._renderers.TryGetValue(pos, out IChunkRenderer? renderer))
            {
                renderer.Dispose();
                this._renderers.Remove(pos);
            }
            this._world.RemoveChunk(pos);
            this._ready.Remove(pos);
            this._failed.Remove(pos);
            Log.Debug("Unloaded {Pos}", pos);
        }
    }

    // The loading screen waits on this. It only needs the VISIBLE surface around
    // spawn to be ready — gating on the full bedrock column would block on the
    // deepest chunks, which carry the lowest load priority in the whole render
    // volume and stream in last (minutes for tall mountain spawn columns). Those
    // deep chunks are occluded at spawn and arrive shortly after.
    private void UpdateSpawnReady()
    {
        if (this._spawnReady || !this._spawnAnchorSet) return;

        int hs = this._cfg.SpawnReadyRadiusHorizontal;
        Vector3i a = this._spawnAnchor;
        for (int dx = -hs; dx <= hs; dx++)
        for (int dz = -hs; dz <= hs; dz++)
        {
            int cx = a.X + dx;
            int cz = a.Z + dz;
            int top = this.SurfaceTopChunk(cx, cz);
            int bottom = Math.Max(0, top - SpawnReadyDepthChunks);
            for (int cy = bottom; cy <= top; cy++)
            {
                Vector3i pos = new(cx, cy, cz);
                if (this._failed.Contains(pos)) continue;
                if (!this._ready.Contains(pos)) return;
            }
        }

        this._spawnReady = true;
        Log.Information("Spawn area ready at anchor {Anchor}", this._spawnAnchor);
    }

    private void InstallResult(ChunkJobResult result)
    {
        if (result.Error != null || result.Mesh == null)
        {
            this._failed.Add(result.Pos);
            Log.Error(result.Error, "Chunk {Pos} failed to load", result.Pos);
            return;
        }

        if (!result.IsRemeshOnly)
        {
            if (result.Chunk == null) { Log.Error("Load result for {Pos} had null chunk", result.Pos); return; }
            this._world.SetChunk(result.Pos, result.Chunk);
        }

        if (this._renderers.TryGetValue(result.Pos, out IChunkRenderer? old))
        {
            old.Dispose();
            this._renderers.Remove(result.Pos);
        }
        // Skip renderer allocation for empty meshes (pure-air chunks above
        // terrain, or chunks whose every face culls). The chunk still counts
        // as ready — DrawVisible's iteration over _renderers naturally no-ops.
        if (result.Mesh.Length > 0)
        {
            this._renderers[result.Pos] = this._rendererFactory(result.Mesh);
        }
        this._ready.Add(result.Pos);
        this._dirty.Remove(result.Pos);

        if (!result.IsRemeshOnly)
        {
            this.MarkCardinalNeighborsDirty(result.Pos);
            Log.Debug("Loaded {Pos}", result.Pos);
        }
        else
        {
            Log.Debug("Remeshed {Pos}", result.Pos);
        }
    }

    private void MarkCardinalNeighborsDirty(Vector3i pos)
    {
        Vector3i[] offsets =
        {
            new(-1, 0, 0), new(1, 0, 0),
            new(0, -1, 0), new(0, 1, 0),
            new(0, 0, -1), new(0, 0, 1),
        };
        foreach (Vector3i off in offsets)
        {
            Vector3i neighbor = pos + off;
            if (this._ready.Contains(neighbor)) this._dirty.Add(neighbor);
        }
    }

    private void DispatchJobs()
    {
        // Dirty re-meshes first (correctness), then new loads (player view).
        while (this._inFlight < this._cfg.MaxGenerateJobsInFlight && this._dirty.Count > 0)
        {
            Vector3i pos = this._dirty.First();
            this._dirty.Remove(pos);
            this.DispatchRemesh(pos);
        }

        while (this._inFlight < this._cfg.MaxGenerateJobsInFlight && this._pendingLoad.Count > 0)
        {
            Vector3i pos = this._pendingLoad.Dequeue();
            this.DispatchLoad(pos);
        }
    }

    private void DispatchRemesh(Vector3i pos)
    {
        VoxelChunk? chunk = this._world.GetChunk(pos);
        if (chunk == null) return;  // got unloaded between MarkDirty and now

        NeighborSnapshots neighbors = this.SnapshotNeighbors(pos);
        VoxelChunk snapshotChunk = CopyChunk(chunk);  // worker reads from copy

        this._inFlight++;

        IChunkMesher mesher = this._mesher;
        BlockRegistry registry = this._registry;
        ConcurrentQueue<ChunkJobResult> completed = this._completed;

        this._executor.Submit(() =>
        {
            try
            {
                float[] mesh = mesher.Build(snapshotChunk, registry, (nx, ny, nz) =>
                    neighbors.IsOpaqueOutside(nx, ny, nz, id => registry.Get(id).IsOpaque));
                completed.Enqueue(new ChunkJobResult(pos, snapshotChunk, mesh, Error: null, IsRemeshOnly: true));
            }
            catch (Exception ex)
            {
                completed.Enqueue(new ChunkJobResult(pos, Chunk: null, Mesh: null, Error: ex, IsRemeshOnly: true));
            }
        });
    }

    private static VoxelChunk CopyChunk(VoxelChunk source)
    {
        VoxelChunk copy = new();
        for (int x = 0; x < VoxelChunk.Size; x++)
        for (int y = 0; y < VoxelChunk.Size; y++)
        for (int z = 0; z < VoxelChunk.Size; z++)
        {
            copy[x, y, z] = source[x, y, z];
        }
        return copy;
    }

    private void DispatchLoad(Vector3i pos)
    {
        NeighborSnapshots neighbors = this.SnapshotNeighbors(pos);
        this._world.ModificationStore.TryGetDiff(pos, out IReadOnlyDictionary<int, ushort>? diff);

        ChunkJob job = new(pos, neighbors, diff);
        this._inFlight++;

        IChunkMesher mesher = this._mesher;
        BlockRegistry registry = this._registry;
        Func<Vector3i, VoxelChunk> generate = this._generate;
        ConcurrentQueue<ChunkJobResult> completed = this._completed;

        this._executor.Submit(() =>
        {
            try
            {
                VoxelChunk chunk = generate(job.Pos);
                if (job.Diff != null)
                {
                    foreach ((int idx, ushort id) in job.Diff)
                    {
                        int lx = idx / (VoxelChunk.Size * VoxelChunk.Size);
                        int ly = (idx / VoxelChunk.Size) % VoxelChunk.Size;
                        int lz = idx % VoxelChunk.Size;
                        chunk[lx, ly, lz] = id;
                    }
                }

                float[] mesh = mesher.Build(chunk, registry, (nx, ny, nz) =>
                    job.Neighbors.IsOpaqueOutside(nx, ny, nz, id => registry.Get(id).IsOpaque));

                completed.Enqueue(new ChunkJobResult(job.Pos, chunk, mesh, Error: null, IsRemeshOnly: false));
            }
            catch (Exception ex)
            {
                completed.Enqueue(new ChunkJobResult(job.Pos, Chunk: null, Mesh: null, Error: ex, IsRemeshOnly: false));
            }
        });
    }

    private NeighborSnapshots SnapshotNeighbors(Vector3i pos)
    {
        ushort[]? Snap(Vector3i p)
        {
            VoxelChunk? c = this._world.GetChunk(p);
            if (c == null) return null;
            ushort[] copy = new ushort[VoxelChunk.Size * VoxelChunk.Size * VoxelChunk.Size];
            for (int x = 0; x < VoxelChunk.Size; x++)
            for (int y = 0; y < VoxelChunk.Size; y++)
            for (int z = 0; z < VoxelChunk.Size; z++)
            {
                copy[(x * VoxelChunk.Size + y) * VoxelChunk.Size + z] = c[x, y, z];
            }
            return copy;
        }

        return new NeighborSnapshots(
            negX: Snap(pos + new Vector3i(-1, 0, 0)),
            posX: Snap(pos + new Vector3i(1, 0, 0)),
            negY: Snap(pos + new Vector3i(0, -1, 0)),
            posY: Snap(pos + new Vector3i(0, 1, 0)),
            negZ: Snap(pos + new Vector3i(0, 0, -1)),
            posZ: Snap(pos + new Vector3i(0, 0, 1)));
    }
}