chameleon/Assets/Scripts/Editor/PolyHaven/PolyHavenImporterWindow.cs

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

namespace Chameleon.Editor.PolyHaven
{
    /// <summary>
    /// Editor window: paste a Poly Haven asset URL, fetch its maps, pick resolution +
    /// maps, then download and auto-build a URP/Lit material under Assets/Poly Haven/{Name}/.
    /// No browser — asset selection stays on polyhaven.com.
    /// </summary>
    public sealed class PolyHavenImporterWindow : EditorWindow
    {
        private const string RootFolder = "Assets/Poly Haven";
        private const string DefaultRes = "2k";

        private string _urlInput = "https://polyhaven.com/a/";
        private PolyHavenInfo _info;
        private PolyHavenFiles _files;
        private Texture2D _thumbnail;
        private string _selectedRes = DefaultRes;
        private readonly Dictionary<PolyHavenMapType, bool> _mapEnabled = new Dictionary<PolyHavenMapType, bool>();

        private bool _busy;
        private string _status = "";
        private Vector2 _scroll;

        // ---- Browser ----------------------------------------------------------
        private const int MaxResults = 150;
        private const int MaxThumbInFlight = 6;

        // Catalog cached for the whole session (survives window close, cleared on domain reload).
        private static List<PolyHavenAsset> s_catalog;

        private bool _browseFoldout = true;
        private bool _autoLoadTried;
        private bool _catalogLoading;
        private string _searchQuery = "";
        private readonly List<PolyHavenAsset> _filtered = new List<PolyHavenAsset>();
        private Vector2 _gridScroll;

        private readonly Dictionary<string, Texture2D> _thumbs = new Dictionary<string, Texture2D>();
        private readonly HashSet<string> _thumbRequested = new HashSet<string>();
        private int _thumbsInFlight;

        [MenuItem("Tools/Poly Haven Importer")]
        public static void Open()
        {
            var w = GetWindow<PolyHavenImporterWindow>("Poly Haven Importer");
            w.minSize = new Vector2(440, 480);
        }

        private void OnEnable()
        {
            // Catalog may already be cached from a previous window instance this session.
            if (s_catalog != null)
            {
                _autoLoadTried = true;
                ApplyFilter();
            }
        }

        private void OnGUI()
        {
            _scroll = EditorGUILayout.BeginScrollView(_scroll);

            EditorGUILayout.LabelField("Poly Haven Importer", EditorStyles.boldLabel);
            EditorGUILayout.HelpBox(
                "Pick a texture on polyhaven.com, paste its URL here (e.g. .../a/broken_brick_wall), " +
                "Fetch, choose resolution + maps, then Download. A URP/Lit material is built automatically.",
                MessageType.Info);

            EditorGUILayout.Space();

            DrawBrowser();

            EditorGUILayout.Space();

            using (new EditorGUI.DisabledScope(_busy))
            {
                EditorGUILayout.LabelField("Asset URL or slug", EditorStyles.miniBoldLabel);
                using (new EditorGUILayout.HorizontalScope())
                {
                    _urlInput = EditorGUILayout.TextField(_urlInput);
                    if (GUILayout.Button("Fetch", GUILayout.Width(70)))
                        Fetch();
                }
            }

            if (_info != null)
                DrawAssetSection();

            EditorGUILayout.Space();
            if (!string.IsNullOrEmpty(_status))
                EditorGUILayout.HelpBox(_status, _busy ? MessageType.None : MessageType.Info);

            EditorGUILayout.EndScrollView();
        }

        private void DrawAssetSection()
        {
            EditorGUILayout.Space();
            using (new EditorGUILayout.HorizontalScope())
            {
                if (_thumbnail)
                {
                    var r = GUILayoutUtility.GetRect(96, 96, GUILayout.Width(96), GUILayout.Height(96));
                    GUI.DrawTexture(r, _thumbnail, ScaleMode.ScaleToFit);
                }
                using (new EditorGUILayout.VerticalScope())
                {
                    EditorGUILayout.LabelField(_info.Name, EditorStyles.boldLabel);
                    EditorGUILayout.LabelField($"Slug: {_info.Slug}", EditorStyles.miniLabel);
                    if (_info.Categories != null && _info.Categories.Length > 0)
                        EditorGUILayout.LabelField($"Categories: {string.Join(", ", _info.Categories)}", EditorStyles.miniLabel);
                    if (_info.Type != 1)
                        EditorGUILayout.HelpBox("This asset is not a texture — only textures are supported.", MessageType.Warning);
                }
            }

            if (_files == null || _files.Resolutions.Count == 0)
                return;

            EditorGUILayout.Space();

            // Resolution popup.
            var resArray = _files.Resolutions.ToArray();
            int resIdx = Mathf.Max(0, System.Array.IndexOf(resArray, _selectedRes));
            resIdx = EditorGUILayout.Popup("Resolution", resIdx, resArray.Select(r => r.ToUpperInvariant()).ToArray());
            _selectedRes = resArray[resIdx];

            EditorGUILayout.Space();
            EditorGUILayout.LabelField("Maps", EditorStyles.miniBoldLabel);
            foreach (var meta in PolyHavenMapInfo.All)
            {
                if (!_files.Has(meta.Type))
                    continue;
                bool cur = _mapEnabled.TryGetValue(meta.Type, out var v) ? v : meta.DefaultEnabled;
                _mapEnabled[meta.Type] = EditorGUILayout.ToggleLeft(meta.Label, cur);
            }

            EditorGUILayout.Space();
            bool canDownload = !_busy && _info.Type == 1 && AnyMapSelected();
            using (new EditorGUI.DisabledScope(!canDownload))
            {
                if (GUILayout.Button("Download / Import", GUILayout.Height(30)))
                    Download();
            }
        }

        private bool AnyMapSelected()
        {
            foreach (var meta in PolyHavenMapInfo.All)
                if (_files.Has(meta.Type) && _mapEnabled.TryGetValue(meta.Type, out var v) && v)
                    return true;
            // Fall back to defaults if the user hasn't toggled anything yet.
            foreach (var meta in PolyHavenMapInfo.All)
                if (_files.Has(meta.Type) && meta.DefaultEnabled && !_mapEnabled.ContainsKey(meta.Type))
                    return true;
            return false;
        }

        // ---- Browser -----------------------------------------------------------

        private void DrawBrowser()
        {
            _browseFoldout = EditorGUILayout.Foldout(_browseFoldout, "Browse catalog", true, EditorStyles.foldoutHeader);
            if (!_browseFoldout)
                return;

            using (new EditorGUI.DisabledScope(_busy))
            {
                if (s_catalog == null)
                {
                    // Auto-load once when the browser is first opened; still offer a manual retry.
                    if (!_catalogLoading && !_autoLoadTried)
                    {
                        _autoLoadTried = true;
                        LoadCatalog();
                    }
                    using (new EditorGUILayout.HorizontalScope())
                    {
                        EditorGUILayout.LabelField(_catalogLoading
                            ? "Loading texture catalog…"
                            : "Catalog not loaded.");
                        using (new EditorGUI.DisabledScope(_catalogLoading))
                            if (GUILayout.Button("Load catalog", GUILayout.Width(110)))
                                LoadCatalog();
                    }
                    return;
                }

                using (new EditorGUILayout.HorizontalScope())
                {
                    EditorGUI.BeginChangeCheck();
                    _searchQuery = EditorGUILayout.TextField("Search", _searchQuery);
                    if (EditorGUI.EndChangeCheck())
                        ApplyFilter();
                    if (GUILayout.Button("↻", GUILayout.Width(26)))
                        LoadCatalog();
                }

                int shown = Mathf.Min(_filtered.Count, MaxResults);
                string cap = _filtered.Count > MaxResults ? $", showing first {MaxResults}" : "";
                EditorGUILayout.LabelField(
                    $"{_filtered.Count} match(es){cap} of {s_catalog.Count} textures.",
                    EditorStyles.miniLabel);

                DrawGrid(shown);
            }
        }

        private void DrawGrid(int count)
        {
            const float cell = 104f;
            const float pad = 6f;
            float viewW = EditorGUIUtility.currentViewWidth - 24f;
            int cols = Mathf.Max(1, Mathf.FloorToInt((viewW + pad) / (cell + pad)));

            _gridScroll = EditorGUILayout.BeginScrollView(_gridScroll, GUILayout.Height(320));

            int rows = Mathf.CeilToInt(count / (float)cols);
            for (int row = 0; row < rows; row++)
            {
                using (new EditorGUILayout.HorizontalScope())
                {
                    for (int c = 0; c < cols; c++)
                    {
                        int i = row * cols + c;
                        if (i >= count)
                        {
                            GUILayout.FlexibleSpace();
                            break;
                        }
                        DrawCell(_filtered[i], cell);
                    }
                }
            }

            EditorGUILayout.EndScrollView();
        }

        private void DrawCell(PolyHavenAsset a, float size)
        {
            using (new EditorGUILayout.VerticalScope(GUILayout.Width(size)))
            {
                var thumbRect = GUILayoutUtility.GetRect(size, size - 22f,
                    GUILayout.Width(size), GUILayout.Height(size - 22f));

                var tex = GetThumb(a);
                if (tex)
                    GUI.DrawTexture(thumbRect, tex, ScaleMode.ScaleToFit);
                else
                    EditorGUI.DrawRect(thumbRect, new Color(0f, 0f, 0f, 0.15f));

                // Transparent hit target over the thumbnail.
                if (GUI.Button(thumbRect, new GUIContent("", a.Name), GUIStyle.none))
                    SelectAsset(a);

                EditorGUILayout.LabelField(new GUIContent(a.Name, a.Name),
                    EditorStyles.miniLabel, GUILayout.Width(size));
            }
        }

        /// <summary>Lazy, throttled thumbnail loader — only visible cells request, capped in-flight.</summary>
        private Texture2D GetThumb(PolyHavenAsset a)
        {
            if (_thumbs.TryGetValue(a.Slug, out var tex))
            {
                if (tex)
                    return tex;
                // Cached texture was destroyed (e.g. play-mode cleanup) — purge and re-request.
                _thumbs.Remove(a.Slug);
                _thumbRequested.Remove(a.Slug);
            }

            // Don't add to _thumbRequested until a slot is free, so it retries next repaint.
            if (!_thumbRequested.Contains(a.Slug) && _thumbsInFlight < MaxThumbInFlight)
            {
                _thumbRequested.Add(a.Slug);
                _thumbsInFlight++;
                PolyHavenApi.FetchThumbnail(a.ThumbnailUrl,
                    t => { _thumbs[a.Slug] = t; _thumbsInFlight--; Repaint(); },
                    _ => { _thumbsInFlight--; Repaint(); }); // keep as requested: no retry-spam on error
            }
            return null;
        }

        private void SelectAsset(PolyHavenAsset a)
        {
            _urlInput = "https://polyhaven.com/a/" + a.Slug;
            GUI.FocusControl(null);
            Fetch(); // url filled -> fetch, then everything as before
        }

        private void LoadCatalog()
        {
            _catalogLoading = true;
            _status = "Loading Poly Haven texture catalog…";
            Repaint();

            PolyHavenApi.FetchAssets("textures",
                list =>
                {
                    s_catalog = list;
                    _catalogLoading = false;
                    ApplyFilter();
                    _status = $"Catalog loaded: {list.Count} textures.";
                    Repaint();
                },
                err =>
                {
                    _catalogLoading = false;
                    _status = "Catalog error: " + err;
                    Debug.LogError("[PolyHaven] " + err);
                    Repaint();
                });
        }

        private void ApplyFilter()
        {
            _filtered.Clear();
            if (s_catalog == null)
                return;

            var q = _searchQuery?.Trim().ToLowerInvariant();
            if (string.IsNullOrEmpty(q))
            {
                _filtered.AddRange(s_catalog);
                return;
            }

            var terms = q.Split((char[])null, StringSplitOptions.RemoveEmptyEntries);
            foreach (var a in s_catalog)
            {
                bool all = true;
                foreach (var t in terms)
                    if (a.Search.IndexOf(t, StringComparison.Ordinal) < 0) { all = false; break; }
                if (all)
                    _filtered.Add(a);
            }
        }

        // ---- Fetch -------------------------------------------------------------

        private void Fetch()
        {
            string slug = PolyHavenApi.ExtractSlug(_urlInput);
            if (string.IsNullOrEmpty(slug))
            {
                EditorUtility.DisplayDialog("Poly Haven", "Could not read a slug from the URL.", "OK");
                return;
            }

            _busy = true;
            _status = $"Fetching '{slug}'…";
            _info = null;
            _files = null;
            _thumbnail = null;
            _mapEnabled.Clear();
            Repaint();

            PolyHavenApi.FetchInfo(slug, info =>
            {
                _info = info;
                Repaint();
                PolyHavenApi.FetchThumbnail(info.ThumbnailUrl, tex => { _thumbnail = tex; Repaint(); }, _ => { });

                PolyHavenApi.FetchFiles(slug, files =>
                {
                    _files = files;
                    if (files.Resolutions.Contains(DefaultRes))
                        _selectedRes = DefaultRes;
                    else if (files.Resolutions.Count > 0)
                        _selectedRes = files.Resolutions[Mathf.Min(1, files.Resolutions.Count - 1)];

                    _busy = false;
                    _status = info.Type == 1
                        ? $"Ready: {files.Maps.Count} map type(s), {files.Resolutions.Count} resolution(s)."
                        : "Fetched, but this asset is not a texture.";
                    Repaint();
                }, err => Fail(err));
            }, err => Fail(err));
        }

        // ---- Download ----------------------------------------------------------

        private void Download()
        {
            // Resolve which maps to fetch.
            var wanted = new List<PolyHavenMapType>();
            foreach (var meta in PolyHavenMapInfo.All)
            {
                if (!_files.Has(meta.Type)) continue;
                bool on = _mapEnabled.TryGetValue(meta.Type, out var v) ? v : meta.DefaultEnabled;
                if (on) wanted.Add(meta.Type);
            }

            // ARM fallback: if the asset packs AO/Rough/Metal only into ARM (no dedicated
            // map for a channel), pull ARM so the builder can unpack the missing channels.
            if (_files.Has(PolyHavenMapType.Arm) && !wanted.Contains(PolyHavenMapType.Arm))
            {
                bool channelMissing = !_files.Has(PolyHavenMapType.Metal)
                    || !_files.Has(PolyHavenMapType.Rough)
                    || !_files.Has(PolyHavenMapType.AO);
                if (channelMissing)
                    wanted.Add(PolyHavenMapType.Arm);
            }

            if (wanted.Count == 0)
                return;

            string displayName = _info.Name;
            string folder = $"{RootFolder}/{PolyHavenMaterialBuilder.SanitizeFileName(displayName)}";

            // Folder-collision handling (non-destructive default).
            if (AssetDatabase.IsValidFolder(folder))
            {
                int choice = EditorUtility.DisplayDialogComplex("Poly Haven",
                    $"Folder already exists:\n{folder}", "New folder", "Cancel", "Overwrite");
                if (choice == 1) return;                       // Cancel
                if (choice == 0)                                // New folder
                    folder = AssetDatabase.GenerateUniqueAssetPath(folder);
                // choice == 2 => Overwrite: reuse folder as-is.
            }
            EnsureFolder(folder);

            int maxSize = ResolutionToMaxSize(_selectedRes);
            _busy = true;
            _status = $"Downloading {wanted.Count} map(s) at {_selectedRes.ToUpperInvariant()}…";
            Repaint();

            var results = new List<PolyHavenMaterialBuilder.Downloaded>();
            var failures = new List<string>();
            int remaining = wanted.Count;

            AssetDatabase.StartAssetEditing();

            void FinishOne()
            {
                remaining--;
                if (remaining > 0)
                    return;

                AssetDatabase.StopAssetEditing();
                AssetDatabase.Refresh();

                if (results.Count == 0)
                {
                    _busy = false;
                    _status = "Download failed — no maps saved.";
                    if (failures.Count > 0)
                        Debug.LogError("[PolyHaven] " + string.Join("; ", failures));
                    Repaint();
                    return;
                }

                PolyHavenMaterialBuilder.Build(folder, displayName, _info.Slug, maxSize, results);
                _busy = false;
                _status = $"Done: {results.Count} map(s) imported, material built in {folder}.";
                if (failures.Count > 0)
                    _status += $" ({failures.Count} map(s) failed — see console.)";
                Repaint();
            }

            foreach (var type in wanted)
            {
                if (!_files.TryGetWithFallback(type, _selectedRes, out var file, out var usedRes))
                {
                    failures.Add($"{type}: no file");
                    FinishOne();
                    continue;
                }
                if (usedRes != _selectedRes)
                    Debug.LogWarning($"[PolyHaven] {type}: {_selectedRes} unavailable, using {usedRes}.");

                var meta = PolyHavenMapInfo.Get(type);
                string fileName = $"{_info.Slug}_{meta.Suffix}.{file.Format}";
                string destPath = $"{folder}/{fileName}";

                PolyHavenApi.DownloadFile(file, Path.GetFullPath(destPath),
                    onDone: () =>
                    {
                        results.Add(new PolyHavenMaterialBuilder.Downloaded { Type = type, AssetPath = destPath });
                        FinishOne();
                    },
                    onError: err =>
                    {
                        failures.Add($"{type}: {err}");
                        FinishOne();
                    });
            }
        }

        // ---- Helpers -----------------------------------------------------------

        private void Fail(string message)
        {
            _busy = false;
            _status = "Error: " + message;
            Debug.LogError("[PolyHaven] " + message);
            Repaint();
        }

        private static int ResolutionToMaxSize(string res)
        {
            int k = PolyHavenFiles.ResolutionRank(res);
            int px = k * 1024;
            // Clamp to a valid Unity max texture size (power of two up to 16384).
            int[] valid = { 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384 };
            foreach (var v in valid)
                if (px <= v) return v;
            return 16384;
        }

        private static void EnsureFolder(string folder)
        {
            if (AssetDatabase.IsValidFolder(folder))
                return;
            var parts = folder.Split('/');
            string cur = parts[0]; // "Assets"
            for (int i = 1; i < parts.Length; i++)
            {
                string next = $"{cur}/{parts[i]}";
                if (!AssetDatabase.IsValidFolder(next))
                    AssetDatabase.CreateFolder(cur, parts[i]);
                cur = next;
            }
        }
    }
}