using System;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEditor.Experimental.GraphView;
using UnityEditor.UIElements;
using UnityEngine;
using UnityEngine.UIElements;
using SearchWindow = UnityEditor.Experimental.GraphView.SearchWindow;

public interface IFoldNode
{
    FoldNodeSerialized Serialize();
    IFoldNode GetPreviousFold();
}

public abstract class FoldNodeView : Node
{
    public FoldNodeData Data { get; private set; }
    protected FoldGraphView graphView;
    readonly Dictionary<string, Port> ports = new();

    public virtual void Initialize(FoldGraphView view, FoldNodeData data)
    {
        graphView = view;
        Data = data;
        title = data.Title;
        viewDataKey = data.guid;
        SetPosition(new Rect(data.position, Vector2.zero));

        RegisterCallback<MouseDownEvent>(evt =>
        {
            if (!evt.shiftKey) return;
            if (!graphView.selection.Contains(this))
                graphView.AddToSelection(this);
            evt.StopImmediatePropagation();
            evt.PreventDefault();
        }, TrickleDown.TrickleDown);
    }

    public override void SetPosition(Rect newPos)
    {
        base.SetPosition(newPos);
        if (Data != null)
        {
            graphView.RecordUndo("Move Node");
            Data.position = newPos.position;
            graphView.MarkDirty();
        }
    }

    protected Port AddPort(Direction direction, Port.Capacity capacity, string name, Type type)
    {
        var port = InstantiatePort(Orientation.Horizontal, direction, capacity, type);
        port.portName = name;
        ports[name] = port;
        graphView.BindEdgeConnector(port);
        (direction == Direction.Input ? inputContainer : outputContainer).Add(port);
        return port;
    }

    public Port GetPort(string name) => ports.TryGetValue(name, out var p) ? p : null;
}

public class ValueNodeView : FoldNodeView
{
    object GetValue() => Data switch
    {
        FloatValueNodeData f => f.value,
        VectorValueNodeData v => v.value,
        GameObjectNodeData g => g.output,
        _ => null
    };

    public T GetTypedValue<T>() => GetValue() is T v ? v : default;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);
        var portType = data switch
        {
            FloatValueNodeData => typeof(float),
            VectorValueNodeData => typeof(Vector4),
            GameObjectNodeData => typeof(GameObject),
            _ => typeof(object)
        };
        AddPort(Direction.Output, Port.Capacity.Multi, "Out", portType);

        VisualElement field = data switch
        {
            FloatValueNodeData f => CreateFloatField(f),
            VectorValueNodeData v => CreateVectorField(v),
            GameObjectNodeData g => CreateObjectField(g),
            _ => null
        };
        if (field != null) extensionContainer.Add(field);
        RefreshExpandedState();
    }

    FloatField CreateFloatField(FloatValueNodeData d)
    {
        var f = new FloatField("Value") { value = d.value };
        f.RegisterValueChangedCallback(e => { graphView.RecordUndo("Edit Float"); d.value = e.newValue; graphView.MarkDirty(); });
        return f;
    }

    Vector4Field CreateVectorField(VectorValueNodeData d)
    {
        var f = new Vector4Field("Value") { value = d.value };
        f.RegisterValueChangedCallback(e => { graphView.RecordUndo("Edit Vector"); d.value = e.newValue; graphView.MarkDirty(); });
        return f;
    }

    ObjectField CreateObjectField(GameObjectNodeData d)
    {
        var f = new ObjectField("Object") { allowSceneObjects = true, objectType = typeof(GameObject), value = d.output };
        f.RegisterValueChangedCallback(e => { graphView.RecordUndo("Edit GameObject"); d.output = e.newValue as GameObject; graphView.MarkDirty(); });
        return f;
    }
}

public class OperationNodeView : FoldNodeView, IFoldNode
{
    Port foldInputPort;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);
        foldInputPort = AddPort(Direction.Input, Port.Capacity.Single, "Input", typeof(IFoldNode));
        AddPort(Direction.Output, Port.Capacity.Single, "Out", typeof(IFoldNode));

        foreach (var (name, type) in ((FoldOperationData)data).GetInputPorts())
            AddPort(Direction.Input, Port.Capacity.Single, name, type);
    }

    public IFoldNode GetPreviousFold() => foldInputPort.connections.FirstOrDefault()?.output?.node as IFoldNode;
    public FoldNodeSerialized Serialize() => ((FoldOperationData)Data).Serialize(graphView);
}

public class KeyframeNodeView : FoldNodeView
{
    Port gameObjectPort, foldPort;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);
        gameObjectPort = AddPort(Direction.Input, Port.Capacity.Single, "GameObject", typeof(GameObject));
        foldPort = AddPort(Direction.Input, Port.Capacity.Single, "Fold Data", typeof(IFoldNode));
        extensionContainer.Add(new Button(OnGenerateClick) { text = "Generate Keyframe" });
        RefreshExpandedState();
    }

    void OnGenerateClick()
    {
        var goEdge = gameObjectPort.connections.FirstOrDefault();
        var go = (goEdge?.output?.node as ValueNodeView)?.GetTypedValue<GameObject>();
        if (go == null) { graphView.ShowNotification("GameObject input missing or empty."); return; }

        var renderer = go.GetComponent<MeshRenderer>();
        if (renderer == null) { graphView.ShowNotification("GameObject must have a MeshRenderer component."); return; }

        var foldNode = foldPort.connections.FirstOrDefault()?.output?.node as IFoldNode;
        if (foldNode == null) { graphView.ShowNotification("Fold input not connected."); return; }

        var stack = new Stack<IFoldNode>();
        for (var cur = foldNode; cur != null; cur = cur.GetPreviousFold())
            stack.Push(cur);

        Undo.RecordObject(renderer, "Generate Keyframe");
        var mat = renderer.material;

        for (int i = 0; i < 16; i++)
        {
            var prefix = $"_Vertex_Deformation_Slot_{i}_";
            var node = stack.Count > 0 ? stack.Pop() : null;
            var s = node?.Serialize();
            bool on = s != null;

            mat.SetFloat(prefix + "Enabled", on ? 1f : 0f);
            mat.SetInt(prefix + "Opcode", on ? s.opcode : 0);
            mat.SetFloat(prefix + "Float_0", on ? s.float0 : 0f);
            mat.SetFloat(prefix + "Float_1", on ? s.float1 : 0f);
            mat.SetFloat(prefix + "Float_2", on ? s.float2 : 0f);
            mat.SetFloat(prefix + "Float_3", on ? s.float3 : 0f);
            mat.SetVector(prefix + "Vector_0", on ? s.vec0 : Vector4.zero);
            mat.SetVector(prefix + "Vector_1", on ? s.vec1 : Vector4.zero);
            mat.SetVector(prefix + "Vector_2", on ? s.vec2 : Vector4.zero);
            mat.SetVector(prefix + "Vector_3", on ? s.vec3 : Vector4.zero);
        }
        EditorUtility.SetDirty(renderer);
    }
}

public class FoldGraphView : GraphView
{
    static readonly (string label, Type dataType)[] AllNodeOptions =
    {
        ("Fold/Axis Align", typeof(AxisAlignNodeData)),
        ("Fold/Plane to Tube", typeof(PlaneToTubeNodeData)),
        ("Fold/Point Align", typeof(PointAlignNodeData)),
        ("Fold/Tube to Plane", typeof(TubeToPlaneNodeData)),
        ("Value/Float", typeof(FloatValueNodeData)),
        ("Value/Vector", typeof(VectorValueNodeData)),
        ("Value/GameObject", typeof(GameObjectNodeData)),
        ("Output/Keyframe", typeof(KeyframeNodeData)),
    };

    readonly FoldWindow window;
    public readonly FoldGraph graphAsset;
    readonly Dictionary<string, FoldNodeView> nodeLookup = new();
    readonly FoldEdgeConnectorListener edgeConnectorListener;
    bool suppressGraphChanges;

    public FoldGraphView(FoldWindow owner, FoldGraph graph)
    {
        window = owner;
        graphAsset = graph;
        edgeConnectorListener = new FoldEdgeConnectorListener(this);

        style.flexGrow = 1;
        this.AddManipulator(new ContentZoomer());
        this.AddManipulator(new ContentDragger());
        this.AddManipulator(new SelectionDragger());
        this.AddManipulator(new RectangleSelector());

        var grid = new GridBackground();
        Insert(0, grid);
        grid.StretchToParentSize();

        graphViewChanged = OnGraphViewChanged;
        nodeCreationRequest = ctx => ShowSearchWindow(null, ctx.screenMousePosition);
        RegisterCallback<KeyDownEvent>(OnKeyDown);
        RegisterCallback<ContextualMenuPopulateEvent>(evt => { evt.menu.ClearItems(); evt.StopImmediatePropagation(); evt.PreventDefault(); ShowSearchWindow(null, evt.mousePosition); });
        Undo.undoRedoPerformed += OnUndoRedo;

        Reload();
    }

    public override List<Port> GetCompatiblePorts(Port startPort, NodeAdapter nodeAdapter) =>
        ports.Where(p => p.node != startPort.node && p.direction != startPort.direction &&
            (startPort.portType == null || p.portType == null || startPort.portType.IsAssignableFrom(p.portType) || p.portType.IsAssignableFrom(startPort.portType))).ToList();

    internal void BindEdgeConnector(Port port) => port.AddManipulator(new EdgeConnector<Edge>(edgeConnectorListener));

    GraphViewChange OnGraphViewChanged(GraphViewChange change)
    {
        if (suppressGraphChanges) return change;
        change.elementsToRemove?.ForEach(e => { if (e is Edge edge) RemoveEdgeData(edge); else if (e is FoldNodeView node) RemoveNode(node); });
        change.edgesToCreate?.ForEach(AddEdgeData);
        return change;
    }

    void RemoveNode(FoldNodeView node)
    {
        RecordUndo("Delete Node");
        graphAsset.nodes.Remove(node.Data);
        graphAsset.edges.RemoveAll(e => e.inputNodeGuid == node.Data.guid || e.outputNodeGuid == node.Data.guid);
        nodeLookup.Remove(node.Data.guid);
        MarkDirty();
    }

    internal void AddEdgeData(Edge edge)
    {
        if (edge.input?.node is not FoldNodeView inp || edge.output?.node is not FoldNodeView outp) return;
        RecordUndo("Add Edge");
        graphAsset.edges.Add(new FoldEdge { inputNodeGuid = inp.Data.guid, inputPortName = edge.input.portName, outputNodeGuid = outp.Data.guid, outputPortName = edge.output.portName });
        MarkDirty();
    }

    void RemoveEdgeData(Edge edge)
    {
        if (edge.input?.node is not FoldNodeView inp || edge.output?.node is not FoldNodeView outp) return;
        RecordUndo("Remove Edge");
        graphAsset.edges.RemoveAll(e => e.inputNodeGuid == inp.Data.guid && e.outputNodeGuid == outp.Data.guid && e.inputPortName == edge.input.portName && e.outputPortName == edge.output.portName);
        MarkDirty();
    }

    public void Reload()
    {
        suppressGraphChanges = true;
        nodeLookup.Clear();
        DeleteElements(graphElements.Where(e => e is Node or Edge).ToList());

        foreach (var d in graphAsset.nodes) { var n = CreateNode(d); nodeLookup[d.guid] = n; AddElement(n); }
        foreach (var e in graphAsset.edges)
        {
            if (nodeLookup.TryGetValue(e.outputNodeGuid, out var o) && nodeLookup.TryGetValue(e.inputNodeGuid, out var i))
                if (o.GetPort(e.outputPortName) is { } op && i.GetPort(e.inputPortName) is { } ip)
                    AddElement(op.ConnectTo(ip));
        }
        suppressGraphChanges = false;
    }

    void OnKeyDown(KeyDownEvent evt)
    {
        if (evt.keyCode == KeyCode.D && (evt.commandKey || evt.ctrlKey)) { evt.StopImmediatePropagation(); DuplicateSelectedNodes(); }
    }

    void DuplicateSelectedNodes()
    {
        var sel = selection.OfType<FoldNodeView>().ToList();
        if (sel.Count == 0) return;

        suppressGraphChanges = true;
        RecordUndo("Duplicate Nodes");

        var map = new Dictionary<string, FoldNodeView>();
        foreach (var orig in sel)
        {
            var clone = JsonUtility.FromJson<FoldNodeData>(JsonUtility.ToJson(orig.Data));
            clone = (FoldNodeData)Activator.CreateInstance(orig.Data.GetType());
            JsonUtility.FromJsonOverwrite(JsonUtility.ToJson(orig.Data), clone);
            clone.guid = Guid.NewGuid().ToString();
            clone.position += new Vector2(30f, 30f);
            graphAsset.nodes.Add(clone);
            var v = CreateNode(clone);
            map[orig.Data.guid] = v;
            nodeLookup[clone.guid] = v;
            AddElement(v);
        }

        foreach (var e in graphAsset.edges.Where(e => map.ContainsKey(e.outputNodeGuid) && map.ContainsKey(e.inputNodeGuid)))
            if (map[e.outputNodeGuid].GetPort(e.outputPortName) is { } op && map[e.inputNodeGuid].GetPort(e.inputPortName) is { } ip)
            { var edge = op.ConnectTo(ip); AddElement(edge); AddEdgeData(edge); }

        suppressGraphChanges = false;
        ClearSelection();
        foreach (var n in map.Values) AddToSelection(n);
        MarkDirty();
    }

    FoldNodeView CreateNode(FoldNodeData data)
    {
        FoldNodeView view = data switch
        {
            FloatValueNodeData or VectorValueNodeData or GameObjectNodeData => new ValueNodeView(),
            FoldOperationData => new OperationNodeView(),
            KeyframeNodeData => new KeyframeNodeView(),
            _ => throw new ArgumentOutOfRangeException(nameof(data))
        };
        view.Initialize(this, data);
        return view;
    }

    internal FoldNodeView CreateNode(Type dataType, Vector2 position)
    {
        if (!suppressGraphChanges) RecordUndo("Create Node");
        var data = (FoldNodeData)Activator.CreateInstance(dataType);
        data.position = position;
        graphAsset.nodes.Add(data);
        var view = CreateNode(data);
        nodeLookup[data.guid] = view;
        AddElement(view);
        view.RefreshExpandedState();
        view.RefreshPorts();
        MarkDirty();
        return view;
    }

    public void MarkDirty() => EditorUtility.SetDirty(graphAsset);
    public void ShowNotification(string msg) => window.ShowNotification(msg);
    public void ShowNodeCreationMenuFromPort(Port port, Vector2 pos) => ShowSearchWindow(port, pos);
    public void Dispose() => Undo.undoRedoPerformed -= OnUndoRedo;
    void OnUndoRedo() => Reload();
    internal void RecordUndo(string label) { if (!suppressGraphChanges) Undo.RegisterCompleteObjectUndo(graphAsset, label); }

    public FoldNodeView CreateNodeAndConnect(Port startPort, Type dataType, Vector2 pos)
    {
        suppressGraphChanges = true;
        var view = CreateNode(dataType, pos);
        var compat = GetCompatiblePorts(startPort, null).FirstOrDefault(p => p.node == view);
        if (compat != null)
        {
            var (o, i) = startPort.direction == Direction.Output ? (startPort, compat) : (compat, startPort);
            var edge = o.ConnectTo(i);
            AddElement(edge);
            AddEdgeData(edge);
        }
        suppressGraphChanges = false;
        MarkDirty();
        return view;
    }

    internal IEnumerable<(string label, Type dataType)> GetNodeOptions(Port port)
    {
        if (port == null) return AllNodeOptions;
        var t = port.portType;
        bool @out = port.direction == Direction.Output;
        return AllNodeOptions.Where(o =>
            (t == typeof(float) || t == typeof(Vector4)) ? (@out ? typeof(FoldOperationData).IsAssignableFrom(o.dataType) : (t == typeof(float) ? o.dataType == typeof(FloatValueNodeData) : o.dataType == typeof(VectorValueNodeData))) :
            t == typeof(GameObject) ? (@out ? o.dataType == typeof(KeyframeNodeData) : o.dataType == typeof(GameObjectNodeData)) :
            typeof(IFoldNode).IsAssignableFrom(t) && (typeof(FoldOperationData).IsAssignableFrom(o.dataType) || (@out && o.dataType == typeof(KeyframeNodeData))));
    }

    public T GetInputValue<T>(string guid, string port, T fallback) =>
        nodeLookup.TryGetValue(guid, out var n) && n.GetPort(port)?.connections.FirstOrDefault()?.output?.node is ValueNodeView v ? v.GetTypedValue<T>() : fallback;

    void ShowSearchWindow(Port port, Vector2 pos)
    {
        var p = ScriptableObject.CreateInstance<FoldSearchProvider>();
        p.Initialize(this, port, pos);
        if (p.HasEntries()) SearchWindow.Open(new SearchWindowContext(GUIUtility.GUIToScreenPoint(pos)), p);
    }
}

class FoldEdgeConnectorListener : IEdgeConnectorListener
{
    readonly FoldGraphView graphView;
    public FoldEdgeConnectorListener(FoldGraphView gv) => graphView = gv;

    public void OnDropOutsidePort(Edge edge, Vector2 pos) => graphView.ShowNodeCreationMenuFromPort(edge.output ?? edge.input, pos);

    public void OnDrop(GraphView gv, Edge edge)
    {
        if (edge.input == null || edge.output == null) return;
        edge.input.Connect(edge);
        edge.output.Connect(edge);
        gv.AddElement(edge);
        graphView?.AddEdgeData(edge);
    }
}

class FoldSearchProvider : ScriptableObject, ISearchWindowProvider
{
    FoldGraphView graphView;
    Port startPort;
    Vector2 worldPosition;
    List<(string label, Type dataType)> options;

    public void Initialize(FoldGraphView gv, Port port, Vector2 pos) { graphView = gv; startPort = port; worldPosition = pos; options = gv.GetNodeOptions(port).ToList(); }
    public bool HasEntries() => options?.Count > 0;

    public List<SearchTreeEntry> CreateSearchTree(SearchWindowContext ctx)
    {
        var entries = new List<SearchTreeEntry> { new SearchTreeGroupEntry(new GUIContent("Create Node"), 0) };
        foreach (var grp in options.GroupBy(o => o.label.Split('/')[0]).OrderBy(g => g.Key))
        {
            entries.Add(new SearchTreeGroupEntry(new GUIContent(grp.Key), 1));
            foreach (var opt in grp)
                entries.Add(new SearchTreeEntry(new GUIContent(opt.label.Split('/')[1])) { level = 2, userData = opt.dataType });
        }
        return entries;
    }

    public bool OnSelectEntry(SearchTreeEntry entry, SearchWindowContext ctx)
    {
        if (entry.userData is not Type dataType) return false;
        var pos = graphView.contentViewContainer.WorldToLocal(worldPosition);
        if (startPort != null) graphView.CreateNodeAndConnect(startPort, dataType, pos);
        else graphView.CreateNode(dataType, pos);
        return true;
    }
}