acdream/src/AcDream.App/Rendering/Wb/EntitySpawnAdapter.cs

using System;
using System.Collections.Generic;
using AcDream.Core.Physics;
using AcDream.Core.World;

namespace AcDream.App.Rendering.Wb;

/// <summary>
/// The exact owner is still inside a reference transition, so its requested
/// logical removal has been queued and must be retried by the live-entity
/// teardown owner after the transition unwinds.
/// </summary>
public sealed class EntityPresentationRemovalDeferredException(uint serverGuid)
    : InvalidOperationException(
        $"Live entity 0x{serverGuid:X8} presentation removal is deferred until its active reference transition completes.");

/// <summary>
/// Routes server-spawned (<c>CreateObject</c>) entities through the
/// per-instance rendering path. Server entities always carry per-instance
/// customizations (palette overrides, texture changes, part swaps). Shared
/// surfaces use the WB atlas while per-instance texture composites are owned
/// by the live entity and released with it.
///
/// <para>
/// Companion to <see cref="LandblockSpawnAdapter"/>: that adapter handles
/// atlas-tier (procedural) entities; this one handles per-instance-tier
/// (server-spawned). The boundary is <c>ServerGuid != 0</c> on
/// <see cref="WorldEntity"/>.
/// </para>
///
/// <para>
/// <b>Sequencer factory</b>: the adapter is constructed with a
/// <c>Func&lt;WorldEntity, AnimationSequencer&gt;</c> factory so tests can
/// inject a stub without needing a live DatCollection or MotionTable.
/// Production callers supply a factory that fetches MotionTable from dats.
/// </para>
///
/// <para>
/// <b>Adjustment 6</b> (resolved Adjustment 4): <see cref="WorldEntity"/> now
/// carries <see cref="WorldEntity.PartOverrides"/> and
/// <see cref="WorldEntity.HiddenPartsMask"/>. <see cref="OnCreate"/> applies
/// both to the created <see cref="AnimatedEntityState"/>.
/// </para>
/// </summary>
public sealed class EntitySpawnAdapter
{
    private readonly IEntityTextureLifetime _textureLifetime;
    private readonly Func<WorldEntity, AnimationSequencer> _sequencerFactory;
    private readonly IWbMeshAdapter? _meshAdapter;

    // One logical owner per server GUID. Animated state survives projection
    // suspension, while the resident bit controls the shorter GPU-presentation
    // lifetime. The exact WorldEntity reference makes delayed visibility edges
    // from a displaced GUID generation harmless.
    // Single-threaded: called only from the render thread (same as GpuWorldState).
    private readonly Dictionary<uint, Owner> _ownersByGuid = new();

    private sealed class Owner(
        WorldEntity entity,
        AnimatedEntityState state,
        HashSet<ulong> meshIds)
    {
        public WorldEntity Entity { get; } = entity;
        public AnimatedEntityState State { get; } = state;
        public HashSet<ulong> MeshIds { get; set; } = meshIds;
        public HashSet<ulong> MeshReferencesHeld { get; } = new();
        public bool IsPresentationResident { get; set; }
        public bool TextureReleaseRequired { get; set; }
        public PresentationTransition Transition { get; set; }
        public bool RemovalPending { get; set; }

        public bool HasPresentationResources
        {
            get
            {
                if (TextureReleaseRequired)
                    return true;

                return MeshReferencesHeld.Count != 0;
            }
        }

        public bool IsFullyResident
        {
            get
            {
                if (!IsPresentationResident || !TextureReleaseRequired)
                    return false;
                return MeshReferencesHeld.SetEquals(MeshIds);
            }
        }

        public bool IsFullySuspended =>
            !IsPresentationResident && !HasPresentationResources;
    }

    private enum PresentationTransition
    {
        None,
        Resuming,
        Suspending,
        ChangingAppearance,
    }

    /// <param name="textureLifetime">
    /// Per-entity texture lifetime owner. Production uses
    /// <see cref="TextureCache"/>; tests use a recording implementation.
    /// </param>
    /// <param name="sequencerFactory">
    /// Factory that builds an <see cref="AnimationSequencer"/> for a given
    /// entity. Receives the full <see cref="WorldEntity"/> so it can look up
    /// the Setup + MotionTable from the entity's <c>SourceGfxObjOrSetupId</c>
    /// and server-supplied motion table override. Tests pass a lambda that
    /// returns a stub sequencer.
    /// </param>
    /// <param name="meshAdapter">
    /// Optional WB mesh adapter. When non-null, presentation residency
    /// registers each unique <c>MeshRef.GfxObjId</c> so WB background-loads
    /// the mesh data. Projection suspension or logical removal balances those
    /// references. When null, the adapter only tracks per-instance state.
    /// </param>
    public EntitySpawnAdapter(
        IEntityTextureLifetime textureLifetime,
        Func<WorldEntity, AnimationSequencer> sequencerFactory,
        IWbMeshAdapter? meshAdapter = null)
    {
        ArgumentNullException.ThrowIfNull(textureLifetime);
        ArgumentNullException.ThrowIfNull(sequencerFactory);
        _textureLifetime = textureLifetime;
        _sequencerFactory = sequencerFactory;
        _meshAdapter = meshAdapter;
    }

    /// <summary>
    /// Process a server-spawned entity. Returns the created
    /// <see cref="AnimatedEntityState"/> for the entity, or <c>null</c> if
    /// <paramref name="entity"/> is atlas-tier (<c>ServerGuid == 0</c>).
    /// </summary>
    public AnimatedEntityState? OnCreate(WorldEntity entity)
    {
        ArgumentNullException.ThrowIfNull(entity);

        // Atlas-tier entities (procedural / dat-hydrated, ServerGuid == 0)
        // are handled by LandblockSpawnAdapter, not here.
        if (entity.ServerGuid == 0) return null;

        // A.5 T18: populate cached AABB so WalkEntities reads from the cache
        // rather than recomputing Position±5 per frame. Called here because
        // all entity-state initialization (position, rotation) is complete
        // by this point via the WorldEntity passed in.
        entity.RefreshAabb();

        // Build the per-entity AnimatedEntityState. The sequencer factory
        // may return a stub (in tests) or a fully-constructed sequencer from
        // the MotionTable (in production). Factory must not return null —
        // if the entity has no motion table the factory should construct a
        // no-op sequencer (Setup + empty MotionTable + NullAnimationLoader).
        var sequencer = _sequencerFactory(entity);
        var state = new AnimatedEntityState(sequencer);

        // Adjustment 6: WorldEntity now carries PartOverrides + HiddenPartsMask.
        state.HideParts(entity.HiddenPartsMask);
        foreach (var po in entity.PartOverrides)
            state.SetPartOverride(po.PartIndex, po.GfxObjId);

        HashSet<ulong> meshIds = _meshAdapter is null
            ? []
            : CollectMeshIds(entity.MeshRefs, entity.PartOverrides);

        // Snapshot each unique GfxObj id for the shorter presentation lifetime.
        // Includes both the entity's natural MeshRefs AND any server-sent
        // PartOverride GfxObjs (weapons, clothing, helmets) — those replace the
        // Setup default and need their own mesh data uploaded.
        // Construct the replacement completely before displacing a live owner.
        // Sequencer/appearance construction is allowed to fail; in that case
        // the prior GUID incarnation and its presentation references remain
        // valid. Retirement is also completed before replacement publication:
        // a failed texture or mesh release therefore leaves the prior owner in
        // the dictionary, with its per-resource progress available to retry.
        var replacementOwner = new Owner(entity, state, meshIds);
        if (_ownersByGuid.TryGetValue(entity.ServerGuid, out Owner? displacedOwner))
        {
            if (displacedOwner.RemovalPending)
            {
                throw new EntityPresentationRemovalDeferredException(entity.ServerGuid);
            }

            if (displacedOwner.Transition != PresentationTransition.None)
            {
                throw new InvalidOperationException(
                    $"Live entity 0x{entity.ServerGuid:X8} replacement was requested while its presentation transition was already in progress.");
            }

            if (displacedOwner.HasPresentationResources)
            {
                if (!SuspendPresentation(displacedOwner))
                {
                    throw new InvalidOperationException(
                        $"Live entity 0x{entity.ServerGuid:X8} replacement was requested while its presentation teardown was already in progress.");
                }
            }

            _ownersByGuid[entity.ServerGuid] = replacementOwner;
        }
        else
        {
            _ownersByGuid.Add(entity.ServerGuid, replacementOwner);
        }

        return state;
    }

    /// <summary>
    /// Changes only the shorter presentation/GPU lifetime of an already-created
    /// live entity. A visible projection acquires WB mesh references; suspension
    /// releases those references and every owner-scoped composite texture. The
    /// animated state remains registered and no create-time scripts are replayed.
    /// Duplicate edges and edges for an older incarnation of a reused server GUID
    /// are ignored. A failed release keeps the published resident state and the
    /// exact unfinished resources so the same edge can be retried safely.
    /// </summary>
    /// <returns><c>true</c> when this call applied a residency edge.</returns>
    public bool SetPresentationResident(WorldEntity entity, bool resident)
    {
        ArgumentNullException.ThrowIfNull(entity);
        if (!_ownersByGuid.TryGetValue(entity.ServerGuid, out Owner? owner)
            || !ReferenceEquals(owner.Entity, entity)
            || owner.RemovalPending
            || (resident ? owner.IsFullyResident : owner.IsFullySuspended))
        {
            return false;
        }

        return resident
            ? ResumePresentation(owner)
            : SuspendPresentation(owner);
    }

    /// <summary>
    /// Reconciles the mesh-reference set for an in-place retail
    /// <c>SmartBox::UpdateVisualDesc</c> mutation. Added meshes are acquired
    /// before <paramref name="publishAppearance"/> makes the new appearance
    /// visible. Only then is the exact new mesh set published and superseded
    /// references released. A failed release remains represented by
    /// <see cref="Owner.MeshReferencesHeld"/> and is retried by the next
    /// residency edge or logical teardown.
    /// <paramref name="afterPublication"/> runs after that commit point but
    /// before retirement, so dependent pose/cache publication cannot leave the
    /// new entity appearance paired with rolled-back mesh ownership.
    /// </summary>
    /// <remarks>
    /// The exact <see cref="WorldEntity"/> reference is the incarnation token.
    /// A delayed appearance callback from a GUID-reused owner is ignored. The
    /// method is render-thread only; a re-entrant request is rejected before
    /// its publication callback can mutate the active owner.
    /// </remarks>
    /// <returns>
    /// <c>true</c> when the matching owner's appearance was published;
    /// otherwise <c>false</c> for a stale, removing, or re-entrant owner.
    /// </returns>
    public bool OnAppearanceChanged(
        WorldEntity entity,
        IReadOnlyList<MeshRef> meshRefs,
        IReadOnlyList<PartOverride> partOverrides,
        Action publishAppearance,
        Action? afterPublication = null)
    {
        ArgumentNullException.ThrowIfNull(entity);
        ArgumentNullException.ThrowIfNull(meshRefs);
        ArgumentNullException.ThrowIfNull(partOverrides);
        ArgumentNullException.ThrowIfNull(publishAppearance);

        if (!_ownersByGuid.TryGetValue(entity.ServerGuid, out Owner? owner)
            || !ReferenceEquals(owner.Entity, entity)
            || owner.RemovalPending
            || owner.Transition != PresentationTransition.None)
        {
            return false;
        }

        HashSet<ulong> nextMeshIds = _meshAdapter is null
            ? []
            : CollectMeshIds(meshRefs, partOverrides);
        owner.Transition = PresentationTransition.ChangingAppearance;
        var acquiredThisTransition = new List<ulong>();
        bool meshSetPublished = false;

        try
        {
            if (owner.IsPresentationResident && _meshAdapter is not null)
                AcquireMissingMeshReferences(owner, nextMeshIds, acquiredThisTransition);

            publishAppearance();

            // Publication is the commit point: every desired resident mesh is
            // owned before this assignment. Superseded references may remain
            // temporarily held if their release reports a retryable failure,
            // but they are no longer part of the published appearance set.
            owner.MeshIds = nextMeshIds;
            meshSetPublished = true;
            afterPublication?.Invoke();

            List<Exception>? releaseFailures = owner.IsPresentationResident
                ? ReleaseMeshReferencesOutside(owner, nextMeshIds)
                : ReleaseAllMeshReferences(owner);
            if (releaseFailures is not null)
            {
                throw new AggregateException(
                    $"Live entity 0x{owner.Entity.ServerGuid:X8} appearance mesh retirement failed.",
                    releaseFailures);
            }

            return true;
        }
        catch (Exception acquireOrPublicationFailure) when (!meshSetPublished)
        {
            // Acquisition failed before publication. Preserve the prior exact
            // mesh set and undo only references acquired by this transition.
            List<Exception>? rollbackFailures = RollBackAcquiredMeshReferences(
                owner,
                acquiredThisTransition);
            if (rollbackFailures is null)
                throw;

            rollbackFailures.Insert(0, acquireOrPublicationFailure);
            throw new AggregateException(
                $"Live entity 0x{owner.Entity.ServerGuid:X8} appearance publication and mesh rollback failed.",
                rollbackFailures);
        }
        finally
        {
            owner.Transition = PresentationTransition.None;
        }
    }

    /// <summary>
    /// Release the per-entity state for <paramref name="serverGuid"/>. Called
    /// on <c>RemoveObject</c>. Unknown guids (never spawned, or already
    /// removed) are silently ignored. Cleanup failure leaves the owner registered
    /// and propagates the exception; a later call resumes its unfinished releases.
    /// </summary>
    public void OnRemove(uint serverGuid)
        => _ = TryRemove(serverGuid, expectedEntity: null);

    private bool TryRemove(uint serverGuid, WorldEntity? expectedEntity)
    {
        if (!_ownersByGuid.TryGetValue(serverGuid, out Owner? owner)
            || (expectedEntity is not null && !ReferenceEquals(owner.Entity, expectedEntity)))
        {
            return false;
        }

        owner.RemovalPending = true;
        if (owner.Transition != PresentationTransition.None)
            throw new EntityPresentationRemovalDeferredException(serverGuid);

        // A resident owner still holds both mesh and potentially-lazy texture
        // resources. A suspended owner released them at the visibility edge,
        // so logical removal must not decrement or release a second time. Do
        // not remove the dictionary entry until cleanup succeeds: otherwise a
        // release exception would orphan its remaining references and make a
        // later OnRemove retry impossible.
        if (owner.HasPresentationResources && !SuspendPresentation(owner))
            return false;

        if (_ownersByGuid.TryGetValue(serverGuid, out Owner? current)
            && ReferenceEquals(current, owner))
        {
            _ownersByGuid.Remove(serverGuid);
            return true;
        }


        return false;
    }

    /// <summary>
    /// Releases a logical owner only when <paramref name="entity"/> is the exact
    /// incarnation currently registered for its server GUID. This is the live
    /// runtime teardown entry point: a delayed callback from an older generation
    /// cannot remove a replacement that reused the same GUID.
    /// </summary>
    /// <returns><c>true</c> when the matching owner was removed.</returns>
    public bool OnRemove(WorldEntity entity)
    {
        ArgumentNullException.ThrowIfNull(entity);
        return TryRemove(entity.ServerGuid, entity);
    }

    /// <summary>
    /// Look up the <see cref="AnimatedEntityState"/> for a server guid.
    /// Returns <c>null</c> if the entity was never spawned or has already
    /// been removed.
    /// </summary>
    public AnimatedEntityState? GetState(uint serverGuid)
        => _ownersByGuid.TryGetValue(serverGuid, out Owner? owner)
            ? owner.State
            : null;

    private bool ResumePresentation(Owner owner)
    {
        if (owner.Transition != PresentationTransition.None)
            return false;

        owner.Transition = PresentationTransition.Resuming;
        var acquiredThisTransition = new List<ulong>();
        bool desiredMeshSetAcquired = false;
        try
        {
            if (_meshAdapter is not null)
                AcquireMissingMeshReferences(owner, owner.MeshIds, acquiredThisTransition);
            desiredMeshSetAcquired = true;

            owner.TextureReleaseRequired = true;
            owner.IsPresentationResident = true;

            List<Exception>? releaseFailures = ReleaseMeshReferencesOutside(
                owner,
                owner.MeshIds);
            if (releaseFailures is not null)
            {
                throw new AggregateException(
                    $"Live entity 0x{owner.Entity.ServerGuid:X8} presentation mesh reconciliation failed.",
                    releaseFailures);
            }

            return true;
        }
        catch (Exception acquireFailure) when (!desiredMeshSetAcquired)
        {
            List<Exception>? rollbackFailures = RollBackAcquiredMeshReferences(
                owner,
                acquiredThisTransition);
            if (rollbackFailures is null)
                throw;

            rollbackFailures.Insert(0, acquireFailure);
            throw new AggregateException(
                $"Live entity 0x{owner.Entity.ServerGuid:X8} presentation resume and rollback failed.",
                rollbackFailures);
        }
        finally
        {
            owner.Transition = PresentationTransition.None;
        }
    }

    private bool SuspendPresentation(Owner owner)
    {
        if (owner.Transition != PresentationTransition.None)
            return false;

        // IsPresentationResident deliberately remains true until every release
        // succeeds. Transition prevents a re-entrant duplicate edge from
        // releasing the same resource twice, while the completion markers keep
        // partial progress retryable after an exception.
        owner.Transition = PresentationTransition.Suspending;

        List<Exception>? failures = null;
        if (owner.TextureReleaseRequired)
        {
            try
            {
                _textureLifetime.ReleaseOwner(owner.Entity.Id);
                owner.TextureReleaseRequired = false;
            }
            catch (Exception error)
            {
                (failures ??= new List<Exception>()).Add(error);
            }
        }

        List<Exception>? meshFailures = ReleaseAllMeshReferences(owner);
        if (meshFailures is not null)
            (failures ??= new List<Exception>()).AddRange(meshFailures);

        if (failures is not null)
        {
            owner.Transition = PresentationTransition.None;
            throw new AggregateException(
                $"Live entity 0x{owner.Entity.ServerGuid:X8} presentation suspension failed.",
                failures);
        }

        owner.IsPresentationResident = false;
        owner.Transition = PresentationTransition.None;
        return true;
    }

    private static HashSet<ulong> CollectMeshIds(
        IReadOnlyList<MeshRef> meshRefs,
        IReadOnlyList<PartOverride> partOverrides)
    {
        var unique = new HashSet<ulong>();
        for (int i = 0; i < meshRefs.Count; i++)
            unique.Add(meshRefs[i].GfxObjId);
        for (int i = 0; i < partOverrides.Count; i++)
            unique.Add(partOverrides[i].GfxObjId);
        return unique;
    }

    private void AcquireMissingMeshReferences(
        Owner owner,
        HashSet<ulong> desiredMeshIds,
        List<ulong> acquiredThisTransition)
    {
        if (_meshAdapter is null)
            return;

        foreach (ulong meshId in desiredMeshIds)
        {
            if (owner.MeshReferencesHeld.Contains(meshId))
                continue;

            try
            {
                _meshAdapter.IncrementRefCount(meshId);
                owner.MeshReferencesHeld.Add(meshId);
                acquiredThisTransition.Add(meshId);
            }
            catch (MeshReferenceMutationException error)
            {
                if (error.MutationCommitted)
                {
                    owner.MeshReferencesHeld.Add(meshId);
                    acquiredThisTransition.Add(meshId);
                }

                throw;
            }
        }
    }

    private List<Exception>? RollBackAcquiredMeshReferences(
        Owner owner,
        List<ulong> acquiredThisTransition)
    {
        if (_meshAdapter is null)
            return null;

        List<Exception>? failures = null;
        for (int i = acquiredThisTransition.Count - 1; i >= 0; i--)
        {
            ulong meshId = acquiredThisTransition[i];
            if (!owner.MeshReferencesHeld.Contains(meshId))
                continue;

            try
            {
                _meshAdapter.DecrementRefCount(meshId);
                owner.MeshReferencesHeld.Remove(meshId);
            }
            catch (Exception error)
            {
                if (error is MeshReferenceMutationException { MutationCommitted: true })
                    owner.MeshReferencesHeld.Remove(meshId);
                (failures ??= new List<Exception>()).Add(error);
            }
        }

        return failures;
    }

    private List<Exception>? ReleaseMeshReferencesOutside(
        Owner owner,
        HashSet<ulong> desiredMeshIds)
    {
        if (_meshAdapter is null || owner.MeshReferencesHeld.Count == 0)
            return null;

        List<Exception>? failures = null;
        ulong[] heldSnapshot = [.. owner.MeshReferencesHeld];
        foreach (ulong meshId in heldSnapshot)
        {
            if (desiredMeshIds.Contains(meshId))
                continue;

            try
            {
                _meshAdapter.DecrementRefCount(meshId);
                owner.MeshReferencesHeld.Remove(meshId);
            }
            catch (Exception error)
            {
                if (error is MeshReferenceMutationException { MutationCommitted: true })
                    owner.MeshReferencesHeld.Remove(meshId);
                (failures ??= new List<Exception>()).Add(error);
            }
        }

        return failures;
    }

    private List<Exception>? ReleaseAllMeshReferences(Owner owner)
    {
        if (_meshAdapter is null || owner.MeshReferencesHeld.Count == 0)
            return null;

        return ReleaseMeshReferencesOutside(owner, []);
    }
}