chameleon/Assets/Scripts/Components/MeshPainter/MeshPainter.Net.cs

using System.IO;
using System.IO.Compression;
using Unity.Netcode;
using UnityEngine;

/// <summary>
/// One committed brush stamp in replicable form. Painting is synced as commands, not pixels: the owner
/// already resolves every stamp into local object space (hit point, radius, tangent frame), so a remote
/// peer replays the SAME stamp through its own backend onto the identical mesh/UVs.
/// Self-contained on purpose: each stamp carries its full brush state, so the stream is stateless.
/// Channel data quantizes to bytes with no real loss (the paint maps are 8-bit anyway).
/// </summary>
public struct PaintStampNet : INetworkSerializable
{
    // Channel/mode bit flags (Channels field).
    public const byte FlagAlbedo = 1 << 0;
    public const byte FlagMetallicSmoothness = 1 << 1;
    public const byte FlagOcclusion = 1 << 2;
    public const byte FlagEmission = 1 << 3;
    public const byte FlagSpray = 1 << 4;

    // Placement (object-local space).
    public Vector3 P;       // hit point
    public float R;         // 3D brush radius
    public Vector3 U, V, N; // tangent frame (right / up / normal)

    public byte Channels;   // active channels + spray flag
    public int BrushIndex;  // active brush mask index (resolved-mask painting)
    public int SpraySeed;   // RNG seed so a spray stamp scatters identically on every peer

    public Color32 BrushColor;     // albedo color (alpha = opacity)
    public Color32 EmissionColor;  // emission color (alpha = opacity)
    public byte Metallic, Smoothness, Occlusion; // 0..255 = 0..1
    public byte MSOpacity, OccOpacity;           // 0..255 = 0..1

    public void NetworkSerialize<T>(BufferSerializer<T> s) where T : IReaderWriter
    {
        s.SerializeValue(ref P);
        s.SerializeValue(ref R);
        s.SerializeValue(ref U);
        s.SerializeValue(ref V);
        s.SerializeValue(ref N);
        s.SerializeValue(ref Channels);
        s.SerializeValue(ref BrushIndex);
        s.SerializeValue(ref SpraySeed);
        s.SerializeValue(ref BrushColor);
        s.SerializeValue(ref EmissionColor);
        s.SerializeValue(ref Metallic);
        s.SerializeValue(ref Smoothness);
        s.SerializeValue(ref Occlusion);
        s.SerializeValue(ref MSOpacity);
        s.SerializeValue(ref OccOpacity);
    }
}

public partial class MeshPainter
{
    /// <summary>
    /// Fires once per committed brush stamp on the owner (after the stamp has been drawn locally),
    /// carrying everything a remote peer needs to replay it. Only the owner's PlayerPaintSync
    /// subscribes; with no subscriber this is a no-op, so single-player keeps zero overhead.
    /// </summary>
    public event System.Action<PaintStampNet> OnStampCommitted;

    // Last spray seed used by the owner this stamp (captured into the replicated stamp).
    private int _lastSpraySeed;

    private static byte ToByte01(float v) => (byte)(Mathf.Clamp01(v) * 255f + 0.5f);

    /// <summary>True if anyone is listening for replicated stamps (the owner is networked + painting).</summary>
    private bool WantsStampCapture => OnStampCommitted != null;

    /// <summary>Builds the replicated stamp from the current brush state + resolved placement and raises the event.</summary>
    private void EmitStamp(Vector3 p, float r, Vector3 u, Vector3 v, Vector3 n)
    {
        var handler = OnStampCommitted;
        if (handler == null) return;

        byte ch = 0;
        if (paintAlbedo) ch |= PaintStampNet.FlagAlbedo;
        if (paintMetallicSmoothness) ch |= PaintStampNet.FlagMetallicSmoothness;
        if (paintOcclusion) ch |= PaintStampNet.FlagOcclusion;
        if (paintEmission) ch |= PaintStampNet.FlagEmission;
        if (spray) ch |= PaintStampNet.FlagSpray;

        handler(new PaintStampNet
        {
            P = p, R = r, U = u, V = v, N = n,
            Channels = ch,
            BrushIndex = activeBrushIndex,
            SpraySeed = _lastSpraySeed,
            BrushColor = brushColor,
            EmissionColor = emissionColor,
            Metallic = ToByte01(metallic),
            Smoothness = ToByte01(smoothness),
            Occlusion = ToByte01(occlusion),
            MSOpacity = ToByte01(metallicSmoothnessOpacity),
            OccOpacity = ToByte01(occlusionOpacity),
        });
    }

    /// <summary>
    /// Bakes the spray mask from a fixed RNG seed so every peer scatters the identical pattern. The
    /// global Unity RNG state is saved and restored so nothing else in the frame is perturbed.
    /// </summary>
    private void BuildSprayStampMaskSeeded(int seed)
    {
        var prev = Random.state;
        Random.InitState(seed);
        BuildSprayStampMask();
        Random.state = prev;
    }

    /// <summary>
    /// Replays a replicated brush stamp on a remote avatar: adopts the stamp's brush state, rebuilds
    /// the matching mask (resolved or seeded spray) and runs it through the active backend exactly like
    /// a local stamp. Mutating the public brush fields is safe here because the remote painter never
    /// paints from local input (PaintingEnabled is false on remotes).
    /// </summary>
    public void ApplyNetworkStamp(in PaintStampNet s)
    {
        if (!_setupValid) return;

        paintAlbedo = (s.Channels & PaintStampNet.FlagAlbedo) != 0;
        paintMetallicSmoothness = (s.Channels & PaintStampNet.FlagMetallicSmoothness) != 0;
        paintOcclusion = (s.Channels & PaintStampNet.FlagOcclusion) != 0;
        paintEmission = (s.Channels & PaintStampNet.FlagEmission) != 0;
        spray = (s.Channels & PaintStampNet.FlagSpray) != 0;

        brushColor = s.BrushColor;
        emissionColor = s.EmissionColor;
        metallic = s.Metallic / 255f;
        smoothness = s.Smoothness / 255f;
        occlusion = s.Occlusion / 255f;
        metallicSmoothnessOpacity = s.MSOpacity / 255f;
        occlusionOpacity = s.OccOpacity / 255f;
        activeBrushIndex = s.BrushIndex;

        if (spray) BuildSprayStampMaskSeeded(s.SpraySeed);
        else { ResolveActiveMask(); UseResolvedStampMask(); }

        BeginBackendStroke();
        DrawBackendStamp(s.P, s.R, s.U, s.V, s.N);
        EndBackendStroke();

        // Accumulate this peer's stroke dirty mask so a network checkpoint marker can commit the same
        // history entry as the owner (same stamps -> same touched maps -> COW-aligned with the owner).
        MarkStrokePainted();
    }

    // --- Late-join snapshot: the whole painted state as a compressed blob ---
    //
    // Live stamps only reach peers connected at paint time; a late joiner needs the accumulated
    // result as a baseline. Instead of streaming 12 MB of raw maps, we dump the canonical maps once
    // and Deflate them (paint is largely flat -> compresses well), then PlayerPaintSync chunks the
    // blob over the network to the single joining client and re-seeds their painter with it.

    private const uint SnapshotMagic = 0x50414e53; // 'SNAP' little-endian
    private const int SnapshotVersion = 1;

    /// <summary>True once Start() has built the backend, so capture/apply are safe.</summary>
    public bool IsPaintReady => _setupValid;

    /// <summary>
    /// Serializes every paint map (canonical layout) and Deflate-compresses the result. Returns false
    /// if the painter is not set up yet. Run on the authority that holds the up-to-date maps.
    /// </summary>
    public bool TryCaptureSnapshot(out byte[] compressed)
    {
        compressed = null;
        if (!_setupValid) return false;

        int n = textureSize * textureSize;
        using (var raw = new MemoryStream(16 + n * 4 * MAP_COUNT))
        {
            using (var w = new BinaryWriter(raw, System.Text.Encoding.UTF8, leaveOpen: true))
            {
                w.Write(SnapshotMagic);
                w.Write(SnapshotVersion);
                w.Write(textureSize);
                w.Write(MAP_COUNT);

                WriteCanonicalMaps(w);
            }
            compressed = Deflate(raw.ToArray());
        }
        return true;
    }

    /// <summary>
    /// Decompresses a snapshot, re-seeds the active backend with it, and makes it this peer's sole history
    /// baseline under the authority's <paramref name="baselineId"/> (so subsequent checkpoint/restore
    /// commands line up). Caller must ensure IsPaintReady.
    /// </summary>
    public void ApplyPaintSnapshot(byte[] compressed, int baselineId)
    {
        if (!TryDecodeSnapshot(compressed, "Paint", out Color32[][] canonical)) return;

        ReseedCanonical(canonical);
        SetHistoryBaseline(canonical, baselineId);
    }

    /// <summary>
    /// Owner-side: re-seeds the painter from an in-memory snapshot blob and resets history to it as a
    /// fresh local baseline (like <see cref="LoadPaintByName"/>, just from bytes) - used to carry a
    /// painting across a scene-switch respawn, then rebroadcast via BroadcastLoadedPaint. Returns false
    /// if the painter is not ready or the blob is incompatible (the current painting stays untouched).
    /// </summary>
    public bool LoadPaintFromSnapshot(byte[] compressed)
    {
        if (!TryDecodeSnapshot(compressed, "Carry", out Color32[][] canonical)) return false;

        ReseedCanonical(canonical);
        ResetHistoryToCurrent(canonical);
        return true;
    }

    /// <summary>Guard + decode for an incoming snapshot blob; false = not ready / empty / undecodable.</summary>
    private bool TryDecodeSnapshot(byte[] compressed, string context, out Color32[][] canonical)
    {
        canonical = null;
        if (!_setupValid || compressed == null || compressed.Length == 0) return false;

        try { canonical = DecodeSnapshot(compressed); }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] {context} snapshot decode failed: {e.Message}", this);
            return false;
        }
        return canonical != null;
    }

    private Color32[][] DecodeSnapshot(byte[] compressed)
    {
        byte[] raw = Inflate(compressed);
        using (var ms = new MemoryStream(raw))
        using (var r = new BinaryReader(ms))
        {
            if (r.ReadUInt32() != SnapshotMagic) return null;
            int ver = r.ReadInt32();
            int size = r.ReadInt32();
            int maps = r.ReadInt32();

            // Resolution / version mismatch: ignore rather than smear; the painter keeps its state.
            if (ver != SnapshotVersion || size != textureSize) return null;

            return ReadCanonicalMaps(r, maps);
        }
    }

    private static byte[] Deflate(byte[] data)
    {
        using (var outp = new MemoryStream())
        {
            using (var ds = new DeflateStream(outp, System.IO.Compression.CompressionLevel.Fastest, leaveOpen: true))
                ds.Write(data, 0, data.Length);
            return outp.ToArray();
        }
    }

    private static byte[] Inflate(byte[] data)
    {
        using (var inp = new MemoryStream(data))
        using (var ds = new DeflateStream(inp, CompressionMode.Decompress))
        using (var outp = new MemoryStream())
        {
            ds.CopyTo(outp);
            return outp.ToArray();
        }
    }
}