chameleon/Assets/Scripts/PlayerPaintSync.cs

using System.Collections;
using System.Collections.Generic;
using Unity.Netcode;
using UnityEngine;

/// <summary>
/// Replicates mesh painting across the network as brush-stamp commands (not pixels). The owner
/// captures every committed stamp from its <see cref="MeshPainter"/>, batches the frame's stamps and
/// forwards them; remote avatars replay each stamp through their own painter onto the identical
/// mesh/UVs. Position/animation come from the other sync components; this only carries the tiny paint
/// commands (~tens of bytes each), so two players see each other paint live without streaming the
/// 12 MB of paint maps.
///
/// Late join: a client that connects AFTER painting happened also gets the accumulated result. On
/// spawn it pulls a one-time compressed snapshot of the whole map set from the server (chunked,
/// throttled to respect the transport's reliable window), buffers any live stamps that arrive during
/// the transfer, and after applying the snapshot replays the buffered stamps newer than the snapshot's
/// cutoff sequence - so no stroke is lost or double-applied across the join window.
/// </summary>
[DisallowMultipleComponent]
public class PlayerPaintSync : NetworkBehaviour
{
    [Tooltip("Max brush stamps forwarded to the network per frame. Pathological fast strokes (hundreds " +
             "of interpolated sub-stamps in one frame) are decimated to this to bound the RPC size; local " +
             "painting stays full fidelity.")]
    [SerializeField] private int maxStampsPerFrame = 64;

    [Tooltip("Snapshot transfer chunk payload (bytes). Kept under the transport MTU so each chunk is a " +
             "single packet (no extra fragmentation).")]
    [SerializeField] private int snapshotChunkBytes = 1200;

    [Tooltip("Snapshot chunks sent per frame. Throttled so the burst stays under the reliable window " +
             "(~32 in-flight) while live stamps keep flowing.")]
    [SerializeField] private int snapshotChunksPerFrame = 8;

    private MeshPainter _painter;
    private readonly List<PaintStampNet> _pending = new List<PaintStampNet>();

    // Monotonic batch counter on the owner; on every peer, the highest batch sequence accounted for
    // (applied to a remote replica / sent by the owner). Used as the snapshot's cutoff marker.
    private int _seq;
    private int _seqHorizon;

    // --- Joining client: snapshot reassembly + live-stamp buffering until the baseline is applied ---
    private bool _awaitingSnapshot;          // buffer live stamps instead of applying them
    private byte[] _snapBuffer;              // reassembled compressed snapshot
    private int _snapCapturedSeq;            // batches <= this are already baked into the snapshot
    private bool _snapReadyToApply;          // last chunk arrived; apply once the painter is ready
    private readonly List<(int seq, PaintStampNet[] stamps)> _bufferedBatches = new List<(int, PaintStampNet[])>();

    private void Awake() => _painter = GetComponentInChildren<MeshPainter>(true);

    public override void OnNetworkSpawn()
    {
        if (!_painter)
            return;

        if (IsOwner)
        {
            _painter.OnStampCommitted += OnStampCommitted;
        }
        else
        {
            // Remote replica: never paint from this machine's mouse; paint only from replicated data.
            _painter.PaintingEnabled = false;
            // Don't persist a replicated avatar: all clones share one save key, so remotes would
            // clobber the local player's saved paint on teardown (same fix as the pose system).
            _painter.persistPaint = false;

            // The server (host) has been present since the start, so its replica already holds the full
            // paint via the live stream; only a genuine late joiner needs the baseline pulled over.
            if (!IsServer)
            {
                _awaitingSnapshot = true;
                RequestSnapshotRpc();
            }
        }
    }

    public override void OnNetworkDespawn()
    {
        if (_painter && IsOwner)
            _painter.OnStampCommitted -= OnStampCommitted;
    }

    // ---- Live stamp stream (owner -> server -> remotes) ----

    private void OnStampCommitted(PaintStampNet stamp) => _pending.Add(stamp);

    private void LateUpdate()
    {
        // Apply a fully-received snapshot once the painter has finished setup, then flush the stamps
        // that arrived during the transfer. Deferred to here so apply+replay happen atomically in order.
        if (_snapReadyToApply && _painter && _painter.IsPaintReady)
        {
            _snapReadyToApply = false;
            _painter.ApplyPaintSnapshot(_snapBuffer);
            FlushBufferedBatches();
            _snapBuffer = null;
            _awaitingSnapshot = false;
        }

        if (!IsOwner || _pending.Count == 0)
            return;

        int seq = ++_seq;
        _seqHorizon = seq;
        SubmitStampsServerRpc(seq, BuildBatch());
        _pending.Clear();
    }

    /// <summary>Returns this frame's stamps, uniformly decimated to maxStampsPerFrame if a fast stroke overran it.</summary>
    private PaintStampNet[] BuildBatch()
    {
        int count = _pending.Count;
        if (count <= maxStampsPerFrame)
            return _pending.ToArray();

        var outp = new PaintStampNet[maxStampsPerFrame];
        for (int i = 0; i < maxStampsPerFrame; i++)
            outp[i] = _pending[(int)((long)i * count / maxStampsPerFrame)];
        return outp;
    }

    [ServerRpc]
    private void SubmitStampsServerRpc(int seq, PaintStampNet[] stamps) => ApplyStampsClientRpc(seq, stamps);

    [ClientRpc]
    private void ApplyStampsClientRpc(int seq, PaintStampNet[] stamps)
    {
        // The owner already painted these locally; only remotes (incl. the host's replica) replay them.
        if (IsOwner || !_painter)
            return;

        _seqHorizon = seq; // server replica tracks this as the snapshot cutoff source

        if (_awaitingSnapshot)
        {
            // Defer until the baseline is applied; the cutoff sorts out which ones to keep.
            _bufferedBatches.Add((seq, stamps));
            return;
        }

        for (int i = 0; i < stamps.Length; i++)
            _painter.ApplyNetworkStamp(stamps[i]);
    }

    private void FlushBufferedBatches()
    {
        if (_painter)
        {
            for (int b = 0; b < _bufferedBatches.Count; b++)
            {
                var (seq, stamps) = _bufferedBatches[b];
                if (seq <= _snapCapturedSeq) continue; // already inside the snapshot
                for (int i = 0; i < stamps.Length; i++)
                    _painter.ApplyNetworkStamp(stamps[i]);
            }
        }
        _bufferedBatches.Clear();
    }

    // ---- Late-join snapshot transfer (server -> the one joining client) ----

    [Rpc(SendTo.Server, InvokePermission = RpcInvokePermission.Everyone)]
    private void RequestSnapshotRpc(RpcParams p = default)
    {
        if (!_painter || !_painter.TryCaptureSnapshot(out byte[] data))
            return;

        // Everything the captured maps already contain: the owner's send counter for our own avatar,
        // otherwise the last batch this server replica has applied.
        int capturedSeq = IsOwner ? _seq : _seqHorizon;
        StartCoroutine(SendSnapshot(p.Receive.SenderClientId, data, capturedSeq));
    }

    private IEnumerator SendSnapshot(ulong targetClientId, byte[] data, int capturedSeq)
    {
        var send = new ClientRpcParams { Send = new ClientRpcSendParams { TargetClientIds = new[] { targetClientId } } };

        BeginSnapshotClientRpc(data.Length, capturedSeq, send);

        int sentThisFrame = 0;
        for (int off = 0; off < data.Length; off += snapshotChunkBytes)
        {
            int len = Mathf.Min(snapshotChunkBytes, data.Length - off);
            var chunk = new byte[len];
            System.Array.Copy(data, off, chunk, 0, len);
            SnapshotChunkClientRpc(off, chunk, send);

            if (++sentThisFrame >= snapshotChunksPerFrame)
            {
                sentThisFrame = 0;
                yield return null; // spread the burst so the reliable window does not overflow
            }
        }

        EndSnapshotClientRpc(send);
    }

    [ClientRpc]
    private void BeginSnapshotClientRpc(int totalBytes, int capturedSeq, ClientRpcParams p = default)
    {
        _snapBuffer = new byte[totalBytes];
        _snapCapturedSeq = capturedSeq;
        _snapReadyToApply = false;
        _awaitingSnapshot = true;
    }

    [ClientRpc]
    private void SnapshotChunkClientRpc(int offset, byte[] chunk, ClientRpcParams p = default)
    {
        if (_snapBuffer == null) return;
        if (offset >= 0 && offset + chunk.Length <= _snapBuffer.Length)
            System.Array.Copy(chunk, 0, _snapBuffer, offset, chunk.Length);
    }

    [ClientRpc]
    private void EndSnapshotClientRpc(ClientRpcParams p = default)
    {
        // Apply happens in LateUpdate (once the painter is ready) so apply + buffered replay stay ordered.
        _snapReadyToApply = _snapBuffer != null;
        if (!_snapReadyToApply)
            _awaitingSnapshot = false;
    }
}