// Initialize Konva stage
const stage = new Konva.Stage({
    container: 'container',
    width: 1000,
    height: 600,
    draggable: true
});

const layer = new Konva.Layer();
stage.add(layer);

// Store nodes, connections, and parsed data
let nodes = [];
let connections = [];
let parsedConnections = [];
let selectedPort = null;
let disconnectMode = false;
let codeWindow = null; // Track open code window
let codeTextarea = null; // Track textarea element

// Zoom functionality
let scale = 1;
const scaleFactor = 1.1;
const minScale = 0.5;
const maxScale = 2.0;

stage.on('wheel', (e) => {
    e.evt.preventDefault();
    const oldScale = scale;
    const pointer = stage.getPointerPosition();
    const delta = e.evt.deltaY > 0 ? 1 / scaleFactor : scaleFactor;
    const newScale = Math.max(minScale, Math.min(maxScale, oldScale * delta));
    if (newScale === scale) return;
    scale = newScale;
    const mousePointTo = {
        x: (pointer.x - stage.x()) / oldScale,
        y: (pointer.y - stage.y()) / oldScale
    };
    stage.scale({ x: scale, y: scale });
    const newPos = {
        x: pointer.x - mousePointTo.x * scale,
        y: pointer.y - mousePointTo.y * scale
    };
    stage.position(newPos);
    stage.batchDraw();
});

// Submit code or file for parsing
function submitCode() {
    const fileInput = document.getElementById('codeFile');
    const codeInput = document.getElementById('codeInput').value;
    const formData = new FormData();

    if (fileInput.files.length > 0) {
        formData.append('file', fileInput.files[0]);
    } else if (codeInput) {
        formData.append('code', codeInput);
    } else {
        alert('Please upload a file or paste code.');
        return;
    }

    fetch('/parse_code', {
        method: 'POST',
        body: formData
    })
    .then(response => response.json())
    .then(data => {
        if (data.error) {
            alert(data.error);
            return;
        }
        clearCanvas();
        parsedConnections = data.connections;
        createNodesFromParsedData(data.nodes, data.connections);
    })
    .catch(error => console.error('Error:', error));
}

// Clear existing nodes and connections
function clearCanvas() {
    nodes.forEach(node => node.destroy());
    layer.find('Shape').forEach(shape => shape.destroy());
    nodes = [];
    connections = [];
    parsedConnections = [];
    if (codeWindow) {
        codeWindow.destroy();
        codeWindow = null;
    }
    if (codeTextarea) {
        codeTextarea.remove();
        codeTextarea = null;
    }
    layer.draw();
}

// Create nodes and connections from parsed data
function createNodesFromParsedData(parsedNodes, parsedConnections) {
    const columns = {
        imports: { x: 50, y: 50, count: 0 }, // Column for imports
        global: { x: 250, y: 50, count: 0 }, // Column for global scope
        functions: {} // Columns for functions
    };

    // First pass: Assign function nodes to columns
    parsedNodes.forEach(nodeData => {
        if (nodeData.type === 'function') {
            const functionId = nodeData.id || `Function_${Object.keys(columns.functions).length + 1}`;
            columns.functions[functionId] = {
                x: 450 + Object.keys(columns.functions).length * 200,
                y: 50,
                count: 0
            };
        }
    });

    // Second pass: Create nodes with column-based positioning
    parsedNodes.forEach(nodeData => {
        const parentPath = nodeData.parent_path || 'global';
        const level = nodeData.level || 0;
        let x, y;

        if (nodeData.type === 'import') {
            // Imports column
            x = columns.imports.x;
            y = columns.imports.y + columns.imports.count * 80;
            columns.imports.count++;
        } else if (nodeData.type === 'function') {
            // Function column
            const functionId = nodeData.id;
            x = columns.functions[functionId].x;
            y = columns.functions[functionId].y + columns.functions[functionId].count * 80;
            columns.functions[functionId].count++;
        } else if (parentPath !== 'global' && parentPath.includes('Function')) {
            // Child of a function
            const functionId = parentPath.split(' -> ')[0];
            if (columns.functions[functionId]) {
                x = columns.functions[functionId].x;
                y = columns.functions[functionId].y + columns.functions[functionId].count * 80;
                columns.functions[functionId].count++;
            } else {
                // Fallback to global if function not found
                x = columns.global.x;
                y = columns.global.y + columns.global.count * 80;
                columns.global.count++;
            }
        } else {
            // Global scope (non-import, non-function)
            x = columns.global.x;
            y = columns.global.y + columns.global.count * 80;
            columns.global.count++;
        }

        const node = createNode(
            x,
            y,
            nodeData.label || 'Unnamed',
            nodeData.type || 'other',
            nodeData.inputs || [],
            nodeData.outputs || [],
            nodeData.id || `node_${nodes.length}`,
            nodeData.source || '',
            parentPath,
            level
        );
        nodes.push(node);
        layer.add(node);
    });

    layer.draw();
    autoConnect();
    saveNodes();
}

// Create a node with inputs, outputs, and code segment
function createNode(x, y, label, type, inputs = [], outputs = [], id, source = '', parent_path = 'global', level = 0) {
    const node = new Konva.Group({
        x: x,
        y: y,
        draggable: true
    });

    // Node rectangle
    const isNumberBox = type === 'number_box';
    const color = isNumberBox ? '#ffcccb' : type === 'function' ? '#ffeb3b' : type.includes('variable') ? '#90caf9' : type === 'import' ? '#a5d6a7' : '#ccc';
    const width = isNumberBox ? 80 : 100;
    const height = isNumberBox ? 40 : 50;

    const box = new Konva.Rect({
        width: width,
        height: height,
        fill: color,
        stroke: 'black',
        strokeWidth: 2,
        cornerRadius: 5
    });

    // Node label
    const text = new Konva.Text({
        text: label,
        fontSize: 12,
        fontFamily: 'Arial',
        fill: 'black',
        width: width,
        align: 'center',
        y: isNumberBox ? 14 : 20
    });

    node.add(box);
    node.add(text);

    // Input/output ports
    const inputPorts = inputs.map((input, i) => ({
        id: `input-${id}-${i}`,
        name: input,
        circle: new Konva.Circle({
            x: 0,
            y: 10 + i * 20,
            radius: 5,
            fill: 'red'
        })
    }));

    const outputPorts = outputs.map((output, i) => ({
        id: `output-${id}-${i}`,
        name: output,
        circle: new Konva.Circle({
            x: width,
            y: 10 + i * 20,
            radius: 5,
            fill: 'green'
        })
    }));

    // Add ports to node and set up click handlers
    inputPorts.forEach(port => {
        node.add(port.circle);
        port.circle.on('click', () => {
            if (!selectedPort) {
                selectedPort = { node, portId: port.id, type: 'input' };
                disconnectMode = true;
            } else if (selectedPort.type === 'output' && selectedPort.node !== node) {
                createSplineConnection(
                    selectedPort.node,
                    selectedPort.portId,
                    node,
                    port.id
                );
                connections.push({
                    fromNodeId: selectedPort.node.data.id,
                    fromPortId: selectedPort.portId,
                    toNodeId: node.data.id,
                    toPortId: port.id
                });
                selectedPort = null;
                disconnectMode = false;
                saveNodes();
            } else if (disconnectMode && selectedPort.type === 'input' && selectedPort.node === node) {
                selectedPort = { node, portId: port.id, type: 'input' };
            } else {
                selectedPort = null;
                disconnectMode = false;
            }
        });
    });

    outputPorts.forEach(port => {
        node.add(port.circle);
        port.circle.on('click', () => {
            if (!selectedPort) {
                selectedPort = { node, portId: port.id, type: 'output' };
                disconnectMode = false;
            } else if (disconnectMode && selectedPort.type === 'input') {
                const connIndex = connections.findIndex(
                    c => c.toNodeId === selectedPort.node.data.id &&
                         c.toPortId === selectedPort.portId &&
                         c.fromNodeId === node.data.id &&
                         c.fromPortId === port.id
                );
                if (connIndex !== -1) {
                    const conn = connections[connIndex];
                    const spline = layer.find('Shape').find(s =>
                        s.data.fromNodeId === conn.fromNodeId &&
                        s.data.fromPortId === conn.fromPortId &&
                        s.data.toNodeId === conn.toNodeId &&
                        s.data.toPortId === conn.toPortId
                    );
                    if (spline) spline.destroy();
                    connections.splice(connIndex, 1);
                    layer.draw();
                    saveNodes();
                }
                selectedPort = null;
                disconnectMode = false;
            } else {
                selectedPort = null;
                disconnectMode = false;
            }
        });
    });

    // Node data
    node.data = {
        id: id,
        type: type,
        label: label,
        inputs: inputPorts,
        outputs: outputPorts,
        x: x,
        y: y,
        source: source,
        parent_path: parent_path,
        level: level
    };

    // Click handler to show code window
    box.on('click', () => {
        if (codeWindow) {
            codeWindow.destroy();
            codeWindow = null;
        }
        if (codeTextarea) {
            codeTextarea.remove();
            codeTextarea = null;
        }

        const nodePos = node.getAbsolutePosition();
      
        // Create textarea for editing
        codeTextarea = document.createElement('textarea');
        codeTextarea.style.position = 'relative';
        const canvasRect = stage.container().getBoundingClientRect();
        const textareaX = nodePos.x;
      //  const textareaX = (nodePos.x + stage.x()) / scale + canvasRect.left;
       // const textareaY = (nodePos.y + height + 10 + stage.y()) / scale + canvasRect.top;
        const textareaY = nodePos.y;
        //codeTextarea.style.right = `${textareaX}px`;
        //codeTextarea.style.top = `${textareaY}px`;
        codeTextarea.style.left = '15px' || `${textareaX}px`;
        codeTextarea.style.top = '-200px' || `${textareaY}px`;
        codeTextarea.style.width = '300px' || `${300 / scale}px`;
        codeTextarea.style.height = '200px' || `${100 / scale}px`;
        codeTextarea.style.fontFamily = 'monospace';
        codeTextarea.style.fontSize = '12px' || `${12 / scale}px`;
        codeTextarea.style.background = '#ebebeb';
        codeTextarea.style.border = 'solid';
        codeTextarea.style.borderColor = 'grey';
        codeTextarea.style.resize = 'none';
        codeTextarea.value = source || '';
        document.body.appendChild(codeTextarea);
        codeWindow = new Konva.Group({
            x: node.x(),
            y: node.y() + height + 10
        });

        const codeBox = new Konva.Rect({
            width: 300,
            height: 100,
            fill: '#f0f0f0',
            stroke: 'black',
            strokeWidth: 1,
            cornerRadius: 5
        });

        // Display source in Konva Text for visual containment
        const codeText = new Konva.Text({
            x: 5,
            y: 5,
            text: source || '',
            fontSize: 12,
            fontFamily: 'monospace',
            fill: 'black',
            width: 290,
            padding: 5
        });

        //codeWindow.add(codeBox);
        //codeWindow.add(codeText);



        codeWindow.add(codeBox);
        codeWindow.add(codeText);
        //codeWindow.add(codeTextarea);
        layer.add(codeWindow);
        // Update code on change
        codeTextarea.addEventListener('change', () => {
            const newSource = codeTextarea.value;
            node.data.source = newSource;
            codeText.text(newSource);
            fetch('/update_node', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    id: node.data.id,
                    source: newSource
                })
            })
            .then(response => response.json())
            .then(data => {
                if (data.error) {
                    alert(data.error);
                } else {
                    console.log('Node updated:', data);
                    updateProgram();
                }
            })
            .catch(error => console.error('Error:', error));
        });

        // Close window on click outside
        stage.on('click', (e) => {
            if (e.target !== box && codeWindow) {
                codeWindow.destroy();
                codeWindow = null;
                if (codeTextarea) {
                    codeTextarea.remove();
                    codeTextarea = null;
                }
                stage.off('click');
            }
        });

        layer.draw();
    });

    // Update position and connections on drag
    node.on('dragmove', () => {
        node.data.x = node.x();
        node.data.y = node.y();
        if (codeWindow && codeTextarea) {
            const nodePos = node.getAbsolutePosition();
            codeWindow.position({ x: node.x(), y: node.y() + height + 10 });
            const canvasRect = stage.container().getBoundingClientRect();
            const textareaX = (nodePos.x + stage.x()) / scale + canvasRect.left;
            const textareaY = (nodePos.y + height + 10 + stage.y()) / scale + canvasRect.top;
            codeTextarea.style.left = `${textareaX}px`;
            codeTextarea.style.top = `${textareaY}px`;
        }
        updateConnections();
        saveNodes();
    });

    return node;
}

// Create a spline connection
function createSplineConnection(fromNode, fromPortId, toNode, toPortId) {
    const fromPort = fromNode.data.outputs.find(p => p.id === fromPortId);
    const toPort = toNode.data.inputs.find(p => p.id === toPortId);
    if (!fromPort || !toPort) return;

    const startX = fromNode.x() + fromPort.circle.x();
    const startY = fromNode.y() + fromPort.circle.y();
    const endX = toNode.x() + toPort.circle.x();
    const endY = toNode.y() + toPort.circle.y();

    const control1X = startX + (endX - startX) / 3;
    const control1Y = startY;
    const control2X = startX + 2 * (endX - startX) / 3;
    const control2Y = endY;

    const spline = new Konva.Shape({
        sceneFunc: function(context, shape) {
            context.beginPath();
            context.moveTo(startX, startY);
            context.bezierCurveTo(control1X, control1Y, control2X, control2Y, endX, endY);
            context.fillStrokeShape(shape);
        },
        stroke: 'black',
        strokeWidth: 2
    });

    spline.data = {
        fromNodeId: fromNode.data.id,
        fromPortId: fromPortId,
        toNodeId: toNode.data.id,
        toPortId: toPortId
    };
    layer.add(spline);
    layer.draw();
}

// Enhanced auto-connect based on hierarchy, position, and role
function autoConnect() {
    layer.find('Shape').forEach(shape => {
        if (shape.data && shape.data.fromNodeId !== undefined) {
            shape.destroy();
        }
    });
    connections = [];

    const sortedNodes = [...nodes].sort((a, b) => {
        if (a.data.level !== b.data.level) return a.data.level - b.data.level;
        return a.data.y - b.data.y; // Sort by y within columns
    });

    const hierarchy = {};
    sortedNodes.forEach(node => {
        const parent = node.data.parent_path.split(' -> ')[0] || 'global';
        if (!hierarchy[parent]) hierarchy[parent] = [];
        hierarchy[parent].push(node);
    });

    parsedConnections.forEach(conn => {
        const fromNode = nodes.find(n => n.data.id === conn.from);
        const toNode = nodes.find(n => n.data.id === conn.to);
        if (fromNode && toNode) {
            const fromPort = fromNode.data.outputs[0];
            const toPort = toNode.data.inputs[0];
            if (fromPort && toPort) {
                createSplineConnection(fromNode, fromPort.id, toNode, toPort.id);
                connections.push({
                    fromNodeId: fromNode.data.id,
                    fromPortId: fromPort.id,
                    toNodeId: toNode.data.id,
                    toPortId: toPort.id
                });
            }
        }
    });

    sortedNodes.forEach(node => {
        const nodeId = node.data.id;
        const parent = node.data.parent_path.split(' -> ')[0] || 'global';
        const role = node.data.type;

        if (parent !== 'global') {
            const parentNode = nodes.find(n => n.data.id === parent || n.data.label === parent.split('[')[0]);
            if (parentNode && parentNode.data.outputs.length > 0 && node.data.inputs.length > 0) {
                const fromPort = parentNode.data.outputs[0];
                const toPort = node.data.inputs[0];
                if (!connections.some(c => c.fromNodeId === parentNode.data.id && c.toNodeId === nodeId)) {
                    createSplineConnection(parentNode, fromPort.id, node, toPort.id);
                    connections.push({
                        fromNodeId: parentNode.data.id,
                        fromPortId: fromPort.id,
                        toNodeId: nodeId,
                        toPortId: toPort.id
                    });
                }
            }
        }

        if (role.includes('variable')) {
            const varName = node.data.label;
            sortedNodes.forEach(otherNode => {
                if (otherNode !== node && otherNode.data.source.includes(varName)) {
                    const fromPort = node.data.outputs[0];
                    const toPort = otherNode.data.inputs[0];
                    if (fromPort && toPort && !connections.some(c => c.fromNodeId === nodeId && c.toNodeId === otherNode.data.id)) {
                        createSplineConnection(node, fromPort.id, otherNode, toPort.id);
                        connections.push({
                            fromNodeId: nodeId,
                            fromPortId: fromPort.id,
                            toNodeId: otherNode.data.id,
                            toPortId: toPort.id
                        });
                    }
                }
            });
        }
    });

    layer.draw();
    saveNodes();
}

// Update full program
function updateProgram() {
    const program = reconstructProgram();
    document.getElementById('codeInput').value = program;
    fetch('/update_program', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ code: program })
    })
    .then(response => response.json())
    .then(data => {
        if (data.error) {
            alert(data.error);
        } else {
            console.log('Program updated:', data);
        }
    })
    .catch(error => console.error('Error:', error));
}

// Reconstruct the original program
function reconstructProgram() {
    const sortedNodes = [...nodes].sort((a, b) => {
        if (a.data.level !== b.data.level) return a.data.level - b.data.level;
        return a.data.y - b.data.y; // Sort by y within columns
    });

    let program = '';
    sortedNodes.forEach(node => {
        const source = node.data.source || '';
        const level = node.data.level || 0;
        const indent = '    '.repeat(level);
        program += indent + source.trim() + '\n';
    });

    return program.trim();
}

// Add a manual node
function addNode() {
    const node = createNode(
        250, // Add to global column
        50 + nodes.filter(n => n.data.parent_path === 'global').length * 80,
        'Function',
        'function',
        ['in1'],
        ['out1'],
        nodes.length,
        'def new_function(): pass',
        'global',
        0
    );
    nodes.push(node);
    layer.add(node);
    layer.draw();
    saveNodes();
}

// Update spline connections when nodes move
function createSplineConnection(fromNode, fromPortId, toNode, toPortId) {
    const fromPort = fromNode.data.outputs.find(p => p.id === fromPortId);
    const toPort = toNode.data.inputs.find(p => p.id === toPortId);
    if (!fromPort || !toPort) return;

    const startX = fromNode.x() + fromPort.circle.x();
    const startY = fromNode.y() + fromPort.circle.y();
    const endX = toNode.x() + toPort.circle.x();
    const endY = toNode.y() + toPort.circle.y();

    const control1X = startX + (endX - startX) / 3;
    const control1Y = startY;
    const control2X = startX + 2 * (endX - startX) / 3;
    const control2Y = endY;

    const spline = new Konva.Shape({
        sceneFunc: function(context, shape) {
            context.beginPath();
            context.moveTo(startX, startY);
            context.bezierCurveTo(control1X, control1Y, control2X, control2Y, endX, endY);
            context.fillStrokeShape(shape);
        },
        stroke: 'black',
        strokeWidth: 2
    });

    spline.data = {
        fromNodeId: fromNode.data.id,
        fromPortId: fromPortId,
        toNodeId: toNode.data.id,
        toPortId: toPortId
    };
    layer.add(spline);
    layer.draw();
}

// Update spline connections when nodes move
function updateConnections() {
    layer.find('Shape').forEach(shape => {
        if (shape.data && shape.data.fromNodeId !== undefined) {
            const fromNode = nodes.find(n => n.data.id === shape.data.fromNodeId);
            const toNode = nodes.find(n => n.data.id === shape.data.toNodeId);
            if (fromNode && toNode) {
                const fromPort = fromNode.data.outputs.find(p => p.id === shape.data.fromPortId);
                const toPort = toNode.data.inputs.find(p => p.id === shape.data.toPortId);
                if (fromPort && toPort) {
                    const startX = fromNode.x() + fromPort.circle.x();
                    const startY = fromNode.y() + fromPort.circle.y();
                    const endX = toNode.x() + toPort.circle.x();
                    const endY = toNode.y() + toPort.circle.y();

                    const control1X = startX + (endX - startX) / 3;
                    const control1Y = startY;
                    const control2X = startX + 2 * (endX - startX) / 3;
                    const control2Y = endY;

                    shape.sceneFunc(function(context, shape) {
                        context.beginPath();
                        context.moveTo(startX, startY);
                        context.bezierCurveTo(control1X, control1Y, control2X, control2Y, endX, endY);
                        context.fillStrokeShape(shape);
                    });
                }
            }
        }
    });
    layer.draw();
}

// Save nodes and connections to backend
function saveNodes() {
    fetch('/save_nodes', {
        method: 'POST',
        headers: {
            'Content-Type': 'application/json'
        },
        body: JSON.stringify({
            nodes: nodes.map(n => ({
                id: n.data.id,
                type: n.data.type,
                label: n.data.label,
                x: n.data.x,
                y: n.data.y,
                inputs: n.data.inputs.map(p => p.name),
                outputs: n.data.outputs.map(p => p.name),
                source: n.data.source,
                parent_path: n.data.parent_path,
                level: n.data.level
            })),
            connections: connections
        })
    })
    .then(response => response.json())
    .then(data => console.log('Saved:', data))
    .catch(error => console.error('Error:', error));
}

// Initial draw
layer.draw();