using UnityEngine;
using UnityEditor;
using System;
using System.IO;
using System.Text;
using System.Collections.Generic;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

public static class InjectToonFireToShaderGraph
{
    [MenuItem("Tools/Inject Toon Fire to Shader Graph")]
    public static void Inject()
    {
        string path = "Assets/Shaders/Fire_Shader.shadergraph";
        if (!File.Exists(path))
        {
            Debug.LogError($"ShaderGraph not found at {path}");
            return;
        }

        string jsonText = File.ReadAllText(path);
        List<JObject> jObjects = new List<JObject>();

        // Parse Multi-JSON using SupportMultipleContent = true
        using (var stringReader = new StringReader(jsonText))
        using (var jsonReader = new JsonTextReader(stringReader))
        {
            jsonReader.SupportMultipleContent = true;
            while (jsonReader.Read())
            {
                if (jsonReader.TokenType == JsonToken.StartObject)
                {
                    JObject obj = JObject.Load(jsonReader);
                    jObjects.Add(obj);
                }
            }
        }

        // Find the main GraphData object (usually the first one, but let's be safe)
        JObject graphData = null;
        foreach (var obj in jObjects)
        {
            if ((string)obj["m_Type"] == "UnityEditor.ShaderGraph.GraphData")
            {
                graphData = obj;
                break;
            }
        }

        if (graphData == null)
        {
            Debug.LogError("Could not find GraphData object in ShaderGraph file!");
            return;
        }

        // 1. Remove connections to BaseColor and Alpha block nodes in m_Edges
        // BaseColor block id: "5795a95f005b46bbbbf4c4e22a0ae9c3"
        // Alpha block id: "6c0a4ae63e3c48ad9e08ffbc4cf86bba"
        JArray edges = (JArray)graphData["m_Edges"];
        List<JToken> edgesToRemove = new List<JToken>();
        foreach (var edge in edges)
        {
            string inputNodeId = (string)edge["m_InputSlot"]["m_Node"]["m_Id"];
            if (inputNodeId == "5795a95f005b46bbbbf4c4e22a0ae9c3" || inputNodeId == "6c0a4ae63e3c48ad9e08ffbc4cf86bba")
            {
                edgesToRemove.Add(edge);
            }
        }

        foreach (var edge in edgesToRemove)
        {
            edges.Remove(edge);
        }

        // 2. Define our Custom Function node and Slot JObjects
        JObject slotInRawValue = JObject.Parse(@"{
            ""m_SGVersion"": 0,
            ""m_Type"": ""UnityEditor.ShaderGraph.Vector1MaterialSlot"",
            ""m_ObjectId"": ""slot_in_raw_value"",
            ""m_Id"": 0,
            ""m_DisplayName"": ""rawValue"",
            ""m_SlotType"": 0,
            ""m_Hidden"": false,
            ""m_ShaderOutputName"": ""rawValue"",
            ""m_StageCapability"": 3,
            ""m_Value"": 0.0,
            ""m_DefaultValue"": 0.0,
            ""m_Labels"": []
        }");

        JObject slotInBaseColor = JObject.Parse(@"{
            ""m_SGVersion"": 0,
            ""m_Type"": ""UnityEditor.ShaderGraph.Vector4MaterialSlot"",
            ""m_ObjectId"": ""slot_in_base_color"",
            ""m_Id"": 1,
            ""m_DisplayName"": ""baseColor"",
            ""m_SlotType"": 0,
            ""m_Hidden"": false,
            ""m_ShaderOutputName"": ""baseColor"",
            ""m_StageCapability"": 3,
            ""m_Value"": {
                ""x"": 1.0,
                ""y"": 1.0,
                ""z"": 1.0,
                ""w"": 1.0
            },
            ""m_DefaultValue"": {
                ""x"": 0.0,
                ""y"": 0.0,
                ""z"": 0.0,
                ""w"": 0.0
            },
            ""m_Labels"": []
        }");

        JObject slotOutColor = JObject.Parse(@"{
            ""m_SGVersion"": 0,
            ""m_Type"": ""UnityEditor.ShaderGraph.Vector3MaterialSlot"",
            ""m_ObjectId"": ""slot_out_color"",
            ""m_Id"": 2,
            ""m_DisplayName"": ""outColor"",
            ""m_SlotType"": 1,
            ""m_Hidden"": false,
            ""m_ShaderOutputName"": ""outColor"",
            ""m_StageCapability"": 3,
            ""m_Value"": {
                ""x"": 0.0,
                ""y"": 0.0,
                ""z"": 0.0
            },
            ""m_DefaultValue"": {
                ""x"": 0.0,
                ""y"": 0.0,
                ""z"": 0.0
            },
            ""m_Labels"": []
        }");

        JObject slotOutAlpha = JObject.Parse(@"{
            ""m_SGVersion"": 0,
            ""m_Type"": ""UnityEditor.ShaderGraph.Vector1MaterialSlot"",
            ""m_ObjectId"": ""slot_out_alpha"",
            ""m_Id"": 3,
            ""m_DisplayName"": ""outAlpha"",
            ""m_SlotType"": 1,
            ""m_Hidden"": false,
            ""m_ShaderOutputName"": ""outAlpha"",
            ""m_StageCapability"": 3,
            ""m_Value"": 0.0,
            ""m_DefaultValue"": 0.0,
            ""m_Labels"": []
        }");

        JObject customFunctionNode = JObject.Parse(@"{
            ""m_SGVersion"": 1,
            ""m_Type"": ""UnityEditor.ShaderGraph.CustomFunctionNode"",
            ""m_ObjectId"": ""node_cartoon_fire_custom"",
            ""m_Group"": {
                ""m_Id"": """"
            },
            ""m_Name"": ""ToonFireColor (Custom Function)"",
            ""m_DrawState"": {
                ""m_Expanded"": true,
                ""m_Position"": {
                    ""serializedVersion"": ""2"",
                    ""x"": -400.0,
                    ""y"": 56.0,
                    ""width"": 220.0,
                    ""height"": 180.0
                }
            },
            ""m_Slots"": [
                { ""m_Id"": ""slot_in_raw_value"" },
                { ""m_Id"": ""slot_in_base_color"" },
                { ""m_Id"": ""slot_out_color"" },
                { ""m_Id"": ""slot_out_alpha"" }
            ],
            ""synonyms"": [],
            ""m_Precision"": 0,
            ""m_PreviewExpanded"": false,
            ""m_DismissedVersion"": 0,
            ""m_PreviewMode"": 0,
            ""m_CustomColors"": {
                ""m_SerializableColors"": []
            },
            ""m_SourceType"": 1,
            ""m_FunctionName"": ""ToonFireColor"",
            ""m_FunctionSource"": """",
            ""m_FunctionBody"": ""float3 colInner = float3(1.0, 0.95, 0.4) * baseColor.rgb * 2.2;\nfloat3 colMid   = float3(1.0, 0.55, 0.0) * baseColor.rgb * 1.8;\nfloat3 colOuter = float3(1.0, 0.15, 0.1) * baseColor.rgb * 1.4;\n\nfloat stepOuter = step(0.15, rawValue);\nfloat stepMid = step(0.38, rawValue);\nfloat stepInner = step(0.58, rawValue);\n\nfloat3 finalColor = float3(0,0,0);\nif (stepInner > 0.5)\n{\n    finalColor = colInner;\n}\nelse if (stepMid > 0.5)\n{\n    finalColor = colMid;\n}\nelse if (stepOuter > 0.5)\n{\n    finalColor = colOuter;\n}\n\noutColor = finalColor;\noutAlpha = stepOuter * baseColor.a;""
        }");

        // Remove previous instances of our nodes from list (if re-running script)
        jObjects.RemoveAll(x => (string)x["m_ObjectId"] == "slot_in_raw_value");
        jObjects.RemoveAll(x => (string)x["m_ObjectId"] == "slot_in_base_color");
        jObjects.RemoveAll(x => (string)x["m_ObjectId"] == "slot_out_color");
        jObjects.RemoveAll(x => (string)x["m_ObjectId"] == "slot_out_alpha");
        jObjects.RemoveAll(x => (string)x["m_ObjectId"] == "node_cartoon_fire_custom");

        // Add them to the list of objects in the file
        jObjects.Add(slotInRawValue);
        jObjects.Add(slotInBaseColor);
        jObjects.Add(slotOutColor);
        jObjects.Add(slotOutAlpha);
        jObjects.Add(customFunctionNode);

        // Add the node reference ID to "m_Nodes" array in GraphData
        JArray nodes = (JArray)graphData["m_Nodes"];
        bool hasNode = false;
        foreach (var node in nodes)
        {
            if ((string)node["m_Id"] == "node_cartoon_fire_custom")
            {
                hasNode = true;
                break;
            }
        }
        if (!hasNode)
        {
            JObject nodeRef = new JObject();
            nodeRef["m_Id"] = "node_cartoon_fire_custom";
            nodes.Add(nodeRef);
        }

        // 3. Connect the new edges in GraphData
        // Edge 1: raw fire shape (MultiplyNode 0d74e8c23ef1428eaf8c38a4c033c563 output slot 2) -> slot_in_raw_value (CustomFunctionNode input slot 0)
        JObject edge1 = JObject.Parse(@"{
            ""m_OutputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""0d74e8c23ef1428eaf8c38a4c033c563""
                },
                ""m_SlotId"": 2
            },
            ""m_InputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""node_cartoon_fire_custom""
                },
                ""m_SlotId"": 0
            }
        }");

        // Edge 2: base color representation (MultiplyNode e48170ea86fb4291b3e4198160139b3e output slot 2) -> slot_in_base_color (CustomFunctionNode input slot 1)
        JObject edge2 = JObject.Parse(@"{
            ""m_OutputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""e48170ea86fb4291b3e4198160139b3e""
                },
                ""m_SlotId"": 2
            },
            ""m_InputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""node_cartoon_fire_custom""
                },
                ""m_SlotId"": 1
            }
        }");

        // Edge 3: CustomFunctionNode outColor (output slot 2) -> BaseColor block node 5795a95f005b46bbbbf4c4e22a0ae9c3 (slot 0)
        JObject edge3 = JObject.Parse(@"{
            ""m_OutputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""node_cartoon_fire_custom""
                },
                ""m_SlotId"": 2
            },
            ""m_InputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""5795a95f005b46bbbbf4c4e22a0ae9c3""
                },
                ""m_SlotId"": 0
            }
        }");

        // Edge 4: CustomFunctionNode outAlpha (output slot 3) -> Alpha block node 6c0a4ae63e3c48ad9e08ffbc4cf86bba (slot 0)
        JObject edge4 = JObject.Parse(@"{
            ""m_OutputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""node_cartoon_fire_custom""
                },
                ""m_SlotId"": 3
            },
            ""m_InputSlot"": {
                ""m_Node"": {
                    ""m_Id"": ""6c0a4ae63e3c48ad9e08ffbc4cf86bba""
                },
                ""m_SlotId"": 0
            }
        }");

        edges.Add(edge1);
        edges.Add(edge2);
        edges.Add(edge3);
        edges.Add(edge4);

        // Write all JObjects back to file separated by two newlines
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < jObjects.Count; i++)
        {
            sb.Append(jObjects[i].ToString(Newtonsoft.Json.Formatting.Indented));
            if (i < jObjects.Count - 1)
            {
                sb.Append("\n\n");
            }
        }

        File.WriteAllText(path, sb.ToString());

        // Restore material shader to original Fire_Shader.shadergraph
        string materialPath = "Assets/Materials/Fire_Material.mat";
        Material mat = AssetDatabase.LoadAssetAtPath<Material>(materialPath);
        if (mat != null)
        {
            Shader origShader = Shader.Find("Shader Graphs/Fire_Shader");
            if (origShader != null)
            {
                Undo.RecordObject(mat, "Restore original Fire ShaderGraph");
                mat.shader = origShader;
                EditorUtility.SetDirty(mat);
            }
        }

        AssetDatabase.ImportAsset(path, ImportAssetOptions.ForceUpdate);
        AssetDatabase.SaveAssets();
        Debug.Log("Successfully injected Toon Fire Custom Function into Fire_Shader.shadergraph, updated edges, and restored material shader!");
    }
}