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

using System.Collections.Generic;
using System.IO;
using System.Linq;
using Chameleon.Editor.WebAssets;
using UnityEditor;
using UnityEngine;

namespace Chameleon.Editor.PolyHaven
{
    /// <summary>
    /// Editor window: browse/paste a Poly Haven asset (texture or model), fetch its files,
    /// pick resolution + maps, then download into Assets/Poly Haven/{Name}/. Textures get an
    /// auto-built URP/Lit material; models additionally download the FBX and remap its
    /// material slots to the built material(s) — one per texture set.
    /// </summary>
    public sealed class PolyHavenImporterWindow : EditorWindow
    {
        private const string RootFolder = "Assets/Poly Haven";
        private const string DefaultRes = "2k";

        private const int TypeTexture = 1;
        private const int TypeModel = 2;

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

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

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

        private static readonly string[] TabLabels = { "Textures", "Models" };
        private static readonly string[] TabTypes = { "textures", "models" };

        // Catalogs cached per type for the whole session (survive window close,
        // cleared on domain reload).
        private static readonly Dictionary<string, List<AssetBrowserGrid.Entry>> s_catalogs =
            new Dictionary<string, List<AssetBrowserGrid.Entry>>();

        private AssetBrowserGrid _grid;
        private bool _browseFoldout = true;
        private int _catalogTab;
        private readonly HashSet<string> _autoLoadTried = new HashSet<string>();
        private bool _catalogLoading;

        private string CatalogType => TabTypes[_catalogTab];

        [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()
        {
            _grid = new AssetBrowserGrid(Repaint);

            // Catalog may already be cached from a previous window instance this session.
            if (s_catalogs.TryGetValue(CatalogType, out var cached))
            {
                _autoLoadTried.Add(CatalogType);
                _grid.SetCatalog(cached);
            }
        }

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

            EditorGUILayout.LabelField("Poly Haven Importer", EditorStyles.boldLabel);
            EditorGUILayout.HelpBox(
                "Pick a texture or model 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; " +
                "models also import the FBX with the material(s) wired up.",
                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 bool IsSupportedType => _info != null && (_info.Type == TypeTexture || _info.Type == TypeModel);
        private bool IsModel => _info != null && _info.Type == TypeModel;

        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}   ({(IsModel ? "Model" : _info.Type == TypeTexture ? "Texture" : "Unsupported")})",
                        EditorStyles.miniLabel);
                    if (_info.Categories != null && _info.Categories.Length > 0)
                        EditorGUILayout.LabelField($"Categories: {string.Join(", ", _info.Categories)}", EditorStyles.miniLabel);
                    if (!IsSupportedType)
                        EditorGUILayout.HelpBox("Only textures and models are supported (no HDRIs).", MessageType.Warning);
                    if (IsModel && _files != null && !_files.HasFbx)
                        EditorGUILayout.HelpBox("This model has no FBX download — cannot import.", 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();
            string mapsLabel = "Maps";
            if (_files.Sets.Count > 1)
                mapsLabel += $" ({_files.Sets.Count} texture sets: {string.Join(", ", _files.Sets.Select(s => s.Name))})";
            EditorGUILayout.LabelField(mapsLabel, EditorStyles.miniBoldLabel);
            foreach (var meta in PbrMapInfo.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 && IsSupportedType && AnyMapSelected()
                               && (!IsModel || _files.HasFbx);
            using (new EditorGUI.DisabledScope(!canDownload))
            {
                if (GUILayout.Button(IsModel ? "Download / Import Model" : "Download / Import", GUILayout.Height(30)))
                    Download();
            }
        }

        private bool AnyMapSelected()
        {
            foreach (var meta in PbrMapInfo.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 PbrMapInfo.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))
            {
                using (new EditorGUI.DisabledScope(_catalogLoading))
                {
                    int newTab = GUILayout.Toolbar(_catalogTab, TabLabels);
                    if (newTab != _catalogTab)
                    {
                        _catalogTab = newTab;
                        _grid.SetCatalog(s_catalogs.TryGetValue(CatalogType, out var cached) ? cached : null);
                    }
                }

                if (!_grid.HasCatalog)
                {
                    // Auto-load once per type when first shown; still offer a manual retry.
                    if (!_catalogLoading && !_autoLoadTried.Contains(CatalogType))
                    {
                        _autoLoadTried.Add(CatalogType);
                        LoadCatalog();
                    }
                    using (new EditorGUILayout.HorizontalScope())
                    {
                        EditorGUILayout.LabelField(_catalogLoading
                            ? $"Loading {CatalogType} catalog…"
                            : "Catalog not loaded.");
                        using (new EditorGUI.DisabledScope(_catalogLoading))
                            if (GUILayout.Button("Load catalog", GUILayout.Width(110)))
                                LoadCatalog();
                    }
                    return;
                }

                _grid.Draw(SelectAsset, LoadCatalog, CatalogType);
            }
        }

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

        private void LoadCatalog()
        {
            string type = CatalogType;
            _catalogLoading = true;
            _status = $"Loading Poly Haven {type} catalog…";
            Repaint();

            PolyHavenApi.FetchAssets(type,
                list =>
                {
                    s_catalogs[type] = list;
                    _catalogLoading = false;
                    if (type == CatalogType)
                        _grid.SetCatalog(list);
                    _status = $"Catalog loaded: {list.Count} {type}.";
                    Repaint();
                },
                err =>
                {
                    _catalogLoading = false;
                    _status = "Catalog error: " + err;
                    Debug.LogError("[PolyHaven] " + err);
                    Repaint();
                });
        }

        // ---- 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();
                EditorWebRequest.GetTexture(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 = IsSupportedType
                        ? $"Ready: {files.MapTypeCount} map type(s), {files.Resolutions.Count} resolution(s)" +
                          (files.Sets.Count > 1 ? $", {files.Sets.Count} texture sets" : "") +
                          (IsModel ? (files.HasFbx ? ", FBX available." : ", but no FBX!") : ".")
                        : "Fetched, but this asset type is not supported (textures/models only).";
                    Repaint();
                }, err => Fail(err));
            }, err => Fail(err));
        }

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

        /// <summary>One pending file download: which set/map it belongs to (Type null for the FBX).</summary>
        private struct Pending
        {
            public PolyHavenMapSet Set;
            public PbrMapType Type;
            public bool IsFbx;
            public PolyHavenFile File;
            public string DestPath;
        }

        private void Download()
        {
            // Resolve which map types to fetch (per files-level availability).
            var wantedTypes = new List<PbrMapType>();
            foreach (var meta in PbrMapInfo.All)
            {
                if (!_files.Has(meta.Type)) continue;
                bool on = _mapEnabled.TryGetValue(meta.Type, out var v) ? v : meta.DefaultEnabled;
                if (on) wantedTypes.Add(meta.Type);
            }
            if (wantedTypes.Count == 0)
                return;

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

            folder = ImportUtil.ResolveFolderCollision(folder, "Poly Haven");
            if (folder == null)
                return;
            ImportUtil.EnsureFolder(folder);

            // Build the per-set download list. ARM fallback per set: if a set packs
            // AO/Rough/Metal only into ARM, pull ARM so the builder can unpack it.
            var pending = new List<Pending>();
            foreach (var set in _files.Sets)
            {
                var setTypes = new List<PbrMapType>();
                foreach (var t in wantedTypes)
                    if (set.Has(t))
                        setTypes.Add(t);

                if (set.Has(PbrMapType.Arm) && !setTypes.Contains(PbrMapType.Arm))
                {
                    bool channelMissing = !set.Has(PbrMapType.Metal)
                        || !set.Has(PbrMapType.Rough)
                        || !set.Has(PbrMapType.AO);
                    if (channelMissing)
                        setTypes.Add(PbrMapType.Arm);
                }

                string setPart = string.IsNullOrEmpty(set.Name)
                    ? ""
                    : "_" + ImportUtil.SanitizeFileName(set.Name);

                foreach (var type in setTypes)
                {
                    if (!set.TryGetWithFallback(type, _selectedRes, out var file, out var usedRes))
                        continue;
                    if (usedRes != _selectedRes)
                        Debug.LogWarning($"[PolyHaven] {set.Name}/{type}: {_selectedRes} unavailable, using {usedRes}.");

                    var meta = PbrMapInfo.Get(type);
                    pending.Add(new Pending
                    {
                        Set = set,
                        Type = type,
                        File = file,
                        DestPath = $"{folder}/{_info.Slug}{setPart}_{meta.Suffix}.{file.Format}",
                    });
                }
            }

            // Model: the FBX geometry itself.
            string fbxPath = null;
            if (IsModel)
            {
                if (!_files.TryGetFbxWithFallback(_selectedRes, out var fbx, out var fbxRes))
                {
                    _status = "No FBX file available for this model.";
                    return;
                }
                if (fbxRes != _selectedRes)
                    Debug.LogWarning($"[PolyHaven] FBX: {_selectedRes} unavailable, using {fbxRes}.");
                fbxPath = $"{folder}/{_info.Slug}_{fbxRes}.fbx";
                pending.Add(new Pending { IsFbx = true, File = fbx, DestPath = fbxPath });
            }

            if (pending.Count == 0)
                return;

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

            var resultsBySet = new Dictionary<PolyHavenMapSet, List<PbrMaterialBuilder.Downloaded>>();
            var failures = new List<string>();
            bool fbxOk = false;
            int remaining = pending.Count;

            AssetDatabase.StartAssetEditing();

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

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

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

                // One material per texture set (single-set assets keep the plain name).
                var setMaterials = new List<(string setName, Material mat)>();
                foreach (var set in _files.Sets)
                {
                    if (!resultsBySet.TryGetValue(set, out var maps) || maps.Count == 0)
                        continue;
                    bool named = !string.IsNullOrEmpty(set.Name) && _files.Sets.Count > 1;
                    string matName = named ? $"{displayName} {set.Name}" : displayName;
                    string setSlug = named ? $"{_info.Slug}_{ImportUtil.SanitizeFileName(set.Name)}" : _info.Slug;
                    var mat = PbrMaterialBuilder.Build(folder, matName, setSlug, maxSize, maps);
                    if (mat)
                        setMaterials.Add((set.Name, mat));
                }

                // Model: wire the built materials into the imported FBX.
                if (IsModel && fbxOk && setMaterials.Count > 0)
                    ModelImportUtil.RemapMaterials(fbxPath, setMaterials);

                _busy = false;
                _status = $"Done: {mapCount} map(s) imported, {setMaterials.Count} material(s) built in {folder}"
                          + (IsModel ? (fbxOk ? ", FBX imported." : ", but the FBX download FAILED.") : ".");
                if (failures.Count > 0)
                    _status += $" ({failures.Count} file(s) failed — see console.)";
                Repaint();
            }

            foreach (var p in pending)
            {
                var item = p; // capture per iteration
                EditorWebRequest.DownloadFile(item.File.Url, Path.GetFullPath(item.DestPath),
                    onDone: () =>
                    {
                        if (item.IsFbx)
                        {
                            fbxOk = true;
                        }
                        else
                        {
                            if (!resultsBySet.TryGetValue(item.Set, out var list))
                                resultsBySet[item.Set] = list = new List<PbrMaterialBuilder.Downloaded>();
                            list.Add(new PbrMaterialBuilder.Downloaded { Type = item.Type, AssetPath = item.DestPath });
                        }
                        FinishOne();
                    },
                    onError: err =>
                    {
                        failures.Add($"{(item.IsFbx ? "FBX" : item.Type.ToString())}: {err}");
                        FinishOne();
                    });
            }
        }

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

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