Editor/Graph/FXEditor.cs
using Editor;
using Sandbox;
using System.Collections.Generic;
using System.Linq;

namespace fxbox.Graph;

/// <summary>
/// Main particle system editor - ShaderGraph style with DockManager
/// </summary>
[EditorForAssetType("fx")]
[EditorApp("FXBox", "auto_awesome", "Create and edit particle systems")]
public class FXBoxEditor : DockWindow, IAssetEditor
{
    public bool CanOpenMultipleAssets => true;
    public static ParticleResource CurrentEditingResource { get; private set; }
    
    private ParticleResource _resource; // Original resource from asset
    private ParticleResource _workingCopy; // Copy we actually edit
    private Asset _asset;
    private bool _isDirty = false;

    // UI Components
    private EmitterList _emitterList;
    private ParticlePreview _preview;
    private PropertiesWidget _properties;
    
    private string _defaultDockState;

    public FXBoxEditor()
    {
        DeleteOnClose = true;
        Title = "FXBox - Particle System Editor";
        Size = new Vector2(1800, 1000);
        
        CreateToolBar();
        CreateUI();
        Show();
    }

    private void CreateToolBar()
    {
        var toolbar = new ToolBar(this, "FXBoxToolbar");
        AddToolBar(toolbar, ToolbarPosition.Top);

        toolbar.AddOption("Save", "save", Save).StatusTip = "Save Particle System (Ctrl+S)";
        toolbar.AddSeparator();
        toolbar.AddOption("Add Emitter", "add_circle", AddEmitter).StatusTip = "Add new emitter";
        
        // Add parameter menu
        toolbar.AddSeparator();
        toolbar.AddOption("Float Parameter", "looks_one", AddFloatParameter);
        toolbar.AddOption("Vector Parameter", "3d_rotation", AddVectorParameter);
        toolbar.AddOption("Color Parameter", "palette", AddColorParameter);
        
        toolbar.AddSeparator();
        toolbar.AddOption("Play", "play_arrow", () => _preview?.TogglePlayback()).StatusTip = "Play/Pause";
        toolbar.AddOption("Restart", "replay", () => _preview?.Restart()).StatusTip = "Restart";
    }

    private void CreateUI()
    {
        BuildMenuBar();

        // Register dock types
        DockManager.RegisterDockType("Emitters", "list", () => CreateEmitterList(), false);
        DockManager.RegisterDockType("Preview", "visibility", () => CreatePreview(), false);
        DockManager.RegisterDockType("Properties", "tune", () => CreateProperties(), false);

        // Create dock widgets
        _emitterList = CreateEmitterList();
        _preview = CreatePreview();
        _properties = CreateProperties();

        // Set up dock layout (ShaderGraph style: Left | Center | Right)
        DockManager.AddDock(null, _preview, DockArea.Left, DockManager.DockProperty.HideOnClose);
        DockManager.AddDock(_preview, _emitterList, DockArea.Left, DockManager.DockProperty.HideOnClose, 0.9f);
        DockManager.AddDock(_preview, _properties, DockArea.Right, DockManager.DockProperty.HideOnClose, 0.1f);

        DockManager.Update();
        
        _defaultDockState = DockManager.State;

        // State cookie management
        if (StateCookie != "FXBoxEditor")
        {
            StateCookie = "FXBoxEditor";
        }
        else
        {
            RestoreFromStateCookie();
        }
    }

    private EmitterList CreateEmitterList()
    {
        var widget = new EmitterList(this);
        widget.Name = "Emitters";
        widget.WindowTitle = "Emitters & Modules";
        widget.SetWindowIcon("list");
        widget.MinimumWidth = 300;
        widget.OnSelectionChanged = OnSelectionChanged;
        widget.OnEmitterDeleted = OnEmitterDeleted;
        widget.OnModuleDeleted = OnModuleDeleted;
        widget.OnAddModule = OnAddModule;
        widget.OnSystemChanged = MarkDirty;
        widget.OnEmitterDuplicated = OnEmitterDuplicated;
        widget.OnModuleDuplicated = OnModuleDuplicated;

        return widget;
    }

    private void OnEmitterDuplicated(ParticleEmitter emitter)
    {
	    if (_workingCopy == null) return;

	    var index = _workingCopy.Emitters.IndexOf(emitter);
	    if (index == -1) return;

	    // Round-trip the entire resource through JSON, then pull out
	    // the emitter at the same index — gives us a full deep clone
	    // without needing Serialize/Deserialize on ParticleEmitter itself.
	    var resourceCopy = DeepCopyResource(_workingCopy);
	    var clone = resourceCopy.Emitters[index];
	    clone.Name = $"{emitter.Name} (Copy)";

	    _workingCopy.Emitters.Insert(index + 1, clone);

	    _emitterList?.SetResource(_workingCopy);
	    _preview?.LoadParticleSystem(_workingCopy);
	    MarkDirty();
    }

    private ParticlePreview CreatePreview()
    {
        var widget = new ParticlePreview(this);
        widget.Name = "Preview";
        widget.WindowTitle = "Preview";
      
        widget.SetWindowIcon("visibility");
        return widget;
    }

    private PropertiesWidget CreateProperties()
    {
        var widget = new PropertiesWidget(this);
        widget.Name = "Properties";
        widget.WindowTitle = "Properties";
        widget.SetWindowIcon("tune");
        widget.MinimumWidth = 300;
        widget.OnPropertyChanged = OnPropertyChanged;
        return widget;
    }

    private void BuildMenuBar()
    {
        var file = MenuBar.AddMenu("File");
        file.AddOption("New", "add", New, "editor.new").StatusTip = "New Particle System";
        file.AddOption("Open", "folder_open", Open, "editor.open").StatusTip = "Open Particle System";
        file.AddOption("Save", "save", Save, "editor.save").StatusTip = "Save Particle System";
        file.AddSeparator();
        file.AddOption("Close", null, () => Close(), "editor.quit").StatusTip = "Close Editor";

        var edit = MenuBar.AddMenu("Edit");
        edit.AddOption("Add Emitter", "add_circle", AddEmitter);
        edit.AddSeparator();
        edit.AddOption("Add Float Parameter", "looks_one", AddFloatParameter);
        edit.AddOption("Add Vector Parameter", "3d_rotation", AddVectorParameter);
        edit.AddOption("Add Color Parameter", "palette", AddColorParameter);

        var view = MenuBar.AddMenu("View");
        view.AboutToShow += () => OnViewMenu(view);
    }

    private void OnViewMenu(Menu view)
    {
        view.Clear();
        view.AddOption("Restore To Default", "settings_backup_restore", RestoreDefaultDockLayout);
        view.AddSeparator();

        foreach (var dock in DockManager.DockTypes)
        {
            var o = view.AddOption(dock.Title, dock.Icon);
            o.Checkable = true;
            o.Checked = DockManager.IsDockOpen(dock.Title);
            o.Toggled += (b) => DockManager.SetDockState(dock.Title, b);
        }
    }

    protected override void RestoreDefaultDockLayout()
    {
        DockManager.State = _defaultDockState;
        SaveToStateCookie();
    }

    private void AddFloatParameter()
    {
        if (_workingCopy == null) return;
        
        var param = new FloatParameter
        {
            Name = $"FloatParam{_workingCopy.FloatParameters.Count + 1}",
            DefaultValue = 1.0f
        };
        
        _workingCopy.FloatParameters.Add(param);
        _properties?.ShowSystemProperties(_workingCopy);
        MarkDirty();
    }

    private void AddVectorParameter()
    {
        if (_workingCopy == null) return;
        
        var param = new VectorParameter
        {
            Name = $"VectorParam{_workingCopy.VectorParameters.Count + 1}",
            DefaultValue = Vector3.One
        };
        
        _workingCopy.VectorParameters.Add(param);
        _properties?.ShowSystemProperties(_workingCopy);
        MarkDirty();
    }

    private void AddColorParameter()
    {
        if (_workingCopy == null) return;
        
        var param = new ColorParameter
        {
            Name = $"ColorParam{_workingCopy.ColorParameters.Count + 1}",
            DefaultValue = Color.White
        };
        
        _workingCopy.ColorParameters.Add(param);
        _properties?.ShowSystemProperties(_workingCopy);
        MarkDirty();
    }

    public void AssetOpen(Asset asset)
    {
        _asset = asset;
        _resource = asset.LoadResource<ParticleResource>();

        if (_resource == null)
        {
            _resource = new ParticleResource();
            var defaultEmitter = new ParticleEmitter { Name = "Emitter 1" };
            AddDefaultModules(defaultEmitter);
            _resource.Emitters.Add(defaultEmitter);
        }

        // Create a deep copy for editing
        _workingCopy = DeepCopyResource(_resource);

        Title = $"FXBox - {asset.Name}";
        LoadResource();
        Focus();
    }

    private ParticleResource DeepCopyResource(ParticleResource source)
    {
        if (source == null) return null;

        var json = source.Serialize().ToJsonString();
        var copy = new ParticleResource();
        copy.Deserialize(Json.ParseToJsonObject(json));
        copy.IsDirty = false;
        
        return copy;
    }

    private void LoadResource()
    {
        CurrentEditingResource = _workingCopy;
        
        _emitterList?.SetResource(_workingCopy);
        _preview?.LoadParticleSystem(_workingCopy);
        _properties?.ShowSystemProperties(_workingCopy);
        _isDirty = false;
        if (_asset != null)
            _asset.HasUnsavedChanges = false;
    }

    [Shortcut("editor.new", "CTRL+N", ShortcutType.Window)]
    private void New()
    {
        PromptSave(() => CreateNew());
    }

    private void CreateNew()
    {
        _asset = null;
        _resource = new ParticleResource();
        var defaultEmitter = new ParticleEmitter { Name = "Emitter 1" };
        AddDefaultModules(defaultEmitter);
        _resource.Emitters.Add(defaultEmitter);
        
        _workingCopy = DeepCopyResource(_resource);
        
        Title = "FXBox - Untitled";
        LoadResource();
    }

    [Shortcut("editor.open", "CTRL+O", ShortcutType.Window)]
    private void Open()
    {
        var fd = new FileDialog(null)
        {
            Title = "Open Particle System",
            DefaultSuffix = ".fx"
        };

        fd.SetNameFilter("Particle System (*.fx)");

        if (!fd.Execute())
            return;

        PromptSave(() => OpenFile(fd.SelectedFile));
    }

    private void OpenFile(string path)
    {
        var asset = AssetSystem.FindByPath(path);
        if (asset != null)
        {
            AssetOpen(asset);
        }
    }

    private void AddEmitter()
    {
        if (_workingCopy == null) return;
        
        var emitter = new ParticleEmitter 
        { 
            Name = $"Emitter {_workingCopy.Emitters.Count + 1}"
        };

        AddDefaultModules(emitter);
        _workingCopy.Emitters.Add(emitter);
        
        _emitterList?.SetResource(_workingCopy);
        _preview?.LoadParticleSystem(_workingCopy);
        MarkDirty();
    }

    private void AddDefaultModules(ParticleEmitter emitter)
    {
        emitter.SpawnModules.Add(new SpawnRateModule { Name = "Spawn Rate" });
        emitter.InitializeModules.Add(new InitializePositionModule { Name = "Initialize Position" });
        emitter.InitializeModules.Add(new InitializeVelocityModule { Name = "Initialize Velocity" });
        emitter.InitializeModules.Add(new InitializeLifetimeModule { Name = "Initialize Lifetime" });
        emitter.InitializeModules.Add(new InitializeSizeModule { Name = "Initialize Size" });
        emitter.InitializeModules.Add(new InitializeColorModule { Name = "Initialize Color" });
        emitter.UpdateModules.Add(new GravityForceModule { Name = "Gravity" });
        emitter.UpdateModules.Add(new DragForceModule { Name = "Drag" });
        emitter.RenderModules.Add(new SpriteRendererModule { Name = "Sprite Renderer" });
    }

    private void OnAddModule(ParticleEmitter emitter, ModuleStage stage)
    {
        var menu = new Menu(this);

        var moduleTypes = EditorTypeLibrary.GetTypes<ParticleModule>()
            .Where(t => !t.IsAbstract)
            .Select(t => t.TargetType)
            .Where(t => {
                var instance = System.Activator.CreateInstance(t) as ParticleModule;
                return instance?.Stage == stage;
            });

        foreach (var moduleType in moduleTypes.OrderBy(t => t.Name))
        {
            var displayInfo = DisplayInfo.ForType(moduleType);
            menu.AddOption(displayInfo.Name, displayInfo.Icon ?? "extension", () => {
                var module = System.Activator.CreateInstance(moduleType) as ParticleModule;
                if (module != null)
                {
                    module.Name = displayInfo.Name;
                    
                    var targetList = stage switch
                    {
                        ModuleStage.Spawn => emitter.SpawnModules,
                        ModuleStage.Initialize => emitter.InitializeModules,
                        ModuleStage.Update => emitter.UpdateModules,
                        ModuleStage.Render => emitter.RenderModules,
                        _ => null
                    };

                    targetList?.Add(module);
                    _emitterList?.SetResource(_workingCopy);
                    _preview?.LoadParticleSystem(_workingCopy);
                    MarkDirty();
                }
            });
        }

        menu.OpenAtCursor();
    }

    private void OnSelectionChanged(object target)
    {
        if (target is ParticleResource resource)
        {
            _properties?.ShowSystemProperties(resource);
        }
        else if (target is ParticleEmitter emitter)
        {
            _properties?.ShowEmitterProperties(emitter);
        }
        else if (target is ParticleModule module)
        {
            _properties?.ShowModuleProperties(module);
        }
    }

    private void OnEmitterDeleted(ParticleEmitter emitter)
    {
        _workingCopy.Emitters.Remove(emitter);
        _emitterList?.SetResource(_workingCopy);
        _preview?.LoadParticleSystem(_workingCopy);
        _properties?.ShowSystemProperties(_workingCopy);
        MarkDirty();
    }

    private void OnModuleDeleted(ParticleModule module)
    {
        foreach (var emitter in _workingCopy.Emitters)
        {
            emitter.SpawnModules.Remove(module);
            emitter.InitializeModules.Remove(module);
            emitter.UpdateModules.Remove(module);
            emitter.RenderModules.Remove(module);
        }

        _emitterList?.SetResource(_workingCopy);
        _preview?.LoadParticleSystem(_workingCopy);
        MarkDirty();
    }

    private void OnPropertyChanged()
    {
        _preview?.LoadParticleSystem(_workingCopy);
        MarkDirty();
    }

    private void MarkDirty()
    {
        _preview?.LoadParticleSystem(_workingCopy);
        
        if (!_isDirty)
        {
            if (_workingCopy != null)
                _workingCopy.IsDirty = true;
            if (_resource != null)
                _resource.IsDirty = true;
                
            _isDirty = true;
            
            if (_asset != null)
                _asset.HasUnsavedChanges = true;
                
            Title = $"FXBox - {_asset?.Name ?? "Untitled"}*";
        }
    }

    [Shortcut("editor.save", "CTRL+S", ShortcutType.Window)]
    private void Save()
    {
        if (_asset != null && _workingCopy != null)
        {
            _workingCopy.IsDirty = false;
            _workingCopy.Version++;
            var json = _workingCopy.Serialize().ToJsonString();
            System.IO.File.WriteAllText(_asset.AbsolutePath, json);
            
            _resource = DeepCopyResource(_workingCopy);
            _resource.IsDirty = false;
            
            _isDirty = false;
            _asset.HasUnsavedChanges = false;
            
            Title = $"FXBox - {_asset.Name}";
            Log.Info($"Saved: {_asset.Name}");
        }
        else if (_workingCopy != null)
        {
            // No asset yet, prompt for save location
            SaveAs();
        }
    }

    private void SaveAs()
    {
        var fd = new FileDialog(null)
        {
            Title = "Save Particle System",
            DefaultSuffix = ".fx"
        };

        fd.SelectFile("untitled.fx");
        fd.SetFindFile();
        fd.SetModeSave();
        fd.SetNameFilter("Particle System (*.fx)");
        
        if (!fd.Execute())
            return;

        var savePath = fd.SelectedFile;
        
        _workingCopy.IsDirty = false;
        _workingCopy.Version++;
        var json = _workingCopy.Serialize().ToJsonString();
        System.IO.File.WriteAllText(savePath, json);
        
        _asset = AssetSystem.RegisterFile(savePath);
        _resource = DeepCopyResource(_workingCopy);
        _resource.IsDirty = false;
        
        _isDirty = false;
        if (_asset != null)
            _asset.HasUnsavedChanges = false;
        
        Title = $"FXBox - {_asset.Name}";
        Log.Info($"Saved: {_asset.Name}");
    }

    
    private void OnModuleDuplicated(ParticleModule module, ModuleStage stage)
    {
	    if (_workingCopy == null) return;

	    foreach (var emitter in _workingCopy.Emitters)
	    {
		    var moduleList = stage switch
		    {
			    ModuleStage.Spawn      => emitter.SpawnModules,
			    ModuleStage.Initialize => emitter.InitializeModules,
			    ModuleStage.Update     => emitter.UpdateModules,
			    ModuleStage.Render     => emitter.RenderModules,
			    _                      => null
		    };

		    if (moduleList == null) continue;

		    var index = moduleList.IndexOf(module);
		    if (index == -1) continue;

		    // Round-trip the whole resource; the module will be at the same
		    // emitter index + stage index in the cloned copy.
		    var emitterIndex = _workingCopy.Emitters.IndexOf(emitter);
		    var resourceCopy = DeepCopyResource(_workingCopy);

		    var clonedList = stage switch
		    {
			    ModuleStage.Spawn      => resourceCopy.Emitters[emitterIndex].SpawnModules,
			    ModuleStage.Initialize => resourceCopy.Emitters[emitterIndex].InitializeModules,
			    ModuleStage.Update     => resourceCopy.Emitters[emitterIndex].UpdateModules,
			    ModuleStage.Render     => resourceCopy.Emitters[emitterIndex].RenderModules,
			    _                      => null
		    };

		    if (clonedList == null) break;

		    var clone = clonedList[index];
		    clone.Name = $"{module.Name} (Copy)";
		    moduleList.Insert(index + 1, clone);

		    break; // modules are unique instances, no need to keep iterating
	    }

	    _emitterList?.SetResource(_workingCopy);
	    _preview?.LoadParticleSystem(_workingCopy);
	    MarkDirty();
    }

    private void PromptSave(System.Action action)
    {
        if (!_isDirty)
        {
            action?.Invoke();
            return;
        }

        var confirm = new PopupWindow(
            "Save Current Particle System", 
            "The particle system has unsaved changes. Would you like to save now?", 
            "Cancel",
            new Dictionary<string, System.Action>()
            {
                { "No", () => { _isDirty = false; action?.Invoke(); } },
                { "Yes", () => { Save(); if (!_isDirty) action?.Invoke(); } }
            }
        );

        confirm.Show();
    }

    public void SelectMember(string memberName) { }

    protected override bool OnClose()
    {
        CurrentEditingResource = null;
        
        if (_isDirty)
        {
            PromptSave(() => { _isDirty = false; Close(); });
            return false;
        }

        return base.OnClose();
    }
}

/// <summary>
/// Left panel showing emitters and their modules in a tree structure
/// </summary>
public class EmitterList : Widget
{
    private TreeView _tree;
    private ParticleResource _resource;

    public System.Action<object> OnSelectionChanged;
    public System.Action<ParticleEmitter> OnEmitterDeleted;
    public System.Action<ParticleModule> OnModuleDeleted;
    public System.Action<ParticleEmitter, ModuleStage> OnAddModule;
    public System.Action<ParticleEmitter> OnEmitterDuplicated;
    public System.Action<ParticleModule, ModuleStage> OnModuleDuplicated;
    public System.Action OnSystemChanged;
    
    public EmitterList(Widget parent) : base(parent)
    {
        Layout = Layout.Column();
        Layout.Spacing = 0;

        // Header
        var header = new Widget(this);
        header.Layout = Layout.Row();
        header.Layout.Spacing = 8;
        header.Layout.Margin = 8;
        header.MinimumHeight = 32;
        header.SetStyles("background-color: #1e1e1e;");

        var label = new Label("Emitters & Modules", header);
        label.SetStyles("font-weight: bold; font-size: 14px;");
        header.Layout.Add(label, 1);

        Layout.Add(header);

        _tree = new TreeView(this);
        _tree.AcceptDrops = true;
        _tree.ItemClicked = OnItemActivated;
        _tree.ItemContextMenu = OnItemContextMenu;
        _tree.ItemSelected = OnTreeMousePress;
        Layout.Add(_tree, 1);
    }

    private void OnTreeMousePress(object item)
    {
        TryHandleStageButtonClick(item);
    }

    private bool TryHandleStageButtonClick(object item)
    {
        if (item is not StageNode stageNode)
            return false;

        if (!_tree.TryGetItemRect(stageNode, out var itemRect))
            return false;

        var buttonRect = new Rect(itemRect.Right - 24, itemRect.Top, 24, itemRect.Height);
        
        if (itemRect.IsInside(buttonRect))
        {
            OnAddModule?.Invoke(stageNode.Emitter, stageNode.Stage);
            return true;
        }

        return false;
    }

    public void SetResource(ParticleResource resource)
    {
        _resource = resource;
        RebuildTree();
    }

    private void RebuildTree()
    {
        _tree.Clear();

        if (_resource == null) return;

        var systemNode = new SystemNode(_resource);
        _tree.AddItem(systemNode);
        _tree.Open(systemNode);

        foreach (var emitter in _resource.Emitters)
        {
            var emitterNode = new EmitterNode(emitter);
            systemNode.AddItem(emitterNode);
            _tree.Open(emitterNode);

            AddStageNode(emitterNode, "Spawn", ModuleStage.Spawn, emitter.SpawnModules, emitter);
            AddStageNode(emitterNode, "Initialize", ModuleStage.Initialize, emitter.InitializeModules, emitter);
            AddStageNode(emitterNode, "Update", ModuleStage.Update, emitter.UpdateModules, emitter);
            AddStageNode(emitterNode, "Render", ModuleStage.Render, emitter.RenderModules, emitter);
        }
    }

    private void AddStageNode(TreeNode parent, string name, ModuleStage stage, List<ParticleModule> modules, ParticleEmitter emitter)
    {
        var stageNode = new StageNode(name, stage, emitter, modules.Count);
        parent.AddItem(stageNode);
        _tree.Open(stageNode);

        foreach (var module in modules)
        {
            var moduleNode = new ModuleNode(module, stage);
            stageNode.AddItem(moduleNode);
        }
    }

    private void OnItemActivated(object item)
    {
        var selected = item;
        if (selected == null) return;
        
        if (selected is SystemNode systemNode)
            OnSelectionChanged?.Invoke(systemNode.Resource);
        else if (selected is EmitterNode emitterNode)
            OnSelectionChanged?.Invoke(emitterNode.Emitter);
        else if (selected is ModuleNode moduleNode)
            OnSelectionChanged?.Invoke(moduleNode.Module);
    }

    private void OnItemContextMenu(object item)
    {
        var selected = item;
        if (selected == null) return;

        var menu = new Menu(this);

        if (selected is EmitterNode emitterNode)
        {
	        menu.AddOption("Duplicate", "content_copy", () => OnEmitterDuplicated?.Invoke(emitterNode.Emitter));
	        menu.AddSeparator();
	        menu.AddOption("Delete Emitter", "delete", () => OnEmitterDeleted?.Invoke(emitterNode.Emitter));
        }
        else if (selected is ModuleNode moduleNode)
        {
	        var module = moduleNode.Module;

	        menu.AddOption("Duplicate", "content_copy", () => OnModuleDuplicated?.Invoke(module, moduleNode.Stage));

	        menu.AddOption(module.Enabled ? "Disable" : "Enable",
		        module.Enabled ? "visibility_off" : "visibility",
		        () => {
			        module.Enabled = !module.Enabled;
			        OnSystemChanged?.Invoke();
			        RebuildTree();
		        });

	        menu.AddSeparator();
	        menu.AddOption("Delete Module", "delete", () => OnModuleDeleted?.Invoke(module));
        }
        else if (selected is StageNode stageNode)
        {
            menu.AddOption("Add Module", "add", () => OnAddModule?.Invoke(stageNode.Emitter, stageNode.Stage));
        }

        menu.OpenAtCursor();
    }

    // TreeNode classes
    private class SystemNode : TreeNode
    {
        public ParticleResource Resource { get; }
        
        public SystemNode(ParticleResource resource)
        {
            Resource = resource;
        }

        public override void OnPaint(VirtualWidget item)
        {
            PaintSelection(item);
            
            var rect = item.Rect.Shrink(4, 2);
            Paint.SetDefaultFont();
            Paint.SetPen(Theme.Text);
            Paint.DrawIcon(rect, "auto_awesome", 16, TextFlag.LeftCenter);
            
            rect.Left += 24;
            Paint.DrawText(rect, "Particle System", TextFlag.LeftCenter);
        }
    }

    private class EmitterNode : TreeNode
    {
        public ParticleEmitter Emitter { get; }
        
        public EmitterNode(ParticleEmitter emitter)
        {
            Emitter = emitter;
        }

        public override void OnPaint(VirtualWidget item)
        {
            PaintSelection(item);
            
            var rect = item.Rect.Shrink(4, 2);
            Paint.SetDefaultFont();
            Paint.SetPen(Theme.Green);
            Paint.DrawText(rect, "● ", TextFlag.LeftCenter);
            
            rect.Left += 20;
            Paint.SetPen(Theme.Text);
            Paint.DrawText(rect, Emitter.Name, TextFlag.LeftCenter);
        }
    }

    private class StageNode : TreeNode
    {
        public new string Name { get; }
        public ModuleStage Stage { get; }
        public ParticleEmitter Emitter { get; }
        public int Count { get; }
        
        public StageNode(string name, ModuleStage stage, ParticleEmitter emitter, int count)
        {
            Name = name;
            Stage = stage;
            Emitter = emitter;
            Count = count;
        }

        public override void OnPaint(VirtualWidget item)
        {
            PaintSelection(item);
            
            var rect = item.Rect.Shrink(4, 2);
            Paint.SetDefaultFont();
            
            var stageColor = Stage switch
            {
                ModuleStage.Spawn => new Color(1f, 0.6f, 0.2f),
                ModuleStage.Initialize => new Color(0.3f, 0.8f, 0.3f),
                ModuleStage.Update => new Color(0.3f, 0.6f, 1f),
                ModuleStage.Render => new Color(0.9f, 0.3f, 0.9f),
                _ => Theme.Text
            };
            
            Paint.SetPen(stageColor);
            Paint.DrawText(rect, $"{Name} ({Count})", TextFlag.LeftCenter);
            
            var buttonRect = new Rect(rect.Right - 24, rect.Top, 24, rect.Height);
            
            if (buttonRect.IsInside(item.Rect))
            {
                Paint.ClearPen();
                Paint.SetBrush(Theme.ControlBackground.Lighten(0.2f));
                Paint.DrawRect(buttonRect.Shrink(2), 2);
            }
            
            Paint.SetPen(stageColor);
            Paint.DrawIcon(buttonRect, "add", 16, TextFlag.Center);
            
            if (item.Dropping)
            {
                Paint.ClearPen();
                Paint.SetBrush(Theme.Blue.WithAlpha(0.2f));
                Paint.DrawRect(item.Rect, 2);
            }
        }

        public override DropAction OnDragDrop(BaseItemWidget.ItemDragEvent e)
        {
            if (e.Data.Object is not ModuleNode draggedNode)
                return DropAction.Ignore;

            var draggedModule = draggedNode.Module;

            if (draggedNode.Stage != Stage)
                return DropAction.Ignore;

            if (e.IsDrop)
            {
                var targetList = Stage switch
                {
                    ModuleStage.Spawn => Emitter.SpawnModules,
                    ModuleStage.Initialize => Emitter.InitializeModules,
                    ModuleStage.Update => Emitter.UpdateModules,
                    ModuleStage.Render => Emitter.RenderModules,
                    _ => null
                };
                
                if (targetList == null) return DropAction.Ignore;

                targetList.Remove(draggedModule);
                targetList.Add(draggedModule);
                
                if (TreeView.Parent is EmitterList list)
                {
                    list.OnSystemChanged?.Invoke();
                    list.RebuildTree();
                }
            }

            return DropAction.Move;
        }
    }

    private class ModuleNode : TreeNode
    {
        public ParticleModule Module { get; }
        public ModuleStage Stage { get; }
        
        public ModuleNode(ParticleModule module, ModuleStage stage)
        {
            Module = module;
            Stage = stage;
        }

        public override void OnPaint(VirtualWidget item)
        {
            PaintSelection(item);
            
            var displayInfo = DisplayInfo.ForType(Module.GetType());
            var rect = item.Rect.Shrink(4, 2);
            
            Paint.SetDefaultFont();
            Paint.SetPen(Module.Enabled ? Theme.Text : Theme.Text.WithAlpha(0.5f));
            
            if (!string.IsNullOrEmpty(displayInfo.Icon))
            {
                Paint.DrawIcon(rect, displayInfo.Icon, 16, TextFlag.LeftCenter);
                rect.Left += 24;
            }
            
            Paint.DrawText(rect, Module.Name ?? displayInfo.Name, TextFlag.LeftCenter);
            
            if (item.Dropping)
            {
                Paint.ClearPen();
                Paint.SetBrush(Theme.Blue.WithAlpha(0.2f));
                
                if (TreeView.CurrentItemDragEvent.DropEdge.HasFlag(BaseItemWidget.ItemEdge.Top))
                {
                    var droprect = item.Rect;
                    droprect.Top -= 1;
                    droprect.Height = 2;
                    Paint.DrawRect(droprect, 2);
                }
                else if (TreeView.CurrentItemDragEvent.DropEdge.HasFlag(BaseItemWidget.ItemEdge.Bottom))
                {
                    var droprect = item.Rect;
                    droprect.Top = droprect.Bottom - 1;
                    droprect.Height = 2;
                    Paint.DrawRect(droprect, 2);
                }
                else
                {
                    Paint.DrawRect(item.Rect, 2);
                }
            }
        }

        public override bool OnDragStart()
        {
            var drag = new Drag(TreeView);
            drag.Data.Object = this;
            drag.Execute();
            return true;
        }

        public override DropAction OnDragDrop(BaseItemWidget.ItemDragEvent e)
        {
            if (e.Data.Object is not ModuleNode draggedNode)
                return DropAction.Ignore;

            var draggedModule = draggedNode.Module;
            var targetModule = Module;

            if (draggedNode.Stage != Stage)
                return DropAction.Ignore;

            var emitterList = TreeView.Parent as EmitterList;
            if (emitterList?._resource == null) return DropAction.Ignore;

            foreach (var emitter in emitterList._resource.Emitters)
            {
                var moduleList = Stage switch
                {
                    ModuleStage.Spawn => emitter.SpawnModules,
                    ModuleStage.Initialize => emitter.InitializeModules,
                    ModuleStage.Update => emitter.UpdateModules,
                    ModuleStage.Render => emitter.RenderModules,
                    _ => null
                };

                if (moduleList == null) continue;

                var draggedIndex = moduleList.IndexOf(draggedModule);
                var targetIndex = moduleList.IndexOf(targetModule);

                if (draggedIndex == -1 || targetIndex == -1) continue;

                if (e.IsDrop)
                {
                    moduleList.RemoveAt(draggedIndex);
                    
                    if (draggedIndex < targetIndex)
                        targetIndex--;
                    
                    if (e.DropEdge.HasFlag(BaseItemWidget.ItemEdge.Top))
                    {
                        moduleList.Insert(targetIndex, draggedModule);
                    }
                    else if (e.DropEdge.HasFlag(BaseItemWidget.ItemEdge.Bottom))
                    {
                        moduleList.Insert(targetIndex + 1, draggedModule);
                    }
                    else
                    {
                        moduleList.Insert(targetIndex, draggedModule);
                    }

                    emitterList.OnSystemChanged?.Invoke();
                    emitterList.RebuildTree();
                }
                
                return DropAction.Move;
            }

            return DropAction.Ignore;
        }
    }
}

/// <summary>
/// Right panel showing properties - ShaderGraph style
/// </summary>
public class PropertiesWidget : Widget
{
    private Label _titleLabel;
    private Layout _contentLayout;
    private object _currentTarget;
    private SerializedObject _currentSerializedObject;

    public System.Action OnPropertyChanged;

    public PropertiesWidget(Widget parent) : base(parent)
    {
        Layout = Layout.Column();
        Layout.Spacing = 0;

        var header = new Widget(this);
        header.Layout = Layout.Column();
        header.Layout.Margin = 8;
        header.MinimumHeight = 32;
        header.SetStyles("background-color: #1e1e1e;");

        _titleLabel = new Label("Properties", header);
        _titleLabel.SetStyles("font-weight: bold; font-size: 14px;");
        header.Layout.Add(_titleLabel);

        Layout.Add(header);

        _contentLayout = Layout.AddColumn(1);
    }

    public void ShowSystemProperties(ParticleResource resource)
    {
        _currentTarget = resource;
        _titleLabel.Text = "System Properties";
        
        RebuildContent(() => {
            if (resource != null)
            {
                var so = resource.GetSerialized();
                so.OnPropertyChanged += OnSerializedPropertyChanged;
                return so;
            }
            return null;
        });
    }

    public void ShowEmitterProperties(ParticleEmitter emitter)
    {
        _currentTarget = emitter;
        _titleLabel.Text = $"Emitter: {emitter.Name}";
        
        RebuildContent(() => {
            if (emitter != null)
            {
                var so = emitter.GetSerialized();
                so.OnPropertyChanged += OnSerializedPropertyChanged;
                return so;
            }
            return null;
        });
    }

    public void ShowModuleProperties(ParticleModule module)
    {
        _currentTarget = module;
        
        var displayInfo = DisplayInfo.ForType(module.GetType());
        _titleLabel.Text = displayInfo.Name;
        
        RebuildContent(() => {
            if (module != null)
            {
                var so = module.GetSerialized();
                so.OnPropertyChanged += OnSerializedPropertyChanged;
                return so;
            }
            return null;
        });
    }

    private void OnSerializedPropertyChanged(SerializedProperty property)
    {
        if (_currentSerializedObject != null && _currentTarget != null)
        {
            _currentSerializedObject.NoteFinishEdit(property);
        }
        
        OnPropertyChanged?.Invoke();
    }

    private void RebuildContent(System.Func<SerializedObject> getSerializedObject)
    {
        _contentLayout.Clear(true);

        if (_currentSerializedObject != null)
        {
            _currentSerializedObject.OnPropertyChanged -= OnSerializedPropertyChanged;
        }

        var scroll = new ScrollArea(this);
        scroll.Canvas = new Widget(scroll);
        scroll.Canvas.Layout = Layout.Column();
        scroll.Canvas.Layout.Margin = 8;
        scroll.Canvas.Layout.Spacing = 4;
        scroll.Canvas.VerticalSizeMode = SizeMode.CanGrow;
        scroll.Canvas.HorizontalSizeMode = SizeMode.Flexible;

        var so = getSerializedObject();
        if (so != null)
        {
            _currentSerializedObject = so;
            
            var sheet = new ControlSheet();
            sheet.AddObject(so);
            
            scroll.Canvas.Layout.Add(sheet);
            scroll.Canvas.Layout.AddStretchCell();
        }
        else
        {
            _currentSerializedObject = null;
        }

        _contentLayout.Add(scroll);
    }
}