chameleon/Assets/Scripts/MeshPainter.Net.cs

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

/// <summary>
/// One committed brush stamp in a form that can be replicated across the network. Painting is
/// synced as commands (these stamps), not as pixels: the owner already resolves every stamp into
/// local object space (hit point, radius, tangent frame) inside <see cref="MeshPainter.PaintAtScreen"/>,
/// so a remote peer just replays the SAME stamp through its own backend onto the identical mesh/UVs
/// and gets a near-identical result - for a fraction of the bandwidth of streaming the 12 MB of maps.
///
/// Self-contained on purpose: each stamp carries its full brush state, so the stream is stateless and
/// robust to ordering/loss (sent reliably). Channel data (metallic/smoothness/occlusion + opacities)
/// quantizes to bytes with no real loss because 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(); }

        if (_gpuActive)
        {
            BeginGpuStroke();
            DrawBrushGpu(s.P, s.R, s.U, s.V, s.N);
            EndGpuStroke();
        }
        else
        {
            ResetCpuDirty();
            PaintStamp3D(s.P, s.R, s.U, s.V, s.N);
            EndCpuStroke();
        }

        // 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);

                var bytes = new byte[n * 4];
                for (int mi = 0; mi < MAP_COUNT; mi++)
                {
                    Color32[] px = ReadCanonical(mi);
                    int count = Mathf.Min(n, px.Length);
                    for (int i = 0; i < count; i++)
                    {
                        int o = i * 4;
                        bytes[o] = px[i].r;
                        bytes[o + 1] = px[i].g;
                        bytes[o + 2] = px[i].b;
                        bytes[o + 3] = px[i].a;
                    }
                    w.Write(bytes, 0, count * 4);
                }
            }
            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 (!_setupValid || compressed == null || compressed.Length == 0) return;

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

        // Re-seed exactly like SwitchBackend, but with externally supplied maps, so the island/dilation
        // seeding path runs and the edge padding is rebuilt correctly for the loaded content.
        ReseedCanonical(canonical);
        SetHistoryBaseline(canonical, baselineId);
    }

    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;

            int n = textureSize * textureSize;
            var result = new Color32[MAP_COUNT][];
            int load = Mathf.Min(maps, MAP_COUNT);
            for (int mi = 0; mi < load; mi++)
            {
                byte[] b = r.ReadBytes(n * 4);
                if (b.Length != n * 4) break; // truncated -> remaining maps stay at the material seed
                var px = new Color32[n];
                for (int i = 0; i < n; i++)
                {
                    int o = i * 4;
                    px[i] = new Color32(b[o], b[o + 1], b[o + 2], b[o + 3]);
                }
                result[mi] = px;
            }
            return result;
        }
    }

    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();
        }
    }
}