// DOM Elements
const sidebar = document.getElementById('sidebar');
const toggleSidebarBtn = document.getElementById('toggle-sidebar');
const settingsBtn = document.getElementById('settings-btn');
const settingsDropdown = document.getElementById('settings-dropdown');

// Modal Elements
const editModal = document.getElementById('edit-modal');
const modalTitle = document.getElementById('modal-title');
const nodeEditForm = document.getElementById('node-edit-form');
const linkEditForm = document.getElementById('link-edit-form');
const editNodeName = document.getElementById('edit-node-name');
const editLinkSource = document.getElementById('edit-link-source');
const editLinkTarget = document.getElementById('edit-link-target');
const editLinkValue = document.getElementById('edit-link-value');
const saveEditBtn = document.getElementById('save-edit');
const cancelEditBtn = document.getElementById('cancel-edit');
const closeModalBtn = document.getElementById('close-modal');

let currentEditItem = null; // Store the item being edited

// Modal Functions
function openModal() {
    editModal.classList.remove('hidden');
    feather.replace();
}

function closeModal() {
    editModal.classList.add('hidden');
    currentEditItem = null;
    nodeEditForm.classList.add('hidden');
    linkEditForm.classList.add('hidden');
}

function populateNodeForm(node) {
    modalTitle.textContent = 'Edit Node';
    nodeEditForm.classList.remove('hidden');
    linkEditForm.classList.add('hidden');
    editNodeName.value = node.name;
}

function populateLinkForm(link) {
    modalTitle.textContent = 'Edit Link';
    linkEditForm.classList.remove('hidden');
    nodeEditForm.classList.add('hidden');

    // Clear and populate source/target dropdowns
    editLinkSource.innerHTML = '';
    editLinkTarget.innerHTML = '';
    sankeyData.nodes.forEach(node => {
        const sourceOption = document.createElement('option');
        sourceOption.value = node.name;
        sourceOption.textContent = node.name;
        editLinkSource.appendChild(sourceOption.cloneNode(true));
        editLinkTarget.appendChild(sourceOption);
    });

    // Handle different link data structures
    let sourceName, targetName, linkValue;
    
    if (link.originalSource) {
        // Case: clicked from diagram
        sourceName = typeof link.originalSource === 'object' ? link.originalSource.name : sankeyData.nodes[link.originalSource].name;
        targetName = typeof link.originalTarget === 'object' ? link.originalTarget.name : sankeyData.nodes[link.originalTarget].name;
        linkValue = link.originalValue;
    } else {
        // Case: clicked from list
        sourceName = typeof link.source === 'object' ? link.source.name : sankeyData.nodes[link.source].name;
        targetName = typeof link.target === 'object' ? link.target.name : sankeyData.nodes[link.target].name;
        linkValue = link.value;
    }
    
    editLinkSource.value = sourceName;
    editLinkTarget.value = targetName;
    editLinkValue.value = linkValue;
}

// Event Handlers for Modal
closeModalBtn.addEventListener('click', closeModal);
cancelEditBtn.addEventListener('click', closeModal);

// Click outside modal to close
editModal.addEventListener('click', (e) => {
    if (e.target === editModal) {
        closeModal();
    }
});

// Toggle sidebar
toggleSidebarBtn.addEventListener('click', () => {
    sidebar.classList.toggle('collapsed');
});

// Toggle settings dropdown
let isSettingsOpen = false;
settingsBtn.addEventListener('click', (e) => {
    e.stopPropagation();
    isSettingsOpen = !isSettingsOpen;
    settingsDropdown.classList.toggle('hidden', !isSettingsOpen);
});

// Close settings dropdown when clicking outside
document.addEventListener('click', (e) => {
    if (!settingsDropdown.contains(e.target) && !settingsBtn.contains(e.target)) {
        settingsDropdown.classList.add('hidden');
        isSettingsOpen = false;
    }
});

// Prevent settings dropdown from closing when clicking inside it
settingsDropdown.addEventListener('click', (e) => {
    e.stopPropagation();
});

// Initialize Feather icons
document.addEventListener('DOMContentLoaded', () => {
    feather.replace();
});

// Define the d3.sankey layout function
d3.sankey = function() {
    var sankey = {},
        nodeWidth = 100,
        nodePadding = 80,
        size = [1, 1],
        nodes = [],
        links = [];

    sankey.nodeWidth = function(_) {
        if (!arguments.length) return nodeWidth;
        nodeWidth = +_;
        return sankey;
    };

    sankey.nodePadding = function(_) {
        if (!arguments.length) return nodePadding;
        nodePadding = +_;
        return sankey;
    };

    sankey.nodes = function(_) {
        if (!arguments.length) return nodes;
        nodes = _;
        return sankey;
    };

    sankey.links = function(_) {
        if (!arguments.length) return links;
        links = _;
        return sankey;
    };

    sankey.size = function(_) {
        if (!arguments.length) return size;
        size = _;
        return sankey;
    };

    sankey.layout = function(iterations) {
        if (!nodes.length) return sankey;
        computeNodeLinks();
        computeNodeValues();
        computeNodeBreadths();
        computeNodeDepths(iterations || 32);
        computeLinkDepths();
        return sankey;
    };

    sankey.relayout = function() {
        computeLinkDepths();
        return sankey;
    };

    sankey.link = function() {
        var curvature = .5;

        function link(d) {
            // Calculate node centers
            var sourceCenter = d.source.dx / 2,
                targetCenter = d.target.dx / 2;
                
            var x0 = d.source.x + sourceCenter,
                x1 = d.target.x + targetCenter,
                xi = d3.interpolateNumber(x0, x1),
                x2 = xi(curvature),
                x3 = xi(1 - curvature),
                y0 = d.source.y + d.sy + d.dy / 2,
                y1 = d.target.y + d.ty + d.dy / 2;
            return "M" + x0 + "," + y0
                 + "C" + x2 + "," + y0
                 + " " + x3 + "," + y1
                 + " " + x1 + "," + y1;
        }

        link.curvature = function(_) {
            if (!arguments.length) return curvature;
            curvature = +_;
            return link;
        };

        return link;
    };

    // Populate the sourceLinks and targetLinks for each node.
    function computeNodeLinks() {
        nodes.forEach(function(node) {
            node.sourceLinks = [];
            node.targetLinks = [];
        });
        links.forEach(function(link) {
            var source = link.source,
                target = link.target;
            if (typeof source === "number") source = link.source = nodes[link.source];
            if (typeof target === "number") target = link.target = nodes[link.target];
            source.sourceLinks.push(link);
            target.targetLinks.push(link);
        });
    }

    // Compute the value (size) of each node by summing the associated links.
    function computeNodeValues() {
        nodes.forEach(function(node) {
            node.value = Math.max(
                d3.sum(node.sourceLinks, value),
                d3.sum(node.targetLinks, value)
            );
        });
    }

    function computeNodeBreadths() {
        // Separate linked and unlinked nodes
        var linkedNodes = nodes.filter(function(node) {
            return node.sourceLinks.length > 0 || node.targetLinks.length > 0;
        });
        var unlinkedNodes = nodes.filter(function(node) {
            return node.sourceLinks.length === 0 && node.targetLinks.length === 0;
        });

        // Process only linked nodes for layout
        var remainingNodes = linkedNodes,
            nextNodes,
            x = 0;

        while (remainingNodes.length) {
            nextNodes = [];
            remainingNodes.forEach(function(node) {
                node.x = x;
                node.dx = nodeWidth;
                node.sourceLinks.forEach(function(link) {
                    if (nextNodes.indexOf(link.target) < 0) {
                        nextNodes.push(link.target);
                    }
                });
            });
            remainingNodes = nextNodes;
            ++x;
        }

        // Move sinks right
        moveSinksRight(x);

        // Scale the layout for linked nodes
        scaleNodeBreadths((size[0] - nodeWidth) / (x - 1));

        // Position unlinked nodes off-screen or at a specific position
        unlinkedNodes.forEach(function(node) {
            node.x = -1; // Position off-screen
            node.dx = 0;  // No width
        });
    }

    function moveSinksRight(x) {
        nodes.forEach(function(node) {
            if (!node.sourceLinks.length) {
                node.x = x - 1;
            }
        });
    }

    function scaleNodeBreadths(kx) {
        nodes.forEach(function(node) {
            node.x *= kx;
        });
    }

    function computeNodeDepths(iterations) {
        // Group nodes by x-coordinate
        var nodesByBreadth = d3.group(nodes, d => d.x);
        // Convert Map to Array and ensure it's sorted
        nodesByBreadth = Array.from(nodesByBreadth.values());

        initializeNodeDepth();
        resolveCollisions();
        for (var alpha = 1; iterations > 0; --iterations) {
            relaxRightToLeft(alpha *= .99);
            resolveCollisions();
            relaxLeftToRight(alpha);
            resolveCollisions();
        }

        function initializeNodeDepth() {
            // Calculate vertical scaling factor
            var ky = d3.min(nodesByBreadth, function(nodes) {
                return (size[1] - (nodes.length - 1) * nodePadding) / d3.sum(nodes, value);
            });

            nodesByBreadth.forEach(function(nodes) {
                nodes.forEach(function(node, i) {
                    node.y = i;
                    node.dy = node.value * ky;
                });
            });

            links.forEach(function(link) {
                link.dy = link.value * ky;
            });
        }

        function relaxLeftToRight(alpha) {
            nodesByBreadth.forEach(function(nodes) {
                nodes.forEach(function(node) {
                    if (node.targetLinks.length) {
                        var y = d3.sum(node.targetLinks, weightedSource) / d3.sum(node.targetLinks, value);
                        node.y += (y - center(node)) * alpha;
                    }
                });
            });

            function weightedSource(link) {
                return center(link.source) * link.value;
            }
        }

        function relaxRightToLeft(alpha) {
            nodesByBreadth.slice().reverse().forEach(function(nodes) {
                nodes.forEach(function(node) {
                    if (node.sourceLinks.length) {
                        var y = d3.sum(node.sourceLinks, weightedTarget) / d3.sum(node.sourceLinks, value);
                        node.y += (y - center(node)) * alpha;
                    }
                });
            });

            function weightedTarget(link) {
                return center(link.target) * link.value;
            }
        }

        function resolveCollisions() {
            nodesByBreadth.forEach(function(nodes) {
                var node,
                    dy,
                    y0 = 0,
                    n = nodes.length,
                    i;

                nodes.sort(ascendingDepth);
                for (i = 0; i < n; ++i) {
                    node = nodes[i];
                    dy = y0 - node.y;
                    if (dy > 0) node.y += dy;
                    y0 = node.y + node.dy + nodePadding;
                }

                dy = y0 - nodePadding - size[1];
                if (dy > 0) {
                    y0 = node.y -= dy;
                    for (i = n - 2; i >= 0; --i) {
                        node = nodes[i];
                        dy = node.y + node.dy + nodePadding - y0;
                        if (dy > 0) node.y -= dy;
                        y0 = node.y;
                    }
                }
            });
        }

        function ascendingDepth(a, b) {
            return a.y - b.y;
        }
    }

    function computeLinkDepths() {
        nodes.forEach(function(node) {
            node.sourceLinks.sort(ascendingTargetDepth);
            node.targetLinks.sort(ascendingSourceDepth);
        });
        nodes.forEach(function(node) {
            var sy = 0, ty = 0;
            node.sourceLinks.forEach(function(link) {
                link.sy = sy;
                sy += link.dy;
            });
            node.targetLinks.forEach(function(link) {
                link.ty = ty;
                ty += link.dy;
            });
        });

        function ascendingSourceDepth(a, b) {
            return a.source.y - b.source.y;
        }

        function ascendingTargetDepth(a, b) {
            return a.target.y - b.target.y;
        }
    }

    function value(link) {
        return link.value;
    }

    function center(node) {
        return node.y + node.dy / 2;
    }

    return sankey;
};

// Sankey diagram data structure with example data
let sankeyData = {
    nodes: [
        { name: "Source", color: "#4f46e5" },
        { name: "Process", color: "#4f46e5" },
        { name: "Output A", color: "#4f46e5" },
        { name: "Output B", color: "#4f46e5" }
    ],
    links: [
        { source: 0, target: 1, value: 100 },
        { source: 1, target: 2, value: 60 },
        { source: 1, target: 3, value: 40 }
    ]
};

// Load saved data from localStorage
const saved = localStorage.getItem('sankeyData');
if (saved) {
    sankeyData = JSON.parse(saved);
}

// Store node positions for dragging
let nodePositions = {};

// Utility function to convert hex to RGB
function hexToRgb(hex) {
    const result = /^#?([a-f\d]{2})([a-f\d]{2})([a-f\d]{2})$/i.exec(hex);
    return result ? {
        r: parseInt(result[1], 16),
        g: parseInt(result[2], 16),
        b: parseInt(result[3], 16)
    } : { r: 0, g: 0, b: 0 };
}

// Save data to localStorage
function saveData() {
    localStorage.setItem('sankeyData', JSON.stringify(sankeyData));
}

// Edit node or link
function handleEditNode(node, index) {
    currentEditItem = { type: 'node', index };
    populateNodeForm(node);
    openModal();
}

function handleEditLink(link, index) {
    currentEditItem = { type: 'link', index };
    populateLinkForm(link);
    openModal();
}

// Save changes from modal
saveEditBtn.addEventListener('click', () => {
    if (currentEditItem) {
        if (currentEditItem.type === 'node') {
            const newName = editNodeName.value.trim();
            if (newName) {
                sankeyData.nodes[currentEditItem.index].name = newName;
            }
        } else if (currentEditItem.type === 'link') {
            const sourceNode = editLinkSource.value;
            const targetNode = editLinkTarget.value;
            const value = parseInt(editLinkValue.value);
            
            // Validate the index exists
            if (currentEditItem.index >= 0 && currentEditItem.index < sankeyData.links.length) {
                // Find the indices of source and target nodes
                const sourceIndex = sankeyData.nodes.findIndex(n => n.name === sourceNode);
                const targetIndex = sankeyData.nodes.findIndex(n => n.name === targetNode);
                
                if (sourceIndex !== -1 && targetIndex !== -1 && value > 0) {
                    // Update the link data
                    sankeyData.links[currentEditItem.index] = {
                        source: sourceIndex,
                        target: targetIndex,
                        value: value
                    };
                }
            }
        }
        
        // Clear diagram cache to force complete redraw
        nodePositions = {};
        
        // Update everything
        updateLists();
        initDiagram();
        saveData();
        closeModal();
    }
});

// Initialize the diagram
function initDiagram() {
        const svg = d3.select("#sankey-diagram");
        svg.selectAll("*").remove();
        
        // Early return if SVG container doesn't exist
        if (!svg.node()) return;
        
        // Get the container dimensions
        const container = document.getElementById('diagram-container');
        const width = container.clientWidth;
        const containerHeight = container.clientHeight;
        
        // Set up the sankey diagram properties
        const baseHeight = parseInt(document.getElementById('diagram-height').value);
        const textPosition = document.getElementById('text-position').value;
        const textPadding = textPosition !== 'middle' ? 50 : 0; // Extra space for text above/below
        
        // Use the larger of container height or base height to prevent shrinking
        const height = Math.max(containerHeight, baseHeight + textPadding);
        const linkAlpha = parseFloat(document.getElementById('link-alpha').value);
        
        // Update SVG dimensions
        svg.attr("width", width)
           .attr("height", height)
           .attr("viewBox", [0, 0, width, height])
           .attr("preserveAspectRatio", "xMinYMin");
        
        // Create Sankey generator with dynamic padding based on node count
        const nodeCount = sankeyData.nodes.length;
        const dynamicPadding = Math.max(20, Math.min(120, Math.floor((height - textPadding) / (nodeCount * 1.5))));
        
        const sankey = d3.sankey()
            .nodeWidth(parseInt(document.getElementById('node-width').value))
            .nodePadding(dynamicPadding)
            .size([width - textPadding, height - textPadding]);    // Create a copy of the data
    const graph = {
        nodes: sankeyData.nodes.map(d => Object.assign({}, d)),
        links: sankeyData.links.map(d => Object.assign({}, d))
    };

    // Generate initial layout
    sankey.nodes(graph.nodes).links(graph.links).layout(32);

    // Filter only connected nodes for ky calculation
    const nodesForKy = graph.nodes.filter(node => 
        node.sourceLinks.length > 0 || node.targetLinks.length > 0
    );

    // Group connected nodes by x coordinate
    const nodesByX = d3.groups(nodesForKy, d => d.x);
    const ky = d3.min(nodesByX, ([x, nodes]) => 
        (height - (nodes.length - 1) * sankey.nodePadding()) / 
        d3.sum(nodes, d => d.value)
    );

    // Apply the scaling to nodes and links
    graph.nodes.forEach(node => {
        // Only apply height to connected nodes
        if (node.sourceLinks.length > 0 || node.targetLinks.length > 0) {
            const nodeHeight = node.value * ky;
            node.dy = nodeHeight;
        } else {
            node.dy = 0; // Unlinked nodes get zero height
        }
    });
    
    graph.links.forEach(link => {
        link.dy = link.value * ky;
    });

    // Relayout to apply the new heights
    sankey.layout(1);
    
    const nodes = sankey.nodes();
    const links = sankey.links();
    
    // Clear previous diagram
    svg.attr("width", width)
       .attr("height", height)
       .attr("viewBox", [0, 0, width, height])
       .attr("style", "max-width: 100%; height: auto;");

    // Add links
    const link = svg.append("g")
        .attr("fill", "none")
        .selectAll("path")
        .data(links)
        .join("path")
        .attr("d", sankey.link())
    link.attr("stroke", d => {
            const linkColorMode = document.getElementById('link-color-mode').value;
            let color;
            if (linkColorMode === 'static') {
                color = document.getElementById('static-link-color').value;
            } else {
                color = linkColorMode === 'source' 
                    ? d.source.color 
                    : d.target.color;
            }
            const finalColor = color || "#93c5fd";
            const rgb = hexToRgb(finalColor);
            return `rgba(${rgb.r},${rgb.g},${rgb.b},${linkAlpha})`;
        })
        .attr("stroke-width", d => Math.max(1, d.dy))
        .style("mix-blend-mode", "multiply")
        .attr("cursor", "pointer")
        .on("mouseover", function() {
            d3.select(this).attr("stroke-opacity", 0.8);
        })
        .on("mouseout", function() {
            d3.select(this).attr("stroke-opacity", linkAlpha);
        })
        .on("dblclick", function(event, d) {
            // Find the link index by matching source name and target name
            const index = sankeyData.links.findIndex(link => {
                const sourceName = typeof d.source === 'object' ? d.source.name : sankeyData.nodes[d.source].name;
                const targetName = typeof d.target === 'object' ? d.target.name : sankeyData.nodes[d.target].name;
                
                const linkSourceName = typeof link.source === 'object' ? 
                    link.source.name : sankeyData.nodes[link.source].name;
                const linkTargetName = typeof link.target === 'object' ? 
                    link.target.name : sankeyData.nodes[link.target].name;
                
                return sourceName === linkSourceName && 
                       targetName === linkTargetName && 
                       link.value === d.value;
            });
            
            // Store the link data for the modal
            const linkData = {
                source: sankeyData.links[index].source,
                target: sankeyData.links[index].target,
                value: sankeyData.links[index].value,
                sourceName: typeof d.source === 'object' ? d.source.name : sankeyData.nodes[d.source].name,
                targetName: typeof d.target === 'object' ? d.target.name : sankeyData.nodes[d.target].name
            };
            
            handleEditLink(linkData, index);
        });

    // Filter out nodes that have no connections
    const connectedNodes = nodes.filter(node => 
        (node.sourceLinks && node.sourceLinks.length > 0) || 
        (node.targetLinks && node.targetLinks.length > 0)
    );

    // Update diagram container background
    const backgroundColor = document.getElementById('background-color').value;
    document.getElementById('diagram-container').style.backgroundColor = backgroundColor;

    // Add nodes with better styling
    const node = svg.append("g")
        .selectAll("g")
        .data(connectedNodes)
        .join("g")
        .attr("transform", d => `translate(${d.x},${d.y})`);

    // Apply stored positions if they exist
    node.attr("transform", d => {
        const pos = nodePositions[d.index];
        return pos ? `translate(${pos.x},${pos.y})` : `translate(${d.x},${d.y})`;
    });

    const nodeEdgeColor = document.getElementById('node-edge-color').value;
    
    // Add node rectangles with click interaction
    node.append("rect")
        .attr("height", d => Math.max(d.dy, 1)) // Ensure minimum height of 1px
        .attr("width", d => d.dx)
        .attr("fill", d => d.color || "#4f46e5")
        .attr("stroke", nodeEdgeColor)
        .attr("stroke-width", 1)
        .attr("rx", 3) // Rounded corners
        .attr("ry", 3)
        .attr("cursor", "pointer")
        .on("dblclick", function(event, d) {
            handleEditNode(d, sankeyData.nodes.findIndex(node => node.name === d.name));
        });

    // Add drag behavior
    node.call(d3.drag()
        .subject(function(event, d) {
            // Get mouse position relative to the node
            const mousePoint = d3.pointer(event, this);
            return { 
                x: d.x,  // Keep x position fixed
                y: d.y + mousePoint[1]  // Add relative y position
            };
        })
        .on("start", function(event) {
            d3.select(this).raise();
        })
        .on("drag", function(event, d) {
            // Calculate new position based on mouse movement
            const yOffset = event.dy;  // Use the change in y position
            d.y += yOffset;  // Update position incrementally
            
            nodePositions[d.index] = { x: d.x, y: d.y };
            d3.select(this).attr("transform", `translate(${d.x},${d.y})`);

            // Update links
            link.attr("d", sankey.link());

            // After dragging, relayout to maintain proper link positions
            sankey.relayout();
        })
    );

    // Add text with customizable position
    let textX, textY, textAnchor;

    // Get the text padding value from the settings
    const textPaddingValue = parseInt(document.getElementById('text-padding').value);
    
    function getDefaultTextPosition(node) {
        const diagramWidth = width - textPadding;
        const xRatio = node.x / diagramWidth;
        
        if (xRatio < 0.33) {  // Left nodes
            return {
                x: node.dx + textPaddingValue,  // Position text on right side
                y: node.dy / 2,
                anchor: "start"
            };
        } else if (xRatio > 0.66) {  // Right nodes
            return {
                x: -textPaddingValue,  // Position text on left side
                y: node.dy / 2,
                anchor: "end"
            };
        } else {  // Middle nodes
            return {
                x: node.dx + textPaddingValue,  // Position text on right side
                y: node.dy / 2,
                anchor: "start"
            };
        }
    }
    
    switch(textPosition) {
        case 'default':
            // Text position will be determined per node
            textX = d => getDefaultTextPosition(d).x;
            textY = d => getDefaultTextPosition(d).y;
            textAnchor = d => getDefaultTextPosition(d).anchor;
            break;
        case 'left':
            textX = -textPaddingValue;
            textY = d => d.dy / 2;
            textAnchor = "end";
            break;
        case 'right':
            textX = d => d.dx + textPaddingValue;
            textY = d => d.dy / 2;
            textAnchor = "start";
            break;
        case 'top':
            textX = d => d.dx / 2;
            textY = -textPaddingValue;
            textAnchor = "middle";
            break;
        case 'bottom':
            textX = d => d.dx / 2;
            textY = d => d.dy + textPaddingValue;
            textAnchor = "middle";
            break;
    }

    // Get font size from settings
    const fontSizeMap = {
        'small': '10px',
        'medium': '12px',
        'large': '14px'
    };
    const fontSize = fontSizeMap[document.getElementById('font-size').value] || '12px';

    // Get text shadow settings
    const textShadowMap = {
        'none': 'none',
        'light': '1px 1px 2px rgba(0,0,0,0.2)',
        'medium': '2px 2px 4px rgba(0,0,0,0.3)',
        'dark': '3px 3px 6px rgba(0,0,0,0.4)'
    };
    const textShadow = textShadowMap[document.getElementById('text-shadow').value] || 'none';

    node.append("text")
        .attr("x", textX)
        .attr("y", textY || (d => d.dy / 2))
        .attr("dy", d => {
            // Adjust vertical alignment based on position
            if (textPosition === 'top') return '-0.5em';
            if (textPosition === 'bottom') return '1em';
            return '.35em';
        })
        .attr("text-anchor", typeof textAnchor === 'function' ? textAnchor : d => textAnchor)
        .text(d => d.name)
        .attr("fill", document.getElementById('text-color').value)
        .attr("font-size", fontSize)
        .attr("font-weight", "bold")
        .style("text-shadow", textShadow)
        .attr("pointer-events", "none");
}

// Update UI lists
function updateLists() {
    // Nodes list
    const nodesList = document.getElementById('nodes-list');
    nodesList.innerHTML = '';
    
    sankeyData.nodes.forEach((node, index) => {
        const nodeElement = document.createElement('div');
        nodeElement.className = 'flex items-center justify-between bg-white p-2 rounded border border-gray-200 cursor-pointer';
        nodeElement.innerHTML = `
            <span>${node.name}</span>
            <div class="flex items-center space-x-2">
                <input type="color" class="node-color-picker w-6 h-6" data-index="${index}" value="${node.color || '#4f46e5'}">
                <button class="text-red-500 hover:text-red-700 delete-node" data-index="${index}">
                    <i data-feather="trash-2" width="16"></i>
                </button>
            </div>
        `;
        // Add double-click handler for editing
        nodeElement.addEventListener('dblclick', (e) => {
            if (!e.target.classList.contains('node-color-picker') && !e.target.closest('.delete-node')) {
                handleEditNode(node, index);
            }
        });
        nodesList.appendChild(nodeElement);
    });

    // Links list
    const linksList = document.getElementById('links-list');
    linksList.innerHTML = '';
    
    sankeyData.links.forEach((link, index) => {
        const sourceNode = sankeyData.nodes[link.source]?.name || 'Unknown';
        const targetNode = sankeyData.nodes[link.target]?.name || 'Unknown';
        
        const linkElement = document.createElement('div');
        linkElement.className = 'flex items-center justify-between bg-white p-2 rounded border border-gray-200 cursor-pointer';
        linkElement.innerHTML = `
            <span>${sourceNode} → ${targetNode} (${link.value})</span>
            <button class="text-red-500 hover:text-red-700 delete-link" data-index="${index}">
                <i data-feather="trash-2" width="16"></i>
            </button>
        `;
        // Add double-click handler for editing
        linkElement.addEventListener('dblclick', (e) => {
            if (!e.target.closest('.delete-link')) {
                handleEditLink(link, index);
            }
        });
        linksList.appendChild(linkElement);
    });

    // Update dropdowns
    updateNodeDropdowns();
    
    // Refresh feather icons
    feather.replace();
}

// Update node dropdowns in links section
function updateNodeDropdowns() {
    const sourceSelect = document.getElementById('source-node');
    const targetSelect = document.getElementById('target-node');
    
    sourceSelect.innerHTML = '';
    targetSelect.innerHTML = '';
    
    sankeyData.nodes.forEach((node, index) => {
        const option = document.createElement('option');
        option.value = index;
        option.textContent = node.name;
        sourceSelect.appendChild(option.cloneNode(true));
        targetSelect.appendChild(option);
    });
}

// Export to CSV function
function exportToCSV() {
    try {
        if (!sankeyData || !sankeyData.links || sankeyData.links.length === 0) {
            alert("No hay datos para exportar.");
            return;
        }

        let csvContent = "Source,Target,Value,Color\r\n";
        sankeyData.links.forEach(link => {
            const src = (sankeyData.nodes[link.source] && sankeyData.nodes[link.source].name) || "";
            const tgt = (sankeyData.nodes[link.target] && sankeyData.nodes[link.target].name) || "";
            const val = link.value || 0;
            const color = (sankeyData.nodes[link.source] && sankeyData.nodes[link.source].color) || "";
            // Escape quotes
            csvContent += `"${src.replace(/"/g,'""')}","${tgt.replace(/"/g,'""')}",${val},"${color}"\r\n`;
        });

        const blob = new Blob([csvContent], { type: 'text/csv;charset=utf-8;' });
        const url = URL.createObjectURL(blob);

        const a = document.createElement('a');
        a.style.display = 'none';
        a.href = url;
        a.download = 'sankey_data.csv';
        document.body.appendChild(a);
        a.click();
        document.body.removeChild(a);
        URL.revokeObjectURL(url);

        console.log("CSV exported");
    } catch (err) {
        console.error("Error exporting CSV:", err);
        alert("Error exporting CSV. See console for details.");
    }
}


// Import from CSV function
function importFromCSV(file) {
    const reader = new FileReader();
    reader.onload = function(e) {
        const text = e.target.result;
        const lines = text.split('\n').filter(line => line.trim() !== '');
        
        if (lines.length < 2) {
            alert("Invalid CSV format");
            return;
        }
        
        // Reset data
        sankeyData = { nodes: [], links: [] };
        nodePositions = {};
        
        // Parse header
        const header = lines[0].split(',').map(h => h.trim().replace(/"/g, ''));
        const sourceIdx = header.indexOf('Source');
        const targetIdx = header.indexOf('Target');
        const valueIdx = header.indexOf('Value');
        const colorIdx = header.indexOf('Color');
        
        if (sourceIdx === -1 || targetIdx === -1 || valueIdx === -1) {
            alert("CSV must have columns: Source, Target, Value");
            return;
        }
        
        // Node name to index mapping
        const nodeMap = {};
        let nodeIndex = 0;
        
        // First pass: collect all unique nodes and their relationships
        const nodeRelations = new Map(); // Track incoming/outgoing links for each node
        
        // Parse data rows
        for (let i = 1; i < lines.length; i++) {
            const values = lines[i].split(',').map(v => v.trim().replace(/"/g, ''));
            if (values.length < 3) continue;
            
            const sourceName = values[sourceIdx];
            const targetName = values[targetIdx];
            const value = parseFloat(values[valueIdx]);
            const color = colorIdx !== -1 ? values[colorIdx] : "";
            
            if (isNaN(value)) continue;
            
            // Track relationships for column assignment
            if (!nodeRelations.has(sourceName)) {
                nodeRelations.set(sourceName, { in: 0, out: 0, totalValue: 0 });
            }
            if (!nodeRelations.has(targetName)) {
                nodeRelations.set(targetName, { in: 0, out: 0, totalValue: 0 });
            }
            
            nodeRelations.get(sourceName).out += 1;
            nodeRelations.get(sourceName).totalValue += value;
            nodeRelations.get(targetName).in += 1;
            nodeRelations.get(targetName).totalValue += value;
            
            // Add nodes if they don't exist
            if (!(sourceName in nodeMap)) {
                nodeMap[sourceName] = nodeIndex++;
                sankeyData.nodes.push({ 
                    name: sourceName,
                    color: color && /^#[0-9A-Fa-f]{6}$/.test(color.trim()) ? color.trim() : "#4f46e5",
                    value: 0 // Will be updated later
                });
            }
            
            if (!(targetName in nodeMap)) {
                nodeMap[targetName] = nodeIndex++;
                sankeyData.nodes.push({ 
                    name: targetName,
                    color: "#4f46e5",
                    value: 0 // Will be updated later
                });
            }
            
            // Add link
            sankeyData.links.push({
                source: nodeMap[sourceName],
                target: nodeMap[targetName],
                value: value
            });
        }
        
        // Update node values based on maximum flow
        sankeyData.nodes.forEach(node => {
            const relation = nodeRelations.get(node.name);
            node.value = Math.max(
                relation.totalValue,
                relation.in ? relation.totalValue / relation.in : 0,
                relation.out ? relation.totalValue / relation.out : 0
            );
        });
        
        updateLists();
        initDiagram();
        saveData();
    };
    reader.readAsText(file);
}


// Center diagram
function centerDiagram() {
    nodePositions = {};
    initDiagram();
}

// Event Listeners
document.addEventListener('DOMContentLoaded', () => {
    // Add node
    document.getElementById('add-node').addEventListener('click', () => {
        const nodeName = document.getElementById('node-name').value.trim();
        if (nodeName) {
            sankeyData.nodes.push({ name: nodeName, color: "#4f46e5" });
            document.getElementById('node-name').value = '';
            updateLists();
            initDiagram();
            saveData();
        }
    });
    
    // Add link
    document.getElementById('add-link').addEventListener('click', () => {
        const sourceIndex = parseInt(document.getElementById('source-node').value);
        const targetIndex = parseInt(document.getElementById('target-node').value);
        const value = parseInt(document.getElementById('link-value').value) || 1;
        
        if (!isNaN(sourceIndex) && !isNaN(targetIndex) && !isNaN(value) && value > 0) {
            sankeyData.links.push({
                source: sourceIndex,
                target: targetIndex,
                value: value
            });
            document.getElementById('link-value').value = '';
            updateLists();
            initDiagram();
            saveData();
        }
    });

    // Export to CSV
    document.getElementById('export-csv').addEventListener('click', exportToCSV);
    
    // Import from CSV
    document.getElementById('import-csv').addEventListener('click', () => {
        document.getElementById('csv-file').click();
    });
    
    document.getElementById('csv-file').addEventListener('change', (e) => {
        const file = e.target.files[0];
        if (file) {
            importFromCSV(file);
        }
        e.target.value = ''; // Reset input
    });

    // Settings panel toggle
    const settingsBtn = document.getElementById('settings-btn');
    const settingsPanel = document.getElementById('settings-panel');
    if (settingsBtn && settingsPanel) {
        settingsBtn.addEventListener('click', () => {
            settingsPanel.classList.toggle('hidden');
        });
    }
    
    // Center diagram
    document.getElementById('center-diagram').addEventListener('click', () => {
        centerDiagram();
        saveData();
    });
    
    // Delete node or link
    document.addEventListener('click', (e) => {
        if (e.target.closest('.delete-node')) {
            const button = e.target.closest('.delete-node');
            const index = parseInt(button.getAttribute('data-index'));
            
            // Delete node and associated links
            sankeyData.links = sankeyData.links.filter(link => 
                link.source !== index && link.target !== index
            );
            sankeyData.nodes.splice(index, 1);
            
            // Update link indices
            sankeyData.links = sankeyData.links.map(link => ({
                source: link.source > index ? link.source - 1 : link.source,
                target: link.target > index ? link.target - 1 : link.target,
                value: link.value
            }));
            
            // Remove node position data
            delete nodePositions[index];
            
            // Update remaining node positions
            const newNodePositions = {};
            Object.keys(nodePositions).forEach(key => {
                const keyNum = parseInt(key);
                if (keyNum > index) {
                    newNodePositions[keyNum - 1] = nodePositions[key];
                } else if (keyNum < index) {
                    newNodePositions[keyNum] = nodePositions[key];
                }
            });
            nodePositions = newNodePositions;
            
            updateLists();
            initDiagram();
            saveData();
        }
        
        if (e.target.closest('.delete-link')) {
            const button = e.target.closest('.delete-link');
            const index = parseInt(button.getAttribute('data-index'));
            
            sankeyData.links.splice(index, 1);
            
            updateLists();
            initDiagram();
            saveData();
        }
    });

    // Node color change
    document.addEventListener('input', (e) => {
        if (e.target.classList.contains('node-color-picker')) {
            const index = parseInt(e.target.getAttribute('data-index'));
            const color = e.target.value;
            sankeyData.nodes[index].color = color;
            initDiagram();
            saveData();
        }
    });

    // Color and options changes
    document.getElementById('text-color').addEventListener('input', initDiagram);
    document.getElementById('text-position').addEventListener('change', initDiagram);
    document.getElementById('text-padding').addEventListener('input', initDiagram);
    document.getElementById('node-width').addEventListener('input', initDiagram);
    document.getElementById('diagram-height').addEventListener('input', initDiagram);
    document.getElementById('link-alpha').addEventListener('input', initDiagram);
    document.getElementById('static-link-color').addEventListener('input', initDiagram);
    document.getElementById('font-size').addEventListener('change', initDiagram);
    document.getElementById('background-color').addEventListener('input', initDiagram);
    document.getElementById('node-edge-color').addEventListener('input', initDiagram);
    document.getElementById('text-shadow').addEventListener('change', initDiagram);
    
    // Handle link color mode changes
    const linkColorMode = document.getElementById('link-color-mode');
    const staticColorContainer = document.getElementById('static-link-color-container');
    
    linkColorMode.addEventListener('change', () => {
        staticColorContainer.classList.toggle('hidden', linkColorMode.value !== 'static');
        initDiagram();
    });

    // Initialize dropdowns and lists
    updateLists();
    initDiagram();
});