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.
///
/// Undo/redo replicates as commands too, not pixels: every replica keeps the SAME per-stroke history
/// (see MeshPainter.Undo), so the owner just forwards a per-stroke CHECKPOINT marker and an undo/redo
/// RESTORE id. Each peer reproduces the move from its own history. A peer that lacks the target id (a
/// late joiner asked to undo past its join) heals by pulling a fresh snapshot from the OWNER - the only
/// guaranteed-correct source of the post-restore state, so this works for a remote client and for the
/// host's own replica alike. Stamps are capped per frame but their overflow carries to the next frame
/// rather than being decimated, and checkpoint/restore markers never overtake the stamps they finalize,
/// so every replica's history stays aligned and the heal path stays a rare exception, not the norm.
///
/// Late join: a client that connects AFTER painting happened 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 events that arrive during the transfer,
/// and after applying the snapshot replays the buffered events newer than the snapshot's cutoff
/// sequence - so no stroke is lost or double-applied across the join window. The applied snapshot also
/// becomes the joiner's history baseline.
///
/// Emission intensity is the one paint property that is NOT baked into the maps (it's a global
/// multiplier on the material's _EmissionColor), so it can't ride the stamp/snapshot path. It rides its
/// own owner-authoritative NetworkVariable; the remote painter's per-frame RefreshEmissionColor picks up
/// the applied field, and a late joiner gets the current value on spawn for free.
/// </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 capped to this to bound the RPC size; the " +
             "overflow is CARRIED OVER to the next frame (never dropped), so remotes replay the identical " +
             "stamp set and undo/redo history stays aligned. Local painting is always full fidelity.")]
    [SerializeField] private int maxStampsPerFrame = 64;

    [Tooltip("Log paint-sync events (stamps/checkpoints/restores sent & applied, snapshot transfers) to " +
             "the console with a per-peer tag. Handy for diagnosing undo/redo replication; turn off for " +
             "quiet play.")]
    [SerializeField] private bool logSync = true;

    [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;

    // Single owner-side FIFO of everything to forward, in the exact order it was produced: brush stamps
    // interleaved with stroke checkpoints and undo/redo restores. Flushed in order in LateUpdate so a
    // checkpoint/restore never overtakes the stamps it refers to, and so the per-frame stamp cap can carry
    // overflow to the next frame without ever reordering or dropping a command.
    private enum OutKind { Stamp, Checkpoint, Restore }
    private struct OutItem
    {
        public OutKind kind;
        public PaintStampNet stamp; // Stamp only
        public int id;              // Checkpoint / Restore only
    }
    private readonly List<OutItem> _outgoing = new List<OutItem>();

    // Global emission-intensity multiplier (MeshPainter.emissionIntensity) — a per-player material value,
    // not part of the painted maps. Owner-authoritative like the rest; the remote applies it to its
    // painter, whose per-frame RefreshEmissionColor pushes it to the material (no extra material poke).
    private readonly NetworkVariable<float> _emissionIntensity =
        new(1f, NetworkVariableReadPermission.Everyone, NetworkVariableWritePermission.Owner);

    // Monotonic message counter on the owner (one per stamp batch / checkpoint / restore); on every peer,
    // the highest sequence accounted for (applied to a remote replica / sent by the owner). Used as the
    // snapshot's cutoff marker so the join window neither loses nor double-applies events.
    private int _seq;
    private int _seqHorizon;

    // --- Joining client: snapshot reassembly + live-event buffering until the baseline is applied ---
    private enum NetKind { Stamps, Checkpoint, Restore }
    private struct NetEvent
    {
        public int seq;
        public NetKind kind;
        public PaintStampNet[] stamps; // Stamps only
        public int id;                 // Checkpoint / Restore only
    }

    private bool _awaitingSnapshot;          // buffer live events instead of applying them
    private byte[] _snapBuffer;              // reassembled compressed snapshot
    private int _snapCapturedSeq;            // events <= this are already baked into the snapshot
    private int _snapBaselineId;             // checkpoint id the snapshot maps represent (history baseline)
    private bool _snapReadyToApply;          // last chunk arrived; apply once the painter is ready
    private readonly List<NetEvent> _buffered = new List<NetEvent>();

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

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

        if (IsOwner)
        {
            _painter.OnStampCommitted += OnStampCommitted;
            _painter.OnHistoryCheckpoint += OnLocalCheckpoint;
            _painter.OnHistoryRestore += OnLocalRestore;
        }
        else
        {
            // Remote replica: never paint from this machine's mouse; paint only from replicated data.
            _painter.PaintingEnabled = false;

            // Track the owner's emission intensity; seed the current (late-join) value right away.
            _emissionIntensity.OnValueChanged += OnEmissionChanged;
            _painter.emissionIntensity = _emissionIntensity.Value;

            // 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;
                LogSync("late joiner: requesting baseline snapshot");
                RequestSnapshotRpc();
            }
        }
    }

    public override void OnNetworkDespawn()
    {
        if (_painter && IsOwner)
        {
            _painter.OnStampCommitted -= OnStampCommitted;
            _painter.OnHistoryCheckpoint -= OnLocalCheckpoint;
            _painter.OnHistoryRestore -= OnLocalRestore;
        }
        _emissionIntensity.OnValueChanged -= OnEmissionChanged;
    }

    // Owner history events -> appended to the ordered outgoing FIFO, flushed in LateUpdate.
    private void OnLocalCheckpoint(int id) => _outgoing.Add(new OutItem { kind = OutKind.Checkpoint, id = id });
    private void OnLocalRestore(int id) => _outgoing.Add(new OutItem { kind = OutKind.Restore, id = id });

    // Per-peer log tag: which client this is and its role, so multi-window logs are easy to read.
    private string SyncTag
    {
        get
        {
            string role = IsOwner ? "owner" : (IsServer ? "host" : "client");
            ulong id = NetworkManager ? NetworkManager.LocalClientId : 0;
            return $"[PaintSync c{id} {role}]";
        }
    }

    private void LogSync(string msg)
    {
        if (logSync) Debug.Log($"{SyncTag} {msg}", this);
    }

    private void OnEmissionChanged(float previous, float current)
    {
        if (_painter) _painter.emissionIntensity = current;
    }

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

    private void OnStampCommitted(PaintStampNet stamp) =>
        _outgoing.Add(new OutItem { kind = OutKind.Stamp, stamp = stamp });

    private void LateUpdate()
    {
        // Apply a fully-received snapshot once the painter has finished setup, then flush the events
        // 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, _snapBaselineId);
            LogSync($"applied snapshot baseline {_snapBaselineId}; replaying buffered events");
            FlushBuffered();
            _snapBuffer = null;
            _awaitingSnapshot = false;
        }

        if (!IsOwner)
            return;

        // Publish the owner's global emission intensity when the HUD slider changes it (delta-only; the
        // NetworkVariable dirties only on an actual change, so assigning the same value is a no-op).
        if (_painter && _emissionIntensity.Value != _painter.emissionIntensity)
            _emissionIntensity.Value = _painter.emissionIntensity;

        FlushOutgoing(unlimited: false);
    }

    /// <summary>
    /// Sends the outgoing FIFO in order: runs of consecutive stamps are batched (capped per frame unless
    /// <paramref name="unlimited"/>), checkpoints/restores go out as tiny markers in their exact position.
    /// Each item gets its own monotonic sequence. With the per-frame cap reached, the unsent tail (stamps
    /// and any markers behind them) stays queued for the next frame - never reordered, never dropped.
    /// </summary>
    private void FlushOutgoing(bool unlimited)
    {
        if (_outgoing.Count == 0)
            return;

        int cap = Mathf.Max(1, maxStampsPerFrame);
        int sentStamps = 0;
        int consumed = 0;

        while (consumed < _outgoing.Count)
        {
            OutKind kind = _outgoing[consumed].kind;

            if (kind == OutKind.Stamp)
            {
                if (!unlimited && sentStamps >= cap)
                    break; // per-frame budget spent; carry the remainder (this stamp on) to next frame

                int budget = unlimited ? int.MaxValue : (cap - sentStamps);
                int runEnd = consumed;
                while (runEnd < _outgoing.Count
                       && _outgoing[runEnd].kind == OutKind.Stamp
                       && (runEnd - consumed) < budget)
                    runEnd++;

                int take = runEnd - consumed;
                var batch = new PaintStampNet[take];
                for (int k = 0; k < take; k++)
                    batch[k] = _outgoing[consumed + k].stamp;

                int seq = ++_seq;
                _seqHorizon = seq;
                SubmitStampsServerRpc(seq, batch);
                sentStamps += take;
                consumed = runEnd;
            }
            else
            {
                int id = _outgoing[consumed].id;
                int seq = ++_seq;
                _seqHorizon = seq;
                if (kind == OutKind.Restore)
                {
                    SubmitRestoreServerRpc(seq, id);
                    LogSync($"sent RESTORE to checkpoint {id} (seq {seq})");
                }
                else
                {
                    SubmitCheckpointServerRpc(seq, id);
                    LogSync($"sent checkpoint {id} (seq {seq})");
                }
                consumed++;
            }
        }

        if (consumed > 0)
            _outgoing.RemoveRange(0, consumed);

        if (!unlimited && _outgoing.Count > 0)
            LogSync($"stamp cap ({cap}) reached; {_outgoing.Count} command(s) carried over to next frame");
    }

    // ---- Stamp / checkpoint / restore commands (owner -> server -> remotes) ----

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

    [ServerRpc]
    private void SubmitCheckpointServerRpc(int seq, int id) => ApplyCheckpointClientRpc(seq, id);

    [ServerRpc]
    private void SubmitRestoreServerRpc(int seq, int id) => ApplyRestoreClientRpc(seq, id);

    [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)
        {
            _buffered.Add(new NetEvent { seq = seq, kind = NetKind.Stamps, stamps = stamps });
            return;
        }

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

    [ClientRpc]
    private void ApplyCheckpointClientRpc(int seq, int id)
    {
        if (IsOwner || !_painter)
            return;

        _seqHorizon = seq;

        if (_awaitingSnapshot)
        {
            _buffered.Add(new NetEvent { seq = seq, kind = NetKind.Checkpoint, id = id });
            return;
        }

        _painter.RemoteCommitCheckpoint(id);
        LogSync($"committed checkpoint {id} (seq {seq})");
    }

    [ClientRpc]
    private void ApplyRestoreClientRpc(int seq, int id)
    {
        if (IsOwner || !_painter)
            return;

        _seqHorizon = seq;

        if (_awaitingSnapshot)
        {
            _buffered.Add(new NetEvent { seq = seq, kind = NetKind.Restore, id = id });
            return;
        }

        // Reproduce the undo/redo from this peer's own history; if it doesn't hold that checkpoint (an
        // undo past this peer's join point), fetch a fresh snapshot of the resulting state instead.
        if (_painter.RemoteRestoreTo(id))
        {
            LogSync($"reproduced RESTORE to checkpoint {id} (seq {seq})");
        }
        else
        {
            LogSync($"cannot reproduce checkpoint {id} (not in local history, seq {seq}) -> owner snapshot fallback");
            RequestSnapshotFallback();
        }
    }

    private void FlushBuffered()
    {
        if (_painter)
        {
            for (int b = 0; b < _buffered.Count; b++)
            {
                var ev = _buffered[b];
                if (ev.seq <= _snapCapturedSeq) continue; // already inside the snapshot baseline

                if (ev.kind == NetKind.Stamps)
                {
                    for (int i = 0; i < ev.stamps.Length; i++)
                        _painter.ApplyNetworkStamp(ev.stamps[i]);
                }
                else if (ev.kind == NetKind.Checkpoint)
                {
                    _painter.RemoteCommitCheckpoint(ev.id);
                }
                else // Restore
                {
                    if (!_painter.RemoteRestoreTo(ev.id))
                    {
                        // Can't reproduce -> the fresh snapshot we request next supersedes the rest.
                        LogSync($"buffered RESTORE to checkpoint {ev.id} not reproducible -> owner snapshot fallback");
                        _buffered.Clear();
                        RequestSnapshotFallback();
                        return;
                    }
                }
            }
        }
        _buffered.Clear();
    }

    /// <summary>
    /// A peer that can't reproduce a restore from its own history asks the OWNER - the only guaranteed
    /// correct source of the post-restore state - for a fresh snapshot, addressed back to just this peer.
    /// Works uniformly for a remote client and for the host's own replica (which can't snapshot itself);
    /// the previous host short-circuit silently dropped the heal and left the host diverged.
    /// </summary>
    private void RequestSnapshotFallback()
    {
        _awaitingSnapshot = true;
        RequestOwnerSnapshotRpc();
    }

    [Rpc(SendTo.Owner, InvokePermission = RpcInvokePermission.Everyone)]
    private void RequestOwnerSnapshotRpc(RpcParams p = default)
    {
        // Runs on the owner. Drain any carried-over commands first so the snapshot's cutoff covers them
        // (they would otherwise replay on top of a snapshot that already contains them), then capture the
        // authoritative current maps and ship them to just the requester.
        if (!_painter)
            return;

        FlushOutgoing(unlimited: true);

        if (!_painter.TryCaptureSnapshot(out byte[] data))
            return;

        ulong requester = p.Receive.SenderClientId;
        int baselineId = _painter.CurrentCheckpointId;
        LogSync($"serving fallback snapshot to client {requester} (cutoff {_seq}, baseline {baselineId}, {data.Length} bytes)");
        StartCoroutine(UploadSnapshot(data, _seq, baselineId, broadcast: false, target: requester));
    }

    // ---- 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 message this server replica has applied. The maps represent the painter's
        // current history entry, so its checkpoint id tags the baseline the joiner adopts.
        int capturedSeq = IsOwner ? _seq : _seqHorizon;
        int baselineId = _painter.CurrentCheckpointId;
        LogSync($"serving late-join snapshot to client {p.Receive.SenderClientId} (cutoff {capturedSeq}, baseline {baselineId}, {data.Length} bytes)");
        StartCoroutine(SendSnapshot(new[] { p.Receive.SenderClientId }, data, capturedSeq, baselineId));
    }

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

        BeginSnapshotClientRpc(data.Length, capturedSeq, baselineId, 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, int baselineId, ClientRpcParams p = default)
    {
        _snapBuffer = new byte[totalBytes];
        _snapCapturedSeq = capturedSeq;
        _snapBaselineId = baselineId;
        _snapReadyToApply = false;
        _awaitingSnapshot = true;
        LogSync($"receiving snapshot ({totalBytes} bytes, cutoff {capturedSeq}, baseline {baselineId})");
    }

    [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;
    }

    // ---- Load re-sync (owner -> server -> every other peer) ----
    //
    // Loading a saved painting replaces the owner's whole map set at once (no stamps) and clears its undo
    // history, so the live stream cannot carry it. The owner captures a snapshot, uploads it to the server
    // chunked (same throttling as the download), and the server re-broadcasts it to every other peer
    // through the existing snapshot client RPCs - including the host's own replica via its client-id
    // target. The snapshot's baseline id makes every peer reset its history to the loaded baseline too. A
    // cutoff sequence drops any in-flight/buffered events already superseded by the loaded baseline.

    private byte[] _uploadBuffer; // server-side reassembly of the owner's uploaded snapshot
    private int _uploadCutoff;
    private int _uploadBaselineId;
    private bool _uploadBroadcast; // true = fan out to every peer but the owner (load); false = single target
    private ulong _uploadTarget;   // the single requester when !_uploadBroadcast (fallback heal)

    /// <summary>Owner-only: broadcast the just-loaded painting to all other peers. No-op offline / single-player.</summary>
    public void BroadcastLoadedPaint()
    {
        if (!IsSpawned || !IsOwner || !_painter) return;
        if (!_painter.TryCaptureSnapshot(out byte[] data)) return;

        // The loaded baseline replaces the whole painting and clears history, so any still-queued local
        // stamps/checkpoints from before the load are stale - drop them rather than stream them on top.
        _outgoing.Clear();

        int cutoff = ++_seq; // every event up to here is now superseded by the loaded baseline
        _seqHorizon = cutoff;
        int baselineId = _painter.CurrentCheckpointId; // the load's fresh baseline id
        LogSync($"broadcasting loaded painting to all peers (cutoff {cutoff}, baseline {baselineId}, {data.Length} bytes)");
        StartCoroutine(UploadSnapshot(data, cutoff, baselineId, broadcast: true, target: 0));
    }

    private IEnumerator UploadSnapshot(byte[] data, int cutoff, int baselineId, bool broadcast, ulong target)
    {
        BeginUploadServerRpc(data.Length, cutoff, baselineId, broadcast, target);

        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);
            UploadChunkServerRpc(off, chunk);

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

        EndUploadServerRpc();
    }

    [ServerRpc]
    private void BeginUploadServerRpc(int totalBytes, int cutoff, int baselineId, bool broadcast, ulong target)
    {
        _uploadBuffer = totalBytes > 0 ? new byte[totalBytes] : null;
        _uploadCutoff = cutoff;
        _uploadBaselineId = baselineId;
        _uploadBroadcast = broadcast;
        _uploadTarget = target;
    }

    [ServerRpc]
    private void UploadChunkServerRpc(int offset, byte[] chunk)
    {
        if (_uploadBuffer == null) return;
        if (offset >= 0 && offset + chunk.Length <= _uploadBuffer.Length)
            System.Array.Copy(chunk, 0, _uploadBuffer, offset, chunk.Length);
    }

    [ServerRpc]
    private void EndUploadServerRpc()
    {
        if (_uploadBuffer == null) return;

        // Load broadcast: fan out to every connected peer except the owner (who already loaded it
        // locally). Fallback heal: send only to the single requester. The host, when it is not the owner,
        // is included by its client id and applies it like any other remote.
        ulong[] targets = _uploadBroadcast ? OtherClientIds(OwnerClientId) : new[] { _uploadTarget };
        if (targets.Length > 0)
            StartCoroutine(SendSnapshot(targets, _uploadBuffer, _uploadCutoff, _uploadBaselineId));

        _uploadBuffer = null;
    }

    /// <summary>All connected client ids except the given one (server-side).</summary>
    private ulong[] OtherClientIds(ulong except)
    {
        var ids = NetworkManager.ConnectedClientsIds;
        var list = new List<ulong>(ids.Count);
        for (int i = 0; i < ids.Count; i++)
            if (ids[i] != except) list.Add(ids[i]);
        return list.ToArray();
    }
}