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          RGBA: RGB = strand colour, A = fur length (the ONE image whose
    //                                     alpha is meaningful - painted length data, not coverage)
    //   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 the
    // exception - its alpha is length data, 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 (RGBA: strand colour + painted length). v5 = split PNG maps + JSON scalars. Loads are
    // not version-gated (only the resolution is), so pre-v6 saves without a fur.png load fine - the fur map
    // stays seeded (bald).
    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 the exception: RGB = strand colour, A = painted length (0 = bald). Its alpha is real
            // data, so store the full RGBA verbatim instead of forcing it 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);

        // Resolution change -> ignore rather than smear mismatched data (the brush re-seeds from the
        // material instead). A missing/old sidecar has width 0 -> skip the header check and trust the PNGs
        // (ReadMapPng validates each image's resolution on its own).
        if (meta != null && meta.width > 0 && (meta.width != textureSize || meta.height != textureSize))
        {
            Debug.LogWarning($"[MeshPainter] Saved paint ignored (got {meta.width}x{meta.height}, expected " +
                             $"{textureSize}x{textureSize}).", this);
            return false;
        }

        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 a real alpha (painted length); loaded verbatim. Missing on pre-v6 saves -> the fur
        // map 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 if missing / wrong resolution / undecodable.</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;
            if (tex.width != textureSize || tex.height != textureSize)
            {
                Debug.LogWarning($"[MeshPainter] Saved map ignored (got {tex.width}x{tex.height}, expected " +
                                 $"{textureSize}x{textureSize}): {Path.GetFileName(path)}", this);
                return false;
            }
            px = tex.GetPixels32();
            return true;
        }
        catch (System.Exception e)
        {
            Debug.LogWarning($"[MeshPainter] Loading map '{Path.GetFileName(path)}' failed: {e.Message}", this);
            return false;
        }
        finally { Destroy(tex); }
    }

    // --- 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++)
        {
            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);
        }
    }

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