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

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

let nodes = [];
let connections = [];
let parsedConnections = [];
let selectedPort = null;
let disconnectMode = false;

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));
}

function clearCanvas() {
    nodes.forEach(node => node.destroy());
    layer.find('Shape').forEach(shape => shape.destroy());
    nodes = [];
    connections = [];
    parsedConnections = [];
    layer.draw();
}

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.value
        );
        nodes.push(node);
        layer.add(node);
    });
    layer.draw();
    saveNodes();
}

function createNode(x, y, label, type, inputs = [], outputs = [], id, value = null) {
    const node = new Konva.Group({
        x: x,
        y: y,
        draggable: true
    });

    // Node appearance based on type
    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 and value
    const textContent = isNumberBox && value ? `${label} = ${value}` : label;
    const text = new Konva.Text({
        text: textContent,
        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'
        })
    }));

    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 = {
        id: id,
        type: type,
        label: label,
        inputs: inputPorts,
        outputs: outputPorts,
        x: x,
        y: y,
        value: value
    };

    node.on('dragmove', () => {
        node.data.x = node.x();
        node.data.y = node.y();
        updateConnections();
        saveNodes();
    });

    return node;
}

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();
}

function autoConnect() {
    layer.find('Shape').forEach(shape => {
        if (shape.data && shape.data.fromNodeId !== undefined) {
            shape.destroy();
        }
    });
    connections = [];

    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
                });
            }
        }
    });

    layer.draw();
    saveNodes();
}

function addNode() {
    const node = createNode(
        Math.random() * (stage.width() - 100),
        Math.random() * (stage.height() - 100),
        'Function',
        'function',
        ['in1'],
        ['out1'],
        nodes.length
    );
    nodes.push(node);
    layer.add(node);
    layer.draw();
    saveNodes();
}

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();
}

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),
                value: n.data.value
            })),
            connections: connections
        })
    }).then(response => response.json())
      .then(data => console.log('Saved:', data))
      .catch(error => console.error('Error:', error));
}

layer.draw();