acdream/src/AcDream.App/Streaming/GpuWorldState.cs

using System.Collections.Generic;
using System.Linq;
using System.Numerics;
using AcDream.App.Rendering.Wb;
using AcDream.Core.World;

namespace AcDream.App.Streaming;

/// <summary>
/// Render-thread-owned registry of currently-loaded landblocks and their
/// entities. All mutation happens in <see cref="StreamingController.Tick"/>
/// on the render thread; the renderer reads <see cref="Entities"/> once per
/// frame.
///
/// <para>
/// Replaces the pre-streaming flat <c>_entities</c> list. This class is the
/// single point of truth for "what's in the world right now" and the only
/// thing that mutates it.
/// </para>
///
/// <para>
/// <b>Pending live entities.</b> Live <c>CreateObject</c> spawns can race
/// against streaming: the server may send a spawn for landblock X before
/// X is loaded into <see cref="_loaded"/> (frequently true on the first
/// frame after login, where the entire post-login spawn flood drains
/// before the streaming controller has finished loading the visible
/// window). To survive this race, <see cref="AppendLiveEntity"/> stores
/// orphaned spawns in a per-landblock pending bucket. When
/// <see cref="AddLandblock"/> later loads the landblock, the matching
/// pending entries are merged into the loaded record before the flat
/// view rebuild. <see cref="RemoveLandblock"/> drops pending entries for
/// the same landblock — if the landblock just left the visible window,
/// any spawns that came with it are no longer relevant.
/// </para>
///
/// <remarks>
/// Threading: not thread-safe. All calls must happen on the render thread.
/// </remarks>
/// </summary>
public sealed class GpuWorldState
{
    private readonly LandblockSpawnAdapter? _wbSpawnAdapter;
    private readonly EntitySpawnAdapter? _wbEntitySpawnAdapter;

    public GpuWorldState(
        LandblockSpawnAdapter? wbSpawnAdapter = null,
        EntitySpawnAdapter? wbEntitySpawnAdapter = null)
    {
        _wbSpawnAdapter = wbSpawnAdapter;
        _wbEntitySpawnAdapter = wbEntitySpawnAdapter;
    }

    private readonly Dictionary<uint, LoadedLandblock> _loaded = new();
    private readonly Dictionary<uint, (Vector3 Min, Vector3 Max)> _aabbs = new();

    /// <summary>
    /// Per-landblock buffer of live entities awaiting their landblock's
    /// arrival. Keyed by canonical landblock id (<c>0xAAAA0xFFFF</c>).
    /// Drained into <see cref="_loaded"/> in <see cref="AddLandblock"/>.
    /// </summary>
    private readonly Dictionary<uint, List<WorldEntity>> _pendingByLandblock = new();

    /// <summary>
    /// Entities that must survive landblock unloads (e.g. the player character).
    /// On RemoveLandblock, these are rescued and re-parked as pending for their
    /// current canonical landblock.
    /// </summary>
    private readonly HashSet<uint> _persistentGuids = new();

    // Cached flat view over all entities across all loaded landblocks,
    // rebuilt on each add/remove. The renderer holds a reference to this
    // list, so rebuilding it replaces the reference atomically.
    private IReadOnlyList<WorldEntity> _flatEntities = System.Array.Empty<WorldEntity>();

    public IReadOnlyList<WorldEntity> Entities => _flatEntities;
    public IReadOnlyCollection<uint> LoadedLandblockIds => _loaded.Keys;

    public bool IsLoaded(uint landblockId) => _loaded.ContainsKey(landblockId);

    /// <summary>
    /// Try to grab the loaded record for a landblock — useful for callers
    /// that need to enumerate entities before the landblock is dropped
    /// (e.g. unregistering dynamic lights on a RemoveLandblock).
    /// </summary>
    public bool TryGetLandblock(uint landblockId, out LoadedLandblock? lb)
    {
        if (_loaded.TryGetValue(landblockId, out var found))
        {
            lb = found;
            return true;
        }
        lb = null;
        return false;
    }

    /// <summary>
    /// Store the axis-aligned bounding box for a loaded landblock. Called from
    /// the render thread after the terrain mesh is built and uploaded.
    /// </summary>
    public void SetLandblockAabb(uint landblockId, Vector3 min, Vector3 max)
    {
        _aabbs[landblockId] = (min, max);
    }

    /// <summary>
    /// Per-landblock iteration with AABB data for use by the frustum-culling
    /// draw path. Landblocks without a stored AABB yield <see cref="Vector3.Zero"/>
    /// for both corners, which the culler will conservatively treat as visible.
    ///
    /// <para>
    /// A.5 T17: also yields an <c>AnimatedById</c> dictionary built on the fly
    /// from the landblock's entity list. This lets <see cref="WbDrawDispatcher"/>
    /// skip the full entity walk when the landblock is frustum-culled but animated
    /// entities inside it must still be processed (Change #1).
    /// Building the dict per-yield is cheap (~132 entities/LB max). A caching
    /// layer is out of A.5 scope.
    /// </para>
    /// </summary>
    public IEnumerable<(uint LandblockId, Vector3 AabbMin, Vector3 AabbMax,
                        IReadOnlyList<WorldEntity> Entities,
                        IReadOnlyDictionary<uint, WorldEntity>? AnimatedById)> LandblockEntries
    {
        get
        {
            foreach (var kvp in _loaded)
            {
                // Build AnimatedById on the fly — cheap (~132 entities/LB max).
                var byId = new Dictionary<uint, WorldEntity>(kvp.Value.Entities.Count);
                foreach (var e in kvp.Value.Entities)
                    byId[e.Id] = e;

                if (_aabbs.TryGetValue(kvp.Key, out var aabb))
                    yield return (kvp.Key, aabb.Min, aabb.Max, kvp.Value.Entities, byId);
                else
                    yield return (kvp.Key, Vector3.Zero, Vector3.Zero, kvp.Value.Entities, byId);
            }
        }
    }

    /// <summary>
    /// Total live entities currently parked in the pending bucket waiting
    /// for their landblock to arrive. Useful diagnostic for verifying the
    /// pending path is doing its job.
    /// </summary>
    public int PendingLiveEntityCount => _pendingByLandblock.Values.Sum(list => list.Count);

    public void AddLandblock(LoadedLandblock landblock)
    {
        // If pending live entities have been waiting for this landblock,
        // merge them into the LoadedLandblock record before storing. The
        // record's Entities field is IReadOnlyList; we replace the whole
        // list rather than try to mutate in place.
        if (_pendingByLandblock.TryGetValue(landblock.LandblockId, out var pending) && pending.Count > 0)
        {
            var merged = new List<WorldEntity>(landblock.Entities.Count + pending.Count);
            merged.AddRange(landblock.Entities);
            merged.AddRange(pending);
            landblock = new LoadedLandblock(landblock.LandblockId, landblock.Heightmap, merged);
            _pendingByLandblock.Remove(landblock.LandblockId);
        }

        _loaded[landblock.LandblockId] = landblock;
        if (_wbSpawnAdapter is not null)
            _wbSpawnAdapter.OnLandblockLoaded(_loaded[landblock.LandblockId]);
        RebuildFlatView();
    }

    /// <summary>
    /// Mark a server-GUID as persistent — this entity survives landblock unloads
    /// and gets re-parked as pending for its current canonical landblock.
    /// </summary>
    public void MarkPersistent(uint serverGuid)
    {
        _persistentGuids.Add(serverGuid);
    }

    /// <summary>
    /// Move a persistent entity from its current landblock slot to a new one.
    /// Called every frame for the player entity so it stays in the landblock
    /// matching its actual position (not its spawn landblock). Without this,
    /// the entity stays in the spawn landblock and gets frustum-culled when
    /// the player walks away.
    /// </summary>
    public void RelocateEntity(WorldEntity entity, uint newCanonicalLb)
    {
        if (entity.ServerGuid == 0) return;

        // Remove from current landblock (find it by scanning)
        foreach (var kvp in _loaded)
        {
            var entities = kvp.Value.Entities;
            for (int i = 0; i < entities.Count; i++)
            {
                if (ReferenceEquals(entities[i], entity))
                {
                    if (kvp.Key == newCanonicalLb) return; // already in the right place

                    // Remove from old
                    var newList = new List<WorldEntity>(entities.Count - 1);
                    for (int j = 0; j < entities.Count; j++)
                        if (j != i) newList.Add(entities[j]);
                    _loaded[kvp.Key] = new LoadedLandblock(kvp.Value.LandblockId, kvp.Value.Heightmap, newList);

                    // Add to new (via AppendLiveEntity which handles pending)
                    AppendLiveEntity(newCanonicalLb, entity);
                    return;
                }
            }
        }
    }

    public void RemoveLandblock(uint landblockId)
    {
        if (_wbSpawnAdapter is not null)
            _wbSpawnAdapter.OnLandblockUnloaded(landblockId);

        // Rescue persistent entities before removal. These get appended
        // to the _persistentRescued list; the caller is responsible for
        // re-injecting them (via AppendLiveEntity) into whatever landblock
        // the player is currently on.
        if (_loaded.TryGetValue(landblockId, out var lb))
        {
            foreach (var entity in lb.Entities)
            {
                if (entity.ServerGuid != 0 && _persistentGuids.Contains(entity.ServerGuid))
                {
                    _persistentRescued.Add(entity);
                }
            }
        }

        _pendingByLandblock.Remove(landblockId);
        _aabbs.Remove(landblockId);

        if (_loaded.Remove(landblockId))
            RebuildFlatView();
    }

    private readonly List<WorldEntity> _persistentRescued = new();

    /// <summary>
    /// Drain entities rescued from unloaded landblocks. The caller should
    /// re-inject each via <see cref="AppendLiveEntity"/> with its current position.
    /// </summary>
    public List<WorldEntity> DrainRescued()
    {
        if (_persistentRescued.Count == 0) return _persistentRescued;
        var result = new List<WorldEntity>(_persistentRescued);
        _persistentRescued.Clear();
        return result;
    }

    /// <summary>
    /// Remove every entity with the given <paramref name="serverGuid"/> from
    /// all loaded landblocks AND all pending buckets, then rebuild the flat
    /// view. Used by the live <c>CreateObject</c> handler to de-duplicate
    /// when the server re-sends a spawn (visibility refresh, landblock
    /// crossing, etc.). Without this, multiple copies of the same NPC
    /// accumulate in the renderer, each with its own <c>PaletteOverride</c>
    /// and <c>MeshRefs</c> — producing "NPC clothing flickers as I turn the
    /// camera" because the depth test picks different duplicates frame-to-frame.
    ///
    /// Safe to call with a server guid that's not currently present — no-op.
    /// </summary>
    public void RemoveEntityByServerGuid(uint serverGuid)
    {
        if (serverGuid == 0) return;

        // Phase N.4 Task 17: release per-instance state for server-spawned
        // entities. No-op for atlas-tier entities (never registered).
        _wbEntitySpawnAdapter?.OnRemove(serverGuid);

        bool rebuiltLoaded = false;

        // Scan loaded landblocks. ToArray() so we can mutate _loaded inside.
        foreach (var kvp in _loaded.ToArray())
        {
            var lb = kvp.Value;
            int foundCount = 0;
            for (int i = 0; i < lb.Entities.Count; i++)
                if (lb.Entities[i].ServerGuid == serverGuid) foundCount++;

            if (foundCount == 0) continue;

            var newList = new List<WorldEntity>(lb.Entities.Count - foundCount);
            foreach (var e in lb.Entities)
                if (e.ServerGuid != serverGuid) newList.Add(e);

            _loaded[kvp.Key] = new LoadedLandblock(lb.LandblockId, lb.Heightmap, newList);
            rebuiltLoaded = true;
        }

        // Scrub pending buckets too — a duplicate CreateObject may arrive
        // while the landblock is still loading.
        foreach (var kvp in _pendingByLandblock)
            kvp.Value.RemoveAll(e => e.ServerGuid == serverGuid);

        if (rebuiltLoaded) RebuildFlatView();
    }

    /// <summary>
    /// Append an entity to a specific landblock's slot. Used by the live
    /// CreateObject path where the server spawns entities at a server-side
    /// position whose landblock may or may not be loaded yet.
    ///
    /// <para>
    /// The server's <c>landblockId</c> is in cell-resolved form
    /// (<c>0xAAAA00CC</c>: high byte X, second byte Y, low 16 bits cell
    /// index within the landblock). The streaming system stores landblocks
    /// keyed by their canonical <c>0xAAAA0xFFFF</c> form. Canonicalize
    /// on the way in so callers don't have to think about it.
    /// </para>
    ///
    /// <para>
    /// Outcome:
    /// <list type="bullet">
    ///   <item>If the landblock is already loaded, the entity is appended
    ///         to its <c>Entities</c> list and the flat view is rebuilt
    ///         immediately.</item>
    ///   <item>If the landblock is not yet loaded, the entity is parked
    ///         in <see cref="_pendingByLandblock"/> and will be merged
    ///         into the next <see cref="AddLandblock"/> for the same id.</item>
    /// </list>
    /// </para>
    /// </summary>
    public void AppendLiveEntity(uint landblockId, WorldEntity entity)
    {
        // Phase N.4 Task 17: route server-spawned entities through the
        // per-instance adapter. Atlas-tier entities (ServerGuid == 0) are
        // skipped by OnCreate — it returns null immediately for those.
        _wbEntitySpawnAdapter?.OnCreate(entity);

        uint canonicalLandblockId = (landblockId & 0xFFFF0000u) | 0xFFFFu;

        if (_loaded.TryGetValue(canonicalLandblockId, out var lb))
        {
            // Hot path — landblock is already loaded. Rebuild the record
            // with the new entity appended.
            var newEntities = new List<WorldEntity>(lb.Entities.Count + 1);
            newEntities.AddRange(lb.Entities);
            newEntities.Add(entity);
            _loaded[canonicalLandblockId] = new LoadedLandblock(lb.LandblockId, lb.Heightmap, newEntities);
            RebuildFlatView();
            return;
        }

        // Cold path — landblock not yet loaded. Park the entity in the
        // pending bucket; AddLandblock will pick it up when the streamer
        // delivers the matching landblock.
        if (!_pendingByLandblock.TryGetValue(canonicalLandblockId, out var bucket))
        {
            bucket = new List<WorldEntity>();
            _pendingByLandblock[canonicalLandblockId] = bucket;
        }
        bucket.Add(entity);
    }

    /// <summary>
    /// Drop all entities from a landblock without removing the terrain. Used
    /// by two-tier streaming when a landblock crosses Near→Far hysteresis.
    /// Per Phase A.5 spec §4.4.
    ///
    /// <para>
    /// <b>Persistent-entity rescue is intentionally omitted</b> (unlike
    /// <see cref="RemoveLandblock"/>): demote-tier entities are atlas-tier
    /// only (procedural scenery, dat-static stabs/buildings) — they never
    /// have <c>ServerGuid != 0</c> and so can never be in <see cref="_persistentGuids"/>.
    /// The local player and other live server-spawned entities live in their
    /// landblock via <c>RelocateEntity</c> per frame and are not affected
    /// by Near→Far demotion of dat-static landblock layers.
    /// </para>
    /// </summary>
    public void RemoveEntitiesFromLandblock(uint landblockId)
    {
        // A.5 T14 follow-up: canonicalize for symmetry with AppendLiveEntity.
        // Streaming callers always pass canonical (0xAAAA0xFFFF) ids; this
        // protects against future callers that mirror AppendLiveEntity's
        // cell-resolved-id pattern.
        uint canonical = (landblockId & 0xFFFF0000u) | 0xFFFFu;
        if (!_loaded.TryGetValue(canonical, out var lb)) return;
        if (_wbSpawnAdapter is not null)
            _wbSpawnAdapter.OnLandblockUnloaded(canonical);
        _loaded[canonical] = new LoadedLandblock(lb.LandblockId, lb.Heightmap, System.Array.Empty<WorldEntity>());
        _pendingByLandblock.Remove(canonical);
        RebuildFlatView();
    }

    /// <summary>
    /// Merge entities into an existing-loaded landblock. Used by two-tier
    /// streaming for the Far→Near promotion case (terrain already loaded;
    /// entity layer streaming in). Falls back to the pending bucket if the
    /// landblock isn't loaded yet (handles the rare "promote arrives before
    /// far load completes" race).
    /// Per Phase A.5 spec §4.4.
    ///
    /// <para>
    /// <b>Landblock id is canonicalized</b> (low 16 bits forced to 0xFFFF) —
    /// callers may pass cell-resolved ids and they will key correctly.
    /// </para>
    /// </summary>
    public void AddEntitiesToExistingLandblock(uint landblockId, IReadOnlyList<WorldEntity> entities)
    {
        // A.5 T14 follow-up: canonicalize for symmetry with AppendLiveEntity.
        uint canonical = (landblockId & 0xFFFF0000u) | 0xFFFFu;
        if (!_loaded.TryGetValue(canonical, out var lb))
        {
            // Park as pending — same pattern as AppendLiveEntity for not-yet-loaded LBs.
            if (!_pendingByLandblock.TryGetValue(canonical, out var bucket))
            {
                bucket = new List<WorldEntity>();
                _pendingByLandblock[canonical] = bucket;
            }
            bucket.AddRange(entities);
            return;
        }
        var merged = new List<WorldEntity>(lb.Entities.Count + entities.Count);
        merged.AddRange(lb.Entities);
        merged.AddRange(entities);
        _loaded[canonical] = new LoadedLandblock(lb.LandblockId, lb.Heightmap, merged);
        if (_wbSpawnAdapter is not null)
            _wbSpawnAdapter.OnLandblockLoaded(_loaded[canonical]);
        RebuildFlatView();
    }

    private void RebuildFlatView()
    {
        _flatEntities = _loaded.Values.SelectMany(lb => lb.Entities).ToArray();
    }
}