static/canvas.js

// 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

// 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;
    }
    layer.draw();
}

// Create nodes and connections from parsed data
function createNodesFromParsedData(parsedNodes, parsedConnections) {
    parsedNodes.forEach(nodeData => {
        const node = createNode(
            nodeData.x,
            nodeData.y,
            nodeData.label,
            nodeData.type,
            nodeData.inputs,
            nodeData.outputs,
            nodeData.id,
            nodeData.source,
            nodeData.parent_path,
            nodeData.level
        );
        nodes.push(node);
        layer.add(node);
    });
    layer.draw();
    autoConnect(); // Call autoConnect after nodes are created
    saveNodes();
}

// Create a node with inputs, outputs, and code segment
function createNode(x, y, label, type, inputs = [], outputs = [], id, source = '', parent_path = '', 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;
        }

        const nodePos = node.getAbsolutePosition();
        codeWindow = new Konva.Group({
            x: nodePos.x,
            y: nodePos.y + height + 10
        });

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

        const textarea = document.createElement('textarea');
        textarea.style.position = 'absolute';
        textarea.style.left = `${nodePos.x / scale + stage.x() / scale}px`;
        textarea.style.top = `${(nodePos.y + height + 10) / scale + stage.y() / scale}px`;
        textarea.style.width = '300px';
        textarea.style.height = '100px';
        textarea.style.fontFamily = 'monospace';
        textarea.value = source || '';
        document.body.appendChild(textarea);

        // Update code on change
        textarea.addEventListener('change', () => {
            const newSource = textarea.value;
            node.data.source = 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();
                document.body.removeChild(textarea);
                codeWindow = null;
                stage.off('click');
            }
        });

        codeWindow.add(codeBox);
        layer.add(codeWindow);
        layer.draw();
    });

    // Update position and connections on drag
    node.on('dragmove', () => {
        node.data.x = node.x();
        node.data.y = node.y();
        if (codeWindow) {
            const nodePos = node.getAbsolutePosition();
            codeWindow.position({ x: nodePos.x, y: nodePos.y + height + 10 });
            const textarea = document.querySelector(`textarea`);
            if (textarea) {
                textarea.style.left = `${nodePos.x / scale + stage.x() / scale}px`;
                textarea.style.top = `${(nodePos.y + height + 10) / scale + stage.y() / scale}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() {
    // Clear existing connections
    layer.find('Shape').forEach(shape => {
        if (shape.data && shape.data.fromNodeId !== undefined) {
            shape.destroy();
        }
    });
    connections = [];

    // Sort nodes by level and y-position to approximate program order
    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;
    });

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

    // Create connections based on parsed connections and hierarchy
    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
                });
            }
        }
    });

    // Additional connections based on hierarchy and role
    sortedNodes.forEach(node => {
        const nodeId = node.data.id;
        const parent = node.data.parent_path.split(' -> ')[0] || 'global';
        const role = node.data.type;

        // Connect to parent if applicable
        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
                    });
                }
            }
        }

        // Connect variables to their uses
        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;
    });

    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(
        Math.random() * (stage.width() - 100),
        Math.random() * (stage.height() - 100),
        '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 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();