using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using UnityEditor;

namespace NanoBrain.Unity {
    public class Dag {

        public class Node {
            public int id;
            public Vector2 position;
            public float radius = 20f; // circle radius
            public Nucleus nucleus;
        }

        public class Edge {
            public int fromId;
            public int toId;
        }

        public List<Node> nodes = new();
        public List<Edge> edges = new();

        public Node FindNode(string name, bool justBaseName = true) {
            if (justBaseName) {
                int colonPos = name.IndexOf(":");
                if (colonPos > 0)
                    name = name[..colonPos];
            }
            foreach (Node node in this.nodes) {
                if (node.nucleus == null)
                    continue;
                    
                string nodeName = node.nucleus.name;
                if (justBaseName) {
                    int colonPos = nodeName.IndexOf(":");
                    if (colonPos > 0)
                        nodeName = nodeName[..colonPos];
                }
                if (nodeName == name)
                    return node;
            }
            return null;
        }

        public static Node GetNodeById(Dag dag, int id) => dag.nodes.FirstOrDefault(x => x.id == id);

        public static void ComputeLayout(Dag dag) {
            BuildAdjacencyAndPredecessors(dag, out Dictionary<int, List<int>> adjacency, out Dictionary<int, List<int>> predecessors);
            List<int> order = TopologicalOrder(adjacency, dag.nodes.Select(n => n.id));
            Dictionary<int, int> column = ComputeLongestPathColumns(adjacency, order, dag.nodes.Select(n => n.id));

            CreateDummiesAndEdges(dag, column, out List<Node> allNodes, out Dictionary<int, int> allColumns, out List<Edge> newEdges);

            List<List<int>> columns = GroupColumns(allColumns);
            PlaceNodes(allNodes, columns);

            // update dag with new lists (dummies included)
            dag.edges = newEdges;
            dag.nodes = allNodes;
        }

        // Helper: build adjacency and predecessor lists
        private static void BuildAdjacencyAndPredecessors(Dag dag, out Dictionary<int, List<int>> adjacency, out Dictionary<int, List<int>> predecessors) {
            adjacency = dag.nodes.ToDictionary(n => n.id, n => new List<int>());
            predecessors = dag.nodes.ToDictionary(n => n.id, n => new List<int>());
            foreach (Edge edge in dag.edges) {
                if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId))
                    continue;
                adjacency[edge.fromId].Add(edge.toId);
                predecessors[edge.toId].Add(edge.fromId);
            }
        }

        // Helper: compute topological order (returns empty list if cycle present)
        private static List<int> TopologicalOrder(Dictionary<int, List<int>> adjacency, IEnumerable<int> nodeIds) {
            Dictionary<int, int> inDegree = nodeIds.ToDictionary(id => id, _ => 0);
            foreach (KeyValuePair<int, List<int>> keyValue in adjacency)
                foreach (int to in keyValue.Value) if (inDegree.ContainsKey(to))
                    inDegree[to]++;

            Queue<int> queue = new(inDegree.Where(keyValue => keyValue.Value == 0).Select(kv => kv.Key));
            List<int> topo = new();
            while (queue.Count > 0) {
                int nodeId = queue.Dequeue();
                topo.Add(nodeId);
                foreach (int adjacentNodeId in adjacency[nodeId]) {
                    if (!inDegree.ContainsKey(adjacentNodeId))
                        continue;
                    inDegree[adjacentNodeId]--;
                    if (inDegree[adjacentNodeId] == 0)
                        queue.Enqueue(adjacentNodeId);
                }
            }
            return topo;
        }

        // Helper: longest-path-from-sinks column assignment (deterministic)
        private static Dictionary<int, int> ComputeLongestPathColumns(Dictionary<int, List<int>> adjacency, List<int> order, IEnumerable<int> nodeIds) {
            Dictionary<int, int> column = nodeIds.ToDictionary(id => id, _ => 0);
            foreach (int nodeId in Enumerable.Reverse(order)) {
                foreach (int child in adjacency[nodeId]) {
                    int cand = column[child] + 1;
                    if (cand > column[nodeId])
                        column[nodeId] = cand;
                }
            }
            return column;
        }

        // Helper: replace long edges with dummy node chains and return augmented node/column/edge lists
        private static void CreateDummiesAndEdges(Dag dag, Dictionary<int, int> column, out List<Node> allNodes, out Dictionary<int, int> allColumns, out List<Edge> newEdges) {
            allColumns = new Dictionary<int, int>(column);
            allNodes = new List<Node>(dag.nodes);
            newEdges = new List<Edge>();
            int nextDummyId = -1;

            foreach (Edge edge in dag.edges) {
                if (!column.ContainsKey(edge.fromId) || !column.ContainsKey(edge.toId)) {
                    newEdges.Add(edge);
                    continue;
                }

                int columnFrom = column[edge.fromId];
                int columnTo = column[edge.toId];
                int span = Mathf.Abs(columnTo - columnFrom);
                if (span <= 1) {
                    newEdges.Add(edge);
                    continue;
                }

                int prev = edge.fromId;
                int direction = columnTo > columnFrom ? 1 : -1;
                for (int step = 1; step < span; step++) {
                    int dummyCol = columnFrom + step * direction;
                    int dummyId = nextDummyId--;
                    Node dummy = new() {
                        id = dummyId,
                        position = Vector2.zero
                    };
                    // System.Reflection.FieldInfo field = typeof(Node).GetField("isDummy");
                    // if (field != null) field.SetValue(dummy, true);
                    allNodes.Add(dummy);
                    allColumns[dummyId] = dummyCol;

                    Edge newDummyEdge = new() {
                        fromId = prev,
                        toId = dummyId
                    };
                    newEdges.Add(newDummyEdge);
                    prev = dummyId;
                }
                Edge newEdge = new() {
                    fromId = prev,
                    toId = edge.toId
                };
                newEdges.Add(newEdge);
            }
        }

        // Helper: group columns into ordered list of lists
        private static List<List<int>> GroupColumns(Dictionary<int, int> allColumns) {
            return allColumns.GroupBy(kv => kv.Value).OrderBy(g => g.Key).Select(g => g.Select(x => x.Key).ToList()).ToList();
        }

        // Helper: place nodes vertically within each column
        private static void PlaceNodes(List<Node> allNodes, List<List<int>> columns) {
            float hSpacing = 100f;
            float totalHeight = 400f;
            for (int columnIx = 0; columnIx < columns.Count; columnIx++) {
                List<int> nodeList = columns[columnIx];
                float spacing = totalHeight / Mathf.Max(1, nodeList.Count);
                float margin = 10 + spacing / 2;
                for (int i = 0; i < nodeList.Count; i++) {
                    int id = nodeList[i];
                    Node node = allNodes.FirstOrDefault(n => n.id == id);
                    if (node == null)
                        continue;
                    float x = (hSpacing * 1.5f) + columnIx * hSpacing;
                    float y = margin + i * spacing;
                    node.position = new Vector2(x, y);
                }
            }
        }

        public static void ComputeLayoutKahn(Dag dag) {
            Dictionary<int, List<int>> adjacency = dag.nodes.ToDictionary(n => n.id, n => new List<int>());
            Dictionary<int, int> outdegree = dag.nodes.ToDictionary(node => node.id, n => 0);
            foreach (Edge edge in dag.edges) {
                if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId))
                    continue;
                adjacency[edge.fromId].Add(edge.toId);
                outdegree[edge.fromId]++;
            }

            // Kahn's algorithm to compute topological layers (horizontal layers)
            // build parent list (reverse adjacency) and parentIndegree = number of children each parent has
            Dictionary<int, List<int>> parents = dag.nodes.ToDictionary(n => n.id, _ => new List<int>());
            Dictionary<int, int> childCount = dag.nodes.ToDictionary(n => n.id, _ => 0);

            foreach (Edge edge in dag.edges) {
                if (!adjacency.ContainsKey(edge.fromId) || !adjacency.ContainsKey(edge.toId)) continue;
                adjacency[edge.fromId].Add(edge.toId);
                parents[edge.toId].Add(edge.fromId);   // parent of 'to' is 'from'
                childCount[edge.fromId]++;              // outdegree
            }

            Dictionary<int, int> column = new();
            Queue<int> queue = new(outdegree.Where(keyValue => keyValue.Value == 0).Select(keyValue => keyValue.Key));
            foreach (int id in queue)
                column[id] = 0;

            // process parents (reverse traversal)
            while (queue.Count > 0) {
                int nodeId = queue.Dequeue();
                int col = column[nodeId];
                foreach (int parentIx in parents[nodeId]) {
                    if (!column.ContainsKey(parentIx) || column[parentIx] < col + 1)
                        column[parentIx] = col + 1;
                    childCount[parentIx]--;               // decrement remaining unprocessed children
                    if (childCount[parentIx] == 0)
                        queue.Enqueue(parentIx);
                }
            }

            // Any unreachable nodes -> assign next layers
            int maxColumn = column.Count > 0 ? column.Values.Max() : 0;
            foreach (Node node in dag.nodes) {
                if (!column.ContainsKey(node.id)) {
                    maxColumn++;
                    column[node.id] = maxColumn;
                }
            }

            // Group nodes by column (left to right)
            List<List<int>> columns =
                column.
                GroupBy(kv => kv.Value).
                OrderBy(g => g.Key).
                Select(g => g.Select(x => x.Key).ToList()).
                ToList();

            float hSpacing = 100f;
            float totalHeight = 400f;

            // Place nodes: x increases with column index, y spaced within column
            for (int columnIx = 0; columnIx < columns.Count; columnIx++) {
                List<int> nodeList = columns[columnIx];
                float spacing = totalHeight / nodeList.Count;
                float margin = 10 + spacing / 2;
                for (int i = 0; i < nodeList.Count; i++) {
                    int index = nodeList[i];
                    Node node = GetNodeById(dag, index);
                    if (node == null)
                        continue;
                    float x = (hSpacing * 1.5f) + columnIx * hSpacing;
                    float y = margin + i * spacing;
                    node.position = new Vector2(x, y);
                }
            }

        }
    }
}