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

using System.Collections.Generic;
using System.IO;
using UnityEngine;

public partial class MeshPainter
{
    // --- Persistence: named paint slots, saved/loaded explicitly from the HUD ---
    //
    // The painter no longer auto-loads on start or auto-saves on quit; painting is volatile until the
    // player explicitly saves it to a named slot, and a fresh session always starts from the material
    // seed. Saved slots live in a single global library folder (role/object independent), so any avatar
    // can load any saved painting as a camouflage preset.
    //
    // On-disk layout: one human-readable folder per painting.
    //   Paintings/<key>/albedo.png       sRGB colour   (the visible camo pattern)
    //   Paintings/<key>/emission.png     sRGB colour
    //   Paintings/<key>/metallic.png     greyscale = canonical MSO.R
    //   Paintings/<key>/occlusion.png    greyscale = canonical MSO.G
    //   Paintings/<key>/smoothness.png   greyscale = canonical MSO.A
    //   Paintings/<key>/fur.png          R = fur length (painted length data, not coverage); G/B/A unused
    //   Paintings/<key>/painting.json    display name + map size + the non-baked material scalars
    // The packed mask map is split into three greyscale images so every file is viewable/editable. This
    // also makes the round-trip byte-exact by construction: the colour maps are forced opaque and the mask
    // channels become opaque greyscale, so there is no PNG alpha / premultiply ambiguity. fur.png is a data
    // map (R = length), stored verbatim. Canonical MSO.B is always 0 (see MeshPaintPackMS.shader) and is
    // reconstructed on load. The downstream reseed path (ReseedCanonical) is unchanged, so pixel fidelity
    // equals the previous .paint format plus exact PNG transport.

    // v6 adds fur.png. v5 = split PNG maps + JSON scalars. Loads are not version-gated (only the resolution
    // is), so saves without a fur.png load fine - the fur map stays seeded (bald). NOTE: the fur.png layout
    // changed to R = length (was RGB colour + A length); pre-change fur.png files load stale length data and
    // should be repainted.
    private const int PaintMetaVersion = 6;

    private const string MetaFileName = "painting.json";
    private const string AlbedoFileName = "albedo.png";
    private const string EmissionFileName = "emission.png";
    private const string MetallicFileName = "metallic.png";
    private const string OcclusionFileName = "occlusion.png";
    private const string SmoothnessFileName = "smoothness.png";
    private const string FurFileName = "fur.png";

    /// <summary>
    /// JSON sidecar for a slot: display name, map resolution, and the per-player material scalars that are
    /// NOT baked into the painted pixels (a global emission multiplier + the PlayerCamoLit camo knobs + the
    /// painted-height knobs). They ride alongside the maps and are re-applied onto the painter on load.
    /// JsonUtility needs public fields on a concrete [Serializable] type.
    /// </summary>
    [System.Serializable]
    private class PaintMeta
    {
        public string name;
        public int version = PaintMetaVersion;
        public int width;
        public int height;
        public float emissionIntensity;
        public float camoHalfLambert;
        public float camoHalfLambertBrightness;
        public float camoSpecularDampen;
        public float occlusionThreshold;
        public float heightNormalIntensity;
        public float parallaxIntensity;
    }

    private static string PaintLibraryDir =>
        Path.Combine(Application.persistentDataPath, "Paintings");

    /// <summary>Filesystem-safe key: invalid filename characters and spaces become underscores.</summary>
    internal static string SanitizeFileKey(string s)
    {
        foreach (char c in Path.GetInvalidFileNameChars()) s = s.Replace(c, '_');
        return s.Replace(' ', '_');
    }

    private static string SanitizeKey(string s) =>
        string.IsNullOrEmpty(s) ? "unnamed" : SanitizeFileKey(s);

    /// <summary>Per-painting folder. The display name lives in the JSON; the folder name is a sanitized key.</summary>
    private static string FolderForName(string name) =>
        Path.Combine(PaintLibraryDir, SanitizeKey(name));

    /// <summary>True if a slot with this name already exists (used to block overwriting on save).</summary>
    public static bool PaintExists(string name) =>
        !string.IsNullOrWhiteSpace(name) && Directory.Exists(FolderForName(name));

    /// <summary>Display names of every saved slot, sorted alphabetically. Each name is read from its JSON sidecar.</summary>
    public static List<string> ListSavedPaintings()
    {
        var names = new List<string>();
        try
        {
            if (!Directory.Exists(PaintLibraryDir)) return names;
            foreach (string dir in Directory.GetDirectories(PaintLibraryDir))
            {
                string n = ReadDisplayName(dir);
                if (!string.IsNullOrEmpty(n)) names.Add(n);
            }
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Listing saved paintings failed: {e.Message}");
        }
        names.Sort(System.StringComparer.CurrentCultureIgnoreCase);
        return names;
    }

    /// <summary>
    /// Writes the current painted maps to a named slot folder. Returns false without writing if the painter
    /// is not ready. An existing slot of the same name is overwritten only when <paramref name="overwrite"/>
    /// is true; otherwise the call fails rather than clobbering it.
    /// </summary>
    public bool SavePaintAs(string name, bool overwrite = false)
    {
        if (!_setupValid || string.IsNullOrWhiteSpace(name)) return false;
        name = name.Trim();
        if (!overwrite && PaintExists(name)) return false;

        try
        {
            string dir = FolderForName(name);
            Directory.CreateDirectory(dir);

            Color32[] albedo = ReadCanonical(MAP_ALBEDO);
            Color32[] mso = ReadCanonical(MAP_MSO);
            Color32[] emission = ReadCanonical(MAP_EM);
            Color32[] fur = ReadCanonical(MAP_FUR);

            // Colour maps are stored opaque (the reseed path re-derives coverage/alpha anyway). The packed
            // mask map is split into one opaque greyscale image per channel: R=metallic, G=occlusion,
            // A=smoothness (canonical MSO.B is always 0).
            WriteMapPng(Path.Combine(dir, AlbedoFileName), albedo, c => new Color32(c.r, c.g, c.b, 255));
            WriteMapPng(Path.Combine(dir, EmissionFileName), emission, c => new Color32(c.r, c.g, c.b, 255));
            WriteMapPng(Path.Combine(dir, MetallicFileName), mso, c => new Color32(c.r, c.r, c.r, 255));
            WriteMapPng(Path.Combine(dir, OcclusionFileName), mso, c => new Color32(c.g, c.g, c.g, 255));
            WriteMapPng(Path.Combine(dir, SmoothnessFileName), mso, c => new Color32(c.a, c.a, c.a, 255));
            // Fur is data, not colour: R = painted length (0 = bald), G/B/A unused. Store verbatim so the
            // length round-trips exactly instead of being forced opaque.
            WriteMapPng(Path.Combine(dir, FurFileName), fur, c => c);

            // The painter is ready here (_setupValid), so the material-backed camo/height getters return the
            // live values rather than the not-ready 0.
            var meta = new PaintMeta
            {
                name = name,
                version = PaintMetaVersion,
                width = textureSize,
                height = textureSize,
                emissionIntensity = emissionIntensity,
                camoHalfLambert = CamoHalfLambert,
                camoHalfLambertBrightness = CamoHalfLambertBrightness,
                camoSpecularDampen = CamoSpecularDampen,
                occlusionThreshold = OcclusionThreshold,
                heightNormalIntensity = HeightNormalIntensity,
                parallaxIntensity = ParallaxIntensity,
            };
            File.WriteAllText(Path.Combine(dir, MetaFileName), JsonUtility.ToJson(meta, true));
            return true;
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Saving paint '{name}' failed: {e.Message}", this);
            return false;
        }
    }

    /// <summary>
    /// Loads a named slot and re-seeds the active backend with it, replacing the current painting.
    /// Returns false if the slot is missing or incompatible (the current painting is left untouched).
    /// </summary>
    public bool LoadPaintByName(string name)
    {
        if (!_setupValid) return false;

        string dir = FolderForName(name);
        if (!Directory.Exists(dir)) return false;

        PaintMeta meta = ReadMeta(dir);

        // No resolution gate: a slot saved at any size is bilinearly resampled to this avatar's textureSize
        // per map (see ReadMapPng), so any painting loads onto any avatar regardless of the size it was
        // saved at. meta.width/height are kept only as informational metadata.

        Color32[][] canonical = ReadFolderMaps(dir);
        if (canonical == null) return false;

        ReseedCanonical(canonical);

        // A load replaces the whole painting and is NOT undoable: the loaded maps become the new history
        // baseline, clearing any prior undo/redo. Peers reset to the same baseline via the load broadcast.
        ResetHistoryToCurrent(canonical);

        // Apply the slot's per-player material scalars so the loaded look is complete. The painter is ready
        // here (_setupValid), so the material-backed setters take effect; PlayerPaintSync replicates the
        // changed values to peers via their owner-authoritative NetworkVariables on its next LateUpdate.
        // A slot with no readable sidecar leaves the current emission/camo/height untouched.
        if (meta != null)
        {
            emissionIntensity = meta.emissionIntensity;
            RefreshEmissionColor();
            CamoHalfLambert = meta.camoHalfLambert;
            CamoHalfLambertBrightness = meta.camoHalfLambertBrightness;
            CamoSpecularDampen = meta.camoSpecularDampen;
            OcclusionThreshold = meta.occlusionThreshold;
            HeightNormalIntensity = meta.heightNormalIntensity;
            ParallaxIntensity = meta.parallaxIntensity;
        }
        return true;
    }

    /// <summary>Deletes a named slot folder and everything in it. Returns true if a folder was removed.</summary>
    public static bool DeletePaint(string name)
    {
        try
        {
            string dir = FolderForName(name);
            if (!Directory.Exists(dir)) return false;
            Directory.Delete(dir, recursive: true);
            return true;
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Deleting paint '{name}' failed: {e.Message}");
            return false;
        }
    }

    /// <summary>Reads a slot's display name from its JSON sidecar (the folder name as a fallback).</summary>
    private static string ReadDisplayName(string dir)
    {
        var meta = ReadMeta(dir);
        return (meta != null && !string.IsNullOrEmpty(meta.name)) ? meta.name : Path.GetFileName(dir);
    }

    /// <summary>Reads and parses a slot's JSON sidecar, or null if it is missing / unreadable.</summary>
    private static PaintMeta ReadMeta(string dir)
    {
        try
        {
            string metaPath = Path.Combine(dir, MetaFileName);
            if (!File.Exists(metaPath)) return null;
            return JsonUtility.FromJson<PaintMeta>(File.ReadAllText(metaPath));
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Reading paint meta in '{Path.GetFileName(dir)}' failed: {e.Message}");
            return null;
        }
    }

    /// <summary>
    /// Reads a slot folder into canonical (material) layout, or null if nothing usable was read. Per-map
    /// null entries are tolerated by both backends (they fall back to the material seed): a missing colour
    /// image leaves that map seeded; the mask map needs all three channel images to be rebuilt.
    /// </summary>
    private Color32[][] ReadFolderMaps(string dir)
    {
        var result = new Color32[MAP_COUNT][];
        int loaded = 0;

        if (ReadMapPng(Path.Combine(dir, AlbedoFileName), out Color32[] albedo))
        { result[MAP_ALBEDO] = albedo; loaded++; }
        if (ReadMapPng(Path.Combine(dir, EmissionFileName), out Color32[] emission))
        { result[MAP_EM] = emission; loaded++; }

        // Fur carries painted length in R; loaded verbatim. Missing on saves without a fur map -> it stays
        // seeded (bald), the correct backward-compatible default.
        if (ReadMapPng(Path.Combine(dir, FurFileName), out Color32[] fur))
        { result[MAP_FUR] = fur; loaded++; }

        // Recombine the split greyscale channels into the packed canonical layout
        // (R=metallic, G=occlusion, B=0, A=smoothness). B is always 0 (see MeshPaintPackMS.shader).
        if (ReadMapPng(Path.Combine(dir, MetallicFileName), out Color32[] metallic) &&
            ReadMapPng(Path.Combine(dir, OcclusionFileName), out Color32[] occlusion) &&
            ReadMapPng(Path.Combine(dir, SmoothnessFileName), out Color32[] smoothness))
        {
            int n = textureSize * textureSize;
            var mso = new Color32[n];
            for (int i = 0; i < n; i++)
                mso[i] = new Color32(metallic[i].r, occlusion[i].r, 0, smoothness[i].r);
            result[MAP_MSO] = mso;
            loaded++;
        }

        return loaded > 0 ? result : null;
    }

    /// <summary>
    /// Writes one map as a PNG, remapping each canonical texel through <paramref name="map"/> (colour maps
    /// forced opaque; a mask channel expanded to opaque greyscale). RGBA32 + opaque alpha makes the round-
    /// trip byte-exact. Returns false if the source array is missing / wrong size.
    /// </summary>
    private bool WriteMapPng(string path, Color32[] canonical, System.Func<Color32, Color32> map)
    {
        int n = textureSize * textureSize;
        if (canonical == null || canonical.Length < n) return false;

        var pixels = new Color32[n];
        for (int i = 0; i < n; i++) pixels[i] = map(canonical[i]);

        var tex = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false);
        try
        {
            tex.SetPixels32(pixels);
            tex.Apply(false);
            File.WriteAllBytes(path, tex.EncodeToPNG());
            return true;
        }
        finally { Destroy(tex); }
    }

    /// <summary>
    /// Loads a map PNG into canonical (material) layout; false only if the file is missing / undecodable.
    /// A PNG whose resolution differs from the current textureSize is bilinearly resampled up/down to
    /// textureSize (see ResampleBilinear), so a slot saved at any resolution loads onto any avatar.
    /// </summary>
    private bool ReadMapPng(string path, out Color32[] px)
    {
        px = null;
        if (!File.Exists(path)) return false;

        var tex = new Texture2D(2, 2, TextureFormat.RGBA32, false);
        try
        {
            if (!tex.LoadImage(File.ReadAllBytes(path))) return false;
            px = (tex.width == textureSize && tex.height == textureSize)
                ? tex.GetPixels32()
                : ResampleBilinear(tex.GetPixels32(), tex.width, tex.height, textureSize, textureSize);
            return true;
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Loading map '{Path.GetFileName(path)}' failed: {e.Message}", this);
            return false;
        }
        finally { Destroy(tex); }
    }

    /// <summary>
    /// Bilinearly resamples a raw Color32 byte grid to a new resolution. Interpolation runs on the stored
    /// bytes with no gamma/linear conversion, matching how the maps are consumed downstream: colour maps
    /// (albedo/emission) rescale smoothly, and the greyscale mask channels + the fur length channel
    /// interpolate as linear scalars. Clamped edge sampling; used to load a slot saved at a different
    /// resolution than the current avatar's textureSize.
    /// </summary>
    private static Color32[] ResampleBilinear(Color32[] src, int srcW, int srcH, int dstW, int dstH)
    {
        var dst = new Color32[dstW * dstH];
        for (int y = 0; y < dstH; y++)
        {
            float sy = (y + 0.5f) * srcH / dstH - 0.5f;
            int y0 = Mathf.Clamp(Mathf.FloorToInt(sy), 0, srcH - 1);
            int y1 = Mathf.Min(y0 + 1, srcH - 1);
            float fy = Mathf.Clamp01(sy - y0);
            for (int x = 0; x < dstW; x++)
            {
                float sx = (x + 0.5f) * srcW / dstW - 0.5f;
                int x0 = Mathf.Clamp(Mathf.FloorToInt(sx), 0, srcW - 1);
                int x1 = Mathf.Min(x0 + 1, srcW - 1);
                float fx = Mathf.Clamp01(sx - x0);

                int r0 = y0 * srcW, r1 = y1 * srcW;
                Color32 c00 = src[r0 + x0], c10 = src[r0 + x1];
                Color32 c01 = src[r1 + x0], c11 = src[r1 + x1];

                dst[y * dstW + x] = new Color32(
                    Bilerp(c00.r, c10.r, c01.r, c11.r, fx, fy),
                    Bilerp(c00.g, c10.g, c01.g, c11.g, fx, fy),
                    Bilerp(c00.b, c10.b, c01.b, c11.b, fx, fy),
                    Bilerp(c00.a, c10.a, c01.a, c11.a, fx, fy));
            }
        }
        return dst;
    }

    /// <summary>Bilinear blend of four byte samples (rounded, clamped to 0..255).</summary>
    private static byte Bilerp(byte c00, byte c10, byte c01, byte c11, float fx, float fy)
    {
        float top = c00 + (c10 - c00) * fx;
        float bot = c01 + (c11 - c01) * fx;
        return (byte)Mathf.Clamp(Mathf.RoundToInt(top + (bot - top) * fy), 0, 255);
    }

    // --- Shared canonical map (de)serialization: also used by the network snapshot path (MeshPainter.Net) ---

    /// <summary>Writes every map's canonical Color32 bytes to the stream (shared by the net snapshot).</summary>
    private void WriteCanonicalMaps(BinaryWriter w)
    {
        int n = textureSize * textureSize;
        var bytes = new byte[n * 4];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            // A deferred backend not yet re-seeded holds uninitialized RTs; serve the material seed instead
            // so a snapshot captured in that window is the base look (consistent with its cutoff), not garbage.
            Color32[] px = _backendSeeded ? ReadCanonical(mi) : _seedCanonical[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);
        }
    }

    /// <summary>Reads up to MAP_COUNT canonical map blocks; a truncated stream leaves the remaining maps null (seeded).</summary>
    private Color32[][] ReadCanonicalMaps(BinaryReader r, int mapsInFile)
    {
        int n = textureSize * textureSize;
        var result = new Color32[MAP_COUNT][];
        int load = Mathf.Min(mapsInFile, MAP_COUNT);
        for (int mi = 0; mi < load; mi++)
        {
            byte[] raw = r.ReadBytes(n * 4);
            if (raw.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(raw[o], raw[o + 1], raw[o + 2], raw[o + 3]);
            }
            result[mi] = px;
        }
        return result;
    }
}