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

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

public interface IValueNode<T>
{
    T GetValue();
}

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

    public override void SetPosition(Rect newPos)
    {
        base.SetPosition(newPos);
        if (Data != null)
        {
            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;
        if (direction == Direction.Input) inputContainer.Add(port);
        else outputContainer.Add(port);
        return port;
    }

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

public abstract class FoldOperationNodeView : FoldNodeView, IFoldNode
{
    protected Port foldInputPort;
    protected Port foldOutputPort;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        foldInputPort = AddPort(Direction.Input, Port.Capacity.Single, "Input", typeof(IFoldNode));
        foldOutputPort = AddPort(Direction.Output, Port.Capacity.Single, "Out", typeof(IFoldNode));
    }

    public IFoldNode GetPreviousFold()
    {
        var edge = foldInputPort.connections.FirstOrDefault();
        return edge?.output?.node as IFoldNode;
    }

    public abstract FoldNodeSerialized Serialize();
}

public class FloatValueNodeView : FoldNodeView, IValueNode<float>
{
    FloatValueNodeData DataTyped => (FloatValueNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Output, Port.Capacity.Multi, "Out", typeof(float));

        var field = new FloatField("Value") { value = DataTyped.value };
        field.RegisterValueChangedCallback(evt =>
        {
            DataTyped.value = evt.newValue;
            graphView.MarkDirty();
        });
        mainContainer.Add(field);
    }

    public float GetValue() => DataTyped.value;
}

public class VectorValueNodeView : FoldNodeView, IValueNode<Vector4>
{
    VectorValueNodeData DataTyped => (VectorValueNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Output, Port.Capacity.Multi, "Out", typeof(Vector4));

        var field = new Vector4Field("Value") { value = DataTyped.value };
        field.RegisterValueChangedCallback(evt =>
        {
            DataTyped.value = evt.newValue;
            graphView.MarkDirty();
        });
        mainContainer.Add(field);
    }

    public Vector4 GetValue() => DataTyped.value;
}

public class GameObjectNodeView : FoldNodeView, IValueNode<GameObject>
{
    GameObjectNodeData DataTyped => (GameObjectNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Output, Port.Capacity.Multi, "Out", typeof(GameObject));

        var field = new ObjectField("Object")
        {
            allowSceneObjects = true,
            objectType = typeof(GameObject),
            value = DataTyped.output
        };
        field.RegisterValueChangedCallback(evt =>
        {
            DataTyped.output = evt.newValue as GameObject;
            graphView.MarkDirty();
        });
        mainContainer.Add(field);
    }

    public GameObject GetValue() => DataTyped.output;
}

public class AxisAlignNodeView : FoldOperationNodeView
{
    AxisAlignNodeData DataTyped => (AxisAlignNodeData)Data;

    Port poPort, ppPort, rPort, tPort;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);
        poPort = AddPort(Direction.Input, Port.Capacity.Single, "po", typeof(Vector4));
        ppPort = AddPort(Direction.Input, Port.Capacity.Single, "pp", typeof(Vector4));
        rPort = AddPort(Direction.Input, Port.Capacity.Single, "r", typeof(Vector4));
        tPort = AddPort(Direction.Input, Port.Capacity.Single, "t", typeof(float));

        mainContainer.Add(CreateVectorField("Origin", () => DataTyped.po, v => DataTyped.po = v));
        mainContainer.Add(CreateVectorField("Pivot", () => DataTyped.pp, v => DataTyped.pp = v));
        mainContainer.Add(CreateVectorField("Radial Axis", () => DataTyped.r, v => DataTyped.r = v));
        mainContainer.Add(CreateFloatField("Strength", () => DataTyped.t, v => DataTyped.t = v));
    }

    public override FoldNodeSerialized Serialize()
    {
        return DataTyped.Serialize(graphView);
    }

    Vector4Field CreateVectorField(string label, Func<Vector4> getter, Action<Vector4> setter)
    {
        var field = new Vector4Field(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }

    FloatField CreateFloatField(string label, Func<float> getter, Action<float> setter)
    {
        var field = new FloatField(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }
}

public class PlaneToTubeNodeView : FoldOperationNodeView
{
    PlaneToTubeNodeData DataTyped => (PlaneToTubeNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Input, Port.Capacity.Single, "p", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "r", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "s", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "t", typeof(float));

        mainContainer.Add(CreateVectorField("Origin", () => DataTyped.p, v => DataTyped.p = v));
        mainContainer.Add(CreateVectorField("Radial Axis", () => DataTyped.r, v => DataTyped.r = v));
        mainContainer.Add(CreateVectorField("Tangent Axis", () => DataTyped.s, v => DataTyped.s = v));
        mainContainer.Add(CreateFloatField("Strength", () => DataTyped.t, v => DataTyped.t = v));
    }

    public override FoldNodeSerialized Serialize() => DataTyped.Serialize(graphView);

    Vector4Field CreateVectorField(string label, Func<Vector4> getter, Action<Vector4> setter)
    {
        var field = new Vector4Field(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }

    FloatField CreateFloatField(string label, Func<float> getter, Action<float> setter)
    {
        var field = new FloatField(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }
}

public class PointAlignNodeView : FoldOperationNodeView
{
    PointAlignNodeData DataTyped => (PointAlignNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Input, Port.Capacity.Single, "po", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "pp", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "r", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "t", typeof(float));

        mainContainer.Add(CreateVectorField("Origin", () => DataTyped.po, v => DataTyped.po = v));
        mainContainer.Add(CreateVectorField("Pivot", () => DataTyped.pp, v => DataTyped.pp = v));
        mainContainer.Add(CreateVectorField("Radial Axis", () => DataTyped.r, v => DataTyped.r = v));
        mainContainer.Add(CreateFloatField("Strength", () => DataTyped.t, v => DataTyped.t = v));
    }

    public override FoldNodeSerialized Serialize() => DataTyped.Serialize(graphView);

    Vector4Field CreateVectorField(string label, Func<Vector4> getter, Action<Vector4> setter)
    {
        var field = new Vector4Field(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }

    FloatField CreateFloatField(string label, Func<float> getter, Action<float> setter)
    {
        var field = new FloatField(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }
}

public class TubeToPlaneNodeView : FoldOperationNodeView
{
    TubeToPlaneNodeData DataTyped => (TubeToPlaneNodeData)Data;

    public override void Initialize(FoldGraphView view, FoldNodeData data)
    {
        base.Initialize(view, data);

        AddPort(Direction.Input, Port.Capacity.Single, "p", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "r", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "s", typeof(Vector4));
        AddPort(Direction.Input, Port.Capacity.Single, "t", typeof(float));

        mainContainer.Add(CreateVectorField("Origin", () => DataTyped.p, v => DataTyped.p = v));
        mainContainer.Add(CreateVectorField("Radial Axis", () => DataTyped.r, v => DataTyped.r = v));
        mainContainer.Add(CreateVectorField("Tangent Axis", () => DataTyped.s, v => DataTyped.s = v));
        mainContainer.Add(CreateFloatField("Strength", () => DataTyped.t, v => DataTyped.t = v));
    }

    public override FoldNodeSerialized Serialize() => DataTyped.Serialize(graphView);

    Vector4Field CreateVectorField(string label, Func<Vector4> getter, Action<Vector4> setter)
    {
        var field = new Vector4Field(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }

    FloatField CreateFloatField(string label, Func<float> getter, Action<float> setter)
    {
        var field = new FloatField(label) { value = getter() };
        field.RegisterValueChangedCallback(evt =>
        {
            setter(evt.newValue);
            graphView.MarkDirty();
        });
        return field;
    }
}

public class KeyframeNodeView : FoldNodeView
{
    Port gameObjectPort;
    Port foldPort;
    readonly List<string> validationMessages = new();

    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));

        var button = new Button(OnGenerateClick) { text = "Generate Keyframe" };
        mainContainer.Add(button);
    }

    GameObject GetTargetObject()
    {
        var edge = gameObjectPort.connections.FirstOrDefault();
        if (edge?.output?.node is IValueNode<GameObject> goProvider)
            return goProvider.GetValue();
        return null;
    }

    IFoldNode GetFoldInput()
    {
        var edge = foldPort.connections.FirstOrDefault();
        return edge?.output?.node as IFoldNode;
    }

    void OnGenerateClick()
    {
        validationMessages.Clear();

        var go = GetTargetObject();
        if (go == null)
        {
            validationMessages.Add("GameObject input missing or empty.");
            ShowValidation();
            return;
        }

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

        var foldNode = GetFoldInput();
        if (foldNode == null)
        {
            validationMessages.Add("Fold input not connected.");
            ShowValidation();
            return;
        }

        var stack = new Stack<IFoldNode>();
        var current = foldNode;
        while (current != null)
        {
            stack.Push(current);
            current = current.GetPreviousFold();
        }

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

        const int kNumSlots = 16;
        for (int i = 0; i < kNumSlots; i++)
        {
            string slotPrefix = $"_Vertex_Deformation_Slot_{i}_";
            IFoldNode node = stack.Count > 0 ? stack.Pop() : null;
            FoldNodeSerialized data = node?.Serialize();

            bool active = data != null;
            mat.SetFloat(slotPrefix + "Enabled", active ? 1.0f : 0.0f);
            mat.SetInt(slotPrefix + "Opcode", active ? data.opcode : 0);

            mat.SetFloat(slotPrefix + "Float_0", active ? data.float0 : 0.0f);
            mat.SetFloat(slotPrefix + "Float_1", active ? data.float1 : 0.0f);
            mat.SetFloat(slotPrefix + "Float_2", active ? data.float2 : 0.0f);
            mat.SetFloat(slotPrefix + "Float_3", active ? data.float3 : 0.0f);

            mat.SetVector(slotPrefix + "Vector_0", active ? data.vec0 : Vector4.zero);
            mat.SetVector(slotPrefix + "Vector_1", active ? data.vec1 : Vector4.zero);
            mat.SetVector(slotPrefix + "Vector_2", active ? data.vec2 : Vector4.zero);
            mat.SetVector(slotPrefix + "Vector_3", active ? data.vec3 : Vector4.zero);
        }

        EditorUtility.SetDirty(renderer);
    }

    void ShowValidation()
    {
        graphView.ShowNotification(string.Join("\n", validationMessages));
    }
}

public class FoldGraphView : GraphView
{
    readonly FoldWindow window;
    public readonly FoldGraph graphAsset;
    readonly Dictionary<string, FoldNodeView> nodeLookup = new();
    bool suppressGraphChanges;

    public FoldGraphView(FoldWindow owner, FoldGraph graph)
    {
        window = owner;
        graphAsset = graph;

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

        gridBackground = new GridBackground();
        Insert(0, gridBackground);
        gridBackground.StretchToParentSize();

        graphViewChanged = OnGraphViewChanged;
        RegisterCallback<KeyDownEvent>(OnKeyDown);

        Reload();
    }

    GridBackground gridBackground;

    GraphViewChange OnGraphViewChanged(GraphViewChange change)
    {
        if (suppressGraphChanges)
            return change;

        if (change.elementsToRemove != null)
        {
            foreach (var element in change.elementsToRemove)
            {
                if (element is Edge e)
                {
                    RemoveEdgeData(e);
                }
                else if (element is FoldNodeView node)
                {
                    RemoveNode(node);
                }
            }
        }

        if (change.edgesToCreate != null)
        {
            foreach (var edge in change.edgesToCreate)
            {
                AddEdgeData(edge);
            }
        }

        return change;
    }

    void RemoveNode(FoldNodeView 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();
    }

    void AddEdgeData(Edge edge)
    {
        if (edge.input?.node is not FoldNodeView inputNode ||
            edge.output?.node is not FoldNodeView outputNode)
            return;

        graphAsset.edges.Add(new FoldEdge
        {
            inputNodeGuid = inputNode.Data.guid,
            inputPortName = edge.input.portName,
            outputNodeGuid = outputNode.Data.guid,
            outputPortName = edge.output.portName
        });
        MarkDirty();
    }

    void RemoveEdgeData(Edge edge)
    {
        if (edge.input?.node is not FoldNodeView inputNode ||
            edge.output?.node is not FoldNodeView outputNode)
            return;

        graphAsset.edges.RemoveAll(e =>
            e.inputNodeGuid == inputNode.Data.guid &&
            e.outputNodeGuid == outputNode.Data.guid &&
            e.inputPortName == edge.input.portName &&
            e.outputPortName == edge.output.portName);
        MarkDirty();
    }

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

        foreach (var nodeData in graphAsset.nodes)
        {
            var node = CreateNode(nodeData);
            nodeLookup[nodeData.guid] = node;
            AddElement(node);
        }

        foreach (var edgeData in graphAsset.edges)
        {
            if (!nodeLookup.TryGetValue(edgeData.outputNodeGuid, out var outputNode) ||
                !nodeLookup.TryGetValue(edgeData.inputNodeGuid, out var inputNode))
                continue;

            var outputPort = outputNode.GetPort(edgeData.outputPortName);
            var inputPort = inputNode.GetPort(edgeData.inputPortName);
            if (outputPort == null || inputPort == null)
                continue;

            var edge = new Edge
            {
                output = outputPort,
                input = inputPort
            };
            edge.input.Connect(edge);
            edge.output.Connect(edge);
            AddElement(edge);
        }

        suppressGraphChanges = false;
    }

    void OnKeyDown(KeyDownEvent evt)
    {
        bool isDuplicate = (evt.keyCode == KeyCode.D) && (evt.commandKey || evt.ctrlKey);
        if (!isDuplicate)
            return;

        evt.StopImmediatePropagation();
        DuplicateSelectedNodes();
    }

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

        suppressGraphChanges = true;

        var cloneMap = new Dictionary<string, FoldNodeView>();
        foreach (var original in selectedNodes)
        {
            var cloneData = CloneNodeData(original.Data);
            cloneData.position += new Vector2(30f, 30f);
            graphAsset.nodes.Add(cloneData);

            var cloneView = CreateNode(cloneData);
            cloneMap[original.Data.guid] = cloneView;
            nodeLookup[cloneData.guid] = cloneView;
            AddElement(cloneView);
        }

        var edgesToClone = graphAsset.edges
            .Where(e => cloneMap.ContainsKey(e.outputNodeGuid) && cloneMap.ContainsKey(e.inputNodeGuid))
            .ToList();

        foreach (var edgeData in edgesToClone)
        {
            var outputNode = cloneMap[edgeData.outputNodeGuid];
            var inputNode = cloneMap[edgeData.inputNodeGuid];
            var outputPort = outputNode.GetPort(edgeData.outputPortName);
            var inputPort = inputNode.GetPort(edgeData.inputPortName);
            if (outputPort == null || inputPort == null)
                continue;

            var edge = new Edge
            {
                output = outputPort,
                input = inputPort
            };
            edge.input.Connect(edge);
            edge.output.Connect(edge);
            AddElement(edge);
            AddEdgeData(edge);
        }

        suppressGraphChanges = false;

        MarkDirty();
    }

    FoldNodeData CloneNodeData(FoldNodeData data)
    {
        var json = JsonUtility.ToJson(data);
        var clone = (FoldNodeData)Activator.CreateInstance(data.GetType());
        JsonUtility.FromJsonOverwrite(json, clone);
        clone.guid = Guid.NewGuid().ToString();
        return clone;
    }

    FoldNodeView CreateNode(FoldNodeData data)
    {
        FoldNodeView view = data switch
        {
            FloatValueNodeData => new FloatValueNodeView(),
            VectorValueNodeData => new VectorValueNodeView(),
            GameObjectNodeData => new GameObjectNodeView(),
            AxisAlignNodeData => new AxisAlignNodeView(),
            PlaneToTubeNodeData => new PlaneToTubeNodeView(),
            PointAlignNodeData => new PointAlignNodeView(),
            TubeToPlaneNodeData => new TubeToPlaneNodeView(),
            KeyframeNodeData => new KeyframeNodeView(),
            _ => throw new ArgumentOutOfRangeException(nameof(data), $"Unknown node data type {data.GetType().Name}")
        };

        view.Initialize(this, data);
        return view;
    }

    public T CreateNode<T>(Vector2 position) where T : FoldNodeData, new()
    {
        var data = new T { position = position };
        graphAsset.nodes.Add(data);
        var view = CreateNode(data);
        nodeLookup[data.guid] = view;
        AddElement(view);
        MarkDirty();
        return data;
    }

    public override void BuildContextualMenu(ContextualMenuPopulateEvent evt)
    {
        base.BuildContextualMenu(evt);
        Vector2 graphPos = contentViewContainer.WorldToLocal(evt.mousePosition);

        evt.menu.AppendAction("Create/Fold/Axis Align", _ => CreateNode<AxisAlignNodeData>(graphPos));
        evt.menu.AppendAction("Create/Fold/Plane to Tube", _ => CreateNode<PlaneToTubeNodeData>(graphPos));
        evt.menu.AppendAction("Create/Fold/Point Align", _ => CreateNode<PointAlignNodeData>(graphPos));
        evt.menu.AppendAction("Create/Fold/Tube to Plane", _ => CreateNode<TubeToPlaneNodeData>(graphPos));
        evt.menu.AppendSeparator();
        evt.menu.AppendAction("Create/Value/Float", _ => CreateNode<FloatValueNodeData>(graphPos));
        evt.menu.AppendAction("Create/Value/Vector", _ => CreateNode<VectorValueNodeData>(graphPos));
        evt.menu.AppendAction("Create/Value/GameObject", _ => CreateNode<GameObjectNodeData>(graphPos));
        evt.menu.AppendSeparator();
        evt.menu.AppendAction("Create/Output/Keyframe", _ => CreateNode<KeyframeNodeData>(graphPos));
    }

    public void MarkDirty()
    {
        EditorUtility.SetDirty(graphAsset);
    }

    public T GetInputValue<T>(string nodeGuid, string portName, T fallback)
    {
        var node = nodeLookup.TryGetValue(nodeGuid, out var found) ? found : null;
        if (node == null) return fallback;

        var port = node.GetPort(portName);
        var edge = port?.connections.FirstOrDefault();
        if (edge?.output?.node is IValueNode<T> provider)
            return provider.GetValue();

        return fallback;
    }

    public void ShowNotification(string message)
    {
        window.ShowNotification(message);
    }
}