chameleon/Assets/Scripts/MeshPainter.Undo.cs

using System.Collections.Generic;
using UnityEngine;
#if ENABLE_INPUT_SYSTEM
using UnityEngine.InputSystem;
#endif

public partial class MeshPainter
{
    // --- Undo / redo: per-stroke history of the canonical maps, mirrored on every peer ---
    //
    // One stroke (mouse-down -> mouse-up) = one undo step; each entry is the full canonical map set AFTER
    // that edit, restored by re-seeding the backend. Index 0 is the baseline (material seed, or a
    // loaded/snapshotted one).
    //
    // Networked: every replica keeps the SAME history, so undo/redo replicates as a tiny command. The
    // owner tags each committed stroke with a monotonic CHECKPOINT ID; remotes capture history at the same
    // boundaries under that id. Undo/redo carries the TARGET id; a peer lacking it (a late joiner) falls
    // back to a snapshot. PlayerPaintSync drives this via OnHistoryCheckpoint / OnHistoryRestore.
    //
    // Memory: unchanged maps are shared by reference with the previous entry (copy-on-write) - safe
    // because both backends CLONE their seed input and ReadCanonical returns fresh arrays.

    [Header("Undo")]
    [Tooltip("Maximum number of undoable edits kept in memory (strokes + loads). Older edits drop off; " +
             "the oldest kept state is the undo floor. Higher = more undo depth but more memory (only the " +
             "maps a stroke changed are stored per step).")]
    [Min(1)]
    public int maxUndoSteps = 20;

    /// <summary>
    /// Fires on the owner when a stroke is committed as a new history entry, carrying its checkpoint id.
    /// PlayerPaintSync forwards a marker so every remote captures the same entry under the same id. No
    /// subscriber = no-op (single-player).
    /// </summary>
    public event System.Action<int> OnHistoryCheckpoint;

    /// <summary>
    /// Fires on the owner when the painting is restored to a history entry (undo / redo), carrying that
    /// entry's checkpoint id. PlayerPaintSync forwards the id so peers reproduce the move locally (or fetch
    /// a snapshot if they lack that id). No subscriber = no-op (single-player).
    /// </summary>
    public event System.Action<int> OnHistoryRestore;

    private List<Color32[][]> _undoStates; // index 0 = oldest kept state; null until InitUndoHistory ran
    private List<int> _undoIds;            // parallel checkpoint id per state (owner-assigned, peer-shared)
    private int _undoIndex;                // the state the painter currently shows
    private int _nextCheckpointId;         // owner: next id to assign; remotes set ids from the network

    private bool _strokeActive;   // a stroke has painted >=1 stamp and is not yet committed to history
    private int _strokeDirtyMaps; // bitmask of map indices touched by the current stroke

    /// <summary>The checkpoint id of the state currently shown. Used as the baseline tag when snapshotting.</summary>
    public int CurrentCheckpointId =>
        (_undoIds != null && _undoIndex >= 0 && _undoIndex < _undoIds.Count) ? _undoIds[_undoIndex] : 0;

    /// <summary>Seeds the history with the material baseline (id 0). Called once after the backend is set up.</summary>
    private void InitUndoHistory()
    {
        if (_seedCanonical == null) return;
        SetSingleBaseline(_seedCanonical, 0);
        _nextCheckpointId = 1;
    }

    /// <summary>
    /// Rebuilds the whole history to a single baseline entry holding <paramref name="maps"/> (per-map null
    /// falls back to the live backend), tagged <paramref name="id"/>. Clears all undo/redo. Used by the
    /// material seed, a load (owner) and an applied snapshot (remote). Takes ownership of the passed arrays.
    /// </summary>
    private void SetSingleBaseline(Color32[][] maps, int id)
    {
        if (_undoStates == null) _undoStates = new List<Color32[][]>(Mathf.Max(2, maxUndoSteps + 1));
        if (_undoIds == null) _undoIds = new List<int>(Mathf.Max(2, maxUndoSteps + 1));
        _undoStates.Clear();
        _undoIds.Clear();

        int n = textureSize * textureSize;
        var baseline = new Color32[MAP_COUNT][];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            var s = (maps != null) ? maps[mi] : null;
            baseline[mi] = (s != null && s.Length == n) ? s : ReadCanonical(mi);
        }
        _undoStates.Add(baseline);
        _undoIds.Add(id);
        _undoIndex = 0;
        _strokeActive = false;
        _strokeDirtyMaps = 0;
    }

    /// <summary>Owner-only: replace the whole history with the current maps as the sole baseline (used by load).</summary>
    private int ResetHistoryToCurrent(Color32[][] maps)
    {
        int id = _nextCheckpointId++;
        SetSingleBaseline(maps, id);
        return id;
    }

    /// <summary>Remote-only: make the just-applied snapshot maps the sole history baseline under the authority's id.</summary>
    public void SetHistoryBaseline(Color32[][] maps, int baselineId)
    {
        if (!_setupValid) return;
        SetSingleBaseline(maps, baselineId);
    }

    /// <summary>Marks that the current stroke painted a stamp (called per committed stamp, local or replicated).</summary>
    private void MarkStrokePainted()
    {
        _strokeActive = true;
        if (paintAlbedo) _strokeDirtyMaps |= 1 << MAP_ALBEDO;
        if (paintMetallicSmoothness || paintOcclusion) _strokeDirtyMaps |= 1 << MAP_MSO;
        if (paintEmission) _strokeDirtyMaps |= 1 << MAP_EM;
    }

    /// <summary>
    /// Owner: commits an in-progress stroke as one undo step and announces its checkpoint id (no-op if
    /// nothing was painted). Remotes use RemoteCommitCheckpoint instead.
    /// </summary>
    private void EndStrokeIfActive()
    {
        if (!_strokeActive) return;

        // Stroke held but nothing actually painted (no active channel / no stamp landed): drop it without
        // consuming a checkpoint id, so the owner only ever announces ids it really holds an entry for.
        if (_strokeDirtyMaps == 0)
        {
            _strokeActive = false;
            return;
        }

        int id = _nextCheckpointId++;
        CaptureUndoState(_strokeDirtyMaps, id);
        _strokeActive = false;
        _strokeDirtyMaps = 0;
        OnHistoryCheckpoint?.Invoke(id);
    }

    /// <summary>
    /// Remote: commits the stamps replayed since the last checkpoint as one undo step under the owner's id,
    /// keeping this peer's history in lockstep. Always records the id - even if (rarely) no stamp landed
    /// here for the stroke - so every replica's history stays 1:1 aligned with the owner's and a later
    /// undo/redo to this id reproduces locally instead of forcing a snapshot fallback.
    /// </summary>
    public void RemoteCommitCheckpoint(int id)
    {
        if (!_setupValid) return;
        CaptureUndoState(_strokeDirtyMaps, id, forceEntry: true);
        _strokeActive = false;
        _strokeDirtyMaps = 0;
    }

    /// <summary>
    /// Snapshots the current maps as a new history entry under <paramref name="id"/>, dropping any redo
    /// branch. With <paramref name="forceEntry"/> an entry is added even when nothing changed (it then
    /// shares every map with the previous entry via COW) - used by remote checkpoints to keep ids aligned.
    /// </summary>
    private void CaptureUndoState(int dirtyMaps, int id, bool forceEntry = false)
    {
        if (_undoStates == null) return;
        if (dirtyMaps == 0 && !forceEntry) return;

        // A fresh edit invalidates everything ahead of the cursor (the redo branch).
        if (_undoIndex < _undoStates.Count - 1)
        {
            int from = _undoIndex + 1;
            int n = _undoStates.Count - from;
            _undoStates.RemoveRange(from, n);
            _undoIds.RemoveRange(from, n);
        }

        Color32[][] prev = _undoStates[_undoIndex];
        var state = new Color32[MAP_COUNT][];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            bool changed = (dirtyMaps & (1 << mi)) != 0;
            state[mi] = changed ? ReadCanonical(mi) : prev[mi]; // COW: share maps the stroke didn't touch
        }

        _undoStates.Add(state);
        _undoIds.Add(id);
        _undoIndex = _undoStates.Count - 1;

        // Bound memory: keep at most maxUndoSteps edits beyond the oldest kept state.
        int maxEntries = Mathf.Max(1, maxUndoSteps) + 1;
        while (_undoStates.Count > maxEntries)
        {
            _undoStates.RemoveAt(0);
            _undoIds.RemoveAt(0);
            _undoIndex--;
        }
    }

    /// <summary>Owner: steps back one edit and announces the restore. Returns false if already at the floor.</summary>
    public bool Undo()
    {
        if (_undoStates == null || _undoIndex <= 0) return false;
        _undoIndex--;
        RestoreCurrentState();
        return true;
    }

    /// <summary>Owner: steps forward one edit (after an undo). Returns false if already at the newest state.</summary>
    public bool Redo()
    {
        if (_undoStates == null || _undoIndex >= _undoStates.Count - 1) return false;
        _undoIndex++;
        RestoreCurrentState();
        return true;
    }

    private void RestoreCurrentState()
    {
        ReseedCanonical(_undoStates[_undoIndex]);
        _strokeActive = false;
        _strokeDirtyMaps = 0;
        OnHistoryRestore?.Invoke(_undoIds[_undoIndex]); // owner -> peers reproduce the move by id
    }

    /// <summary>
    /// Remote: restores to the history entry tagged <paramref name="id"/>. Returns false (and does nothing)
    /// if this peer doesn't hold that id - the caller then falls back to fetching a snapshot.
    /// </summary>
    public bool RemoteRestoreTo(int id)
    {
        if (_undoStates == null || _undoIds == null) return false;
        int idx = _undoIds.IndexOf(id);
        if (idx < 0) return false;

        _undoIndex = idx;
        ReseedCanonical(_undoStates[idx]);
        _strokeActive = false;
        _strokeDirtyMaps = 0;
        return true;
    }

    /// <summary>Reads Ctrl+Z / Ctrl+Y this frame. Modifier-gated so plain Z/Y never trigger it.</summary>
    private void GetUndoRedoInput(out bool undo, out bool redo)
    {
        undo = false;
        redo = false;
        if (TypingGuard.IsTyping) return; // let Ctrl+Z/Y edit text instead of the painting
#if ENABLE_INPUT_SYSTEM
        var kb = Keyboard.current;
        if (kb != null)
        {
            if (kb.leftCtrlKey.isPressed || kb.rightCtrlKey.isPressed)
            {
                undo = kb.zKey.wasPressedThisFrame;
                redo = kb.yKey.wasPressedThisFrame;
            }
            return;
        }
#endif
#if ENABLE_LEGACY_INPUT_MANAGER
        if (Input.GetKey(KeyCode.LeftControl) || Input.GetKey(KeyCode.RightControl))
        {
            undo = Input.GetKeyDown(KeyCode.Z);
            redo = Input.GetKeyDown(KeyCode.Y);
        }
#endif
    }
}