// Initialize map - will be set to saved position or F. Dharanboodhoo, Maldives
const saved = localStorage.getItem('mapMakerData');
let initialView = [3.0667, 73.1833];
let initialZoom = 15;
let initialIsSatellite = false;

if (saved) {
    const data = JSON.parse(saved);
    if (data.mapView) {
        initialView = [data.mapView.center.lat, data.mapView.center.lng];
        initialZoom = data.mapView.zoom;
        initialIsSatellite = data.mapView.isSatellite || false;
    }
}

const map = L.map('map', {
    maxZoom: 25,
    dragging: false  // Disable default left-click dragging
}).setView(initialView, initialZoom);

// Create custom panes for proper z-index ordering
map.createPane('linesPane');
map.createPane('segmentsPane');
map.createPane('switchesPane');
map.createPane('nodesPane');

// Set z-index for each pane (higher = on top)
map.getPane('linesPane').style.zIndex = 400;
map.getPane('segmentsPane').style.zIndex = 410; // Above lines, below switches
map.getPane('switchesPane').style.zIndex = 450;
map.getPane('nodesPane').style.zIndex = 500;

// Enable middle mouse button and right-click dragging
let isPanning = false;
let panStart = null;

map.getContainer().addEventListener('mousedown', (e) => {
    // Middle mouse button OR right-click in select/move mode OR left-click in move mode
    if (e.button === 1 ||
        (e.button === 2 && (currentTool === 'select' || currentTool === 'move')) ||
        (e.button === 0 && currentTool === 'move')) {
        e.preventDefault();
        isPanning = true;
        panStart = { x: e.clientX, y: e.clientY };
        map.getContainer().style.cursor = 'grabbing';
    }
});

map.getContainer().addEventListener('mousemove', (e) => {
    if (isPanning && panStart) {
        const dx = e.clientX - panStart.x;
        const dy = e.clientY - panStart.y;

        const center = map.getCenter();
        const zoom = map.getZoom();
        const scale = 256 * Math.pow(2, zoom);

        const newCenter = map.unproject(
            map.project(center, zoom).subtract([dx, dy]),
            zoom
        );

        map.panTo(newCenter, { animate: false });
        panStart = { x: e.clientX, y: e.clientY };
    }
});

map.getContainer().addEventListener('mouseup', (e) => {
    if (e.button === 1 || e.button === 2 || e.button === 0) {
        isPanning = false;
        panStart = null;
        map.getContainer().style.cursor = '';
    }
});

map.getContainer().addEventListener('mouseleave', () => {
    if (isPanning) {
        isPanning = false;
        panStart = null;
        map.getContainer().style.cursor = '';
    }
});

// Base layers
const streetLayer = L.tileLayer('https://{s}.tile.openstreetmap.org/{z}/{x}/{y}.png', {
    attribution: '&copy; <a href="https://www.openstreetmap.org/copyright">OpenStreetMap</a> contributors',
    maxZoom: 25,
    maxNativeZoom: 19
});

const satelliteLayer = L.tileLayer('https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}', {
    attribution: '&copy; Google',
    maxZoom: 25,
    maxNativeZoom: 22
});

// Start with saved view or street view
let isSatelliteView = initialIsSatellite;
if (isSatelliteView) {
    satelliteLayer.addTo(map);
} else {
    streetLayer.addTo(map);
}

// Update button text when DOM is ready
window.addEventListener('DOMContentLoaded', () => {
    if (isSatelliteView) {
        document.getElementById('satelliteToggle').textContent = 'Street View';
    }
});

// State management
let currentTool = 'select';
let currentColor = '#FF0000';
let drawingLine = false;
let linePoints = [];
let rulerPoints = [];
let selectedObject = null;
let tempMarkers = [];
let selectedSegments = new Set(); // For sum tool
let sumDisplay = null;

// Storage for all drawn objects
let drawnObjects = {
    lines: [],
    nodes: [],
    switches: []
};

// Layer groups
const layerGroup = L.layerGroup().addTo(map);

// Tool button handlers
document.querySelectorAll('.tool-btn').forEach(btn => {
    btn.addEventListener('click', (e) => {
        document.querySelectorAll('.tool-btn').forEach(b => b.classList.remove('active'));
        e.target.classList.add('active');

        const toolId = e.target.id;
        if (toolId === 'lineTool') currentTool = 'line';
        else if (toolId === 'nodeTool') currentTool = 'node';
        else if (toolId === 'switchTool') currentTool = 'switch';
        else if (toolId === 'rulerTool') currentTool = 'ruler';
        else if (toolId === 'sumTool') currentTool = 'sum';
        else if (toolId === 'selectTool') currentTool = 'select';
        else if (toolId === 'moveTool') currentTool = 'move';
        else if (toolId === 'deleteTool') currentTool = 'delete';

        // Reset temporary states
        finishLineDrawing();
        rulerPoints = [];
        document.getElementById('rulerDisplay').style.display = 'none';

        // Clear sum selections when switching tools
        if (toolId !== 'sumTool') {
            clearSumSelections();
        } else {
            showSumDisplay();
        }
    });
});

// Satellite toggle handler
document.getElementById('satelliteToggle').addEventListener('click', () => {
    if (isSatelliteView) {
        map.removeLayer(satelliteLayer);
        map.addLayer(streetLayer);
        document.getElementById('satelliteToggle').textContent = 'Satellite View';
        isSatelliteView = false;
    } else {
        map.removeLayer(streetLayer);
        map.addLayer(satelliteLayer);
        document.getElementById('satelliteToggle').textContent = 'Street View';
        isSatelliteView = true;
    }
    saveToLocalStorage();
});

// Color picker handlers
document.querySelectorAll('.color-btn').forEach(btn => {
    btn.addEventListener('click', (e) => {
        document.querySelectorAll('.color-btn').forEach(b => b.classList.remove('active'));
        e.target.classList.add('active');
        currentColor = e.target.dataset.color;

        // If a node is selected, change its color
        if (selectedObject && selectedObject.type === 'node') {
            changeNodeColor(selectedObject.data, currentColor);
        }
    });
});

// Prevent popup clicks from triggering map clicks
map.on('popupopen', (e) => {
    const popup = e.popup;
    const container = popup.getElement();

    if (container) {
        container.addEventListener('click', (e) => {
            e.stopPropagation();
        });
        container.addEventListener('mousedown', (e) => {
            e.stopPropagation();
        });
        container.addEventListener('mouseup', (e) => {
            e.stopPropagation();
        });
    }
});

// Map click handler
map.on('click', (e) => {
    const latlng = e.latlng;

    if (currentTool === 'node') {
        addNode(latlng, currentColor);
    } else if (currentTool === 'switch') {
        addSwitch(latlng);
    } else if (currentTool === 'line') {
        handleLineDrawing(latlng);
    } else if (currentTool === 'ruler') {
        handleRulerMeasurement(latlng);
    }
});

// Map right-click to perform current tool action
map.on('contextmenu', (e) => {
    L.DomEvent.preventDefault(e);
    const latlng = e.latlng;

    if (currentTool === 'select' || currentTool === 'move') {
        // Allow map panning with right-click in select and move mode
        return;
    } else if (currentTool === 'line') {
        if (drawingLine) {
            finishLineDrawing();
        } else {
            handleLineDrawing(latlng);
        }
    } else if (currentTool === 'node') {
        addNode(latlng, currentColor);
    } else if (currentTool === 'switch') {
        addSwitch(latlng);
    } else if (currentTool === 'ruler') {
        handleRulerMeasurement(latlng);
    }
});

// Add node (now a point marker)
function addNode(latlng, color) {
    const marker = L.circleMarker(latlng, {
        radius: 6,
        fillColor: color,
        color: '#fff',
        weight: 2,
        fillOpacity: 1,
        pane: 'nodesPane',
        interactive: true
    }).addTo(layerGroup);

    const nodeData = {
        id: Date.now(),
        latlng: latlng,
        color: color,
        layer: marker
    };

    // Track attached lines
    marker._attachedLines = [];
    marker._nodeData = nodeData;

    updateNodePopup(nodeData);

    marker.on('click', (e) => {
        L.DomEvent.stopPropagation(e);
        if (currentTool === 'select') {
            selectObject(nodeData, 'node');
        } else if (currentTool === 'move') {
            // Start dragging the node
            const startLatLng = marker.getLatLng();
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                nodeData.latlng = newLatLng;

                // Update all attached lines
                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else if (currentTool === 'line') {
            // Use this node as a line point
            handleLineDrawing(marker.getLatLng(), nodeData);
            // Auto-finish the line when clicking on a node
            if (drawingLine && linePoints.length >= 2) {
                finishLineDrawing();
            }
        } else if (currentTool === 'delete') {
            // Delete the node
            if (confirm('Are you sure you want to delete this node?')) {
                // Remove all attached lines
                if (marker._attachedLines && marker._attachedLines.length > 0) {
                    const linesToDelete = [...marker._attachedLines];
                    linesToDelete.forEach(lineInfo => {
                        const lineData = lineInfo.line;
                        layerGroup.removeLayer(lineData.layer);
                        lineData.markers.forEach(m => layerGroup.removeLayer(m));
                        lineData.labels.forEach(l => layerGroup.removeLayer(l));
                        lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
                        drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);
                    });
                }

                // Remove the node
                layerGroup.removeLayer(marker);
                drawnObjects.nodes = drawnObjects.nodes.filter(n => n.id !== nodeData.id);
                updateNodeCount();
                saveToLocalStorage();
            }
        }
    });

    marker.on('contextmenu', (e) => {
        L.DomEvent.stopPropagation(e);
        L.DomEvent.preventDefault(e);

        if (currentTool === 'move') {
            // Move the node with right-click in move mode
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                nodeData.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else {
            // Delete the node in other modes
            if (confirm('Are you sure you want to delete this node?')) {
                layerGroup.removeLayer(marker);
                drawnObjects.nodes = drawnObjects.nodes.filter(n => n.id !== nodeData.id);
                if (selectedObject && selectedObject.data.id === nodeData.id) {
                    selectedObject = null;
                }
                updateNodeCount();
                saveToLocalStorage();
            }
        }
    });

    // Make marker draggable in move mode
    marker.options.draggable = false;

    drawnObjects.nodes.push(nodeData);
    updateNodeCount();
    saveToLocalStorage();
}

function updateNodePopup(nodeData) {
    const colorOptions = [
        '#FF0000', '#00FF00', '#0000FF', '#FFFF00',
        '#FF00FF', '#00FFFF', '#FFA500', '#800080'
    ];

    const colorButtons = colorOptions.map(color =>
        `<button onclick="window.changeNodeColorFromPopup('${nodeData.id}', '${color}')"
                 style="width: 25px; height: 25px; background: ${color}; border: 2px solid ${nodeData.color === color ? '#000' : '#fff'};
                        cursor: pointer; margin: 2px; border-radius: 3px;"></button>`
    ).join('');

    nodeData.layer.bindPopup(`
        <div style="min-width: 150px;">
            <strong>Node</strong><br>
            <div style="margin: 8px 0;">
                <strong>Change Color:</strong><br>
                <div style="display: flex; flex-wrap: wrap; gap: 2px; margin-top: 5px;">
                    ${colorButtons}
                </div>
            </div>
            <button onclick="window.deleteObject('${nodeData.id}', 'node')"
                    style="width: 100%; padding: 8px; background: #e74c3c; color: white; border: none;
                           cursor: pointer; border-radius: 4px; margin-top: 8px;">
                Delete Node
            </button>
        </div>
    `);
}

function changeNodeColor(nodeData, newColor) {
    nodeData.color = newColor;
    nodeData.layer.setStyle({ fillColor: newColor });
    updateNodePopup(nodeData);
    updateNodeCount();
    saveToLocalStorage();
}

// Add switch
function addSwitch(latlng) {
    const ports = prompt("How many ports for this switch?", "8");
    if (!ports || isNaN(ports) || parseInt(ports) < 1) return;

    const portCount = parseInt(ports);

    const marker = L.marker(latlng, {
        icon: L.divIcon({
            className: 'switch-marker',
            html: `<div style="width: 30px; height: 30px; background: #2c3e50; border: 3px solid #3498db; border-radius: 5px; display: flex; align-items: center; justify-content: center; color: white; font-weight: bold; font-size: 10px;">${portCount}</div>`,
            iconSize: [30, 30],
            iconAnchor: [15, 15]
        }),
        pane: 'switchesPane'
    }).addTo(layerGroup);

    const switchData = {
        id: Date.now(),
        latlng: latlng,
        ports: portCount,
        layer: marker
    };

    // Track attached lines
    marker._attachedLines = [];
    marker._switchData = switchData;
    marker._maxPorts = portCount;

    updateSwitchPopup(switchData);

    marker.on('click', (e) => {
        L.DomEvent.stopPropagation(e);
        if (currentTool === 'select') {
            selectObject(switchData, 'switch');
        } else if (currentTool === 'move') {
            // Close popup for move mode
            marker.closePopup();

            // Start dragging the switch
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                switchData.latlng = newLatLng;

                // Update all attached lines
                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else if (currentTool === 'line') {
            // Close popup for line mode
            marker.closePopup();

            // Check port limit
            if (marker._attachedLines.length >= marker._maxPorts) {
                alert(`This switch only has ${marker._maxPorts} ports and they are all in use!`);
                return;
            }
            handleLineDrawing(marker.getLatLng(), switchData);
            // Auto-finish the line when clicking on a switch
            if (drawingLine && linePoints.length >= 2) {
                finishLineDrawing();
            }
        } else if (currentTool === 'delete') {
            // Delete the switch
            marker.closePopup();
            if (confirm('Are you sure you want to delete this switch?')) {
                // Remove all attached lines
                if (marker._attachedLines && marker._attachedLines.length > 0) {
                    const linesToDelete = [...marker._attachedLines];
                    linesToDelete.forEach(lineInfo => {
                        const lineData = lineInfo.line;
                        layerGroup.removeLayer(lineData.layer);
                        lineData.markers.forEach(m => layerGroup.removeLayer(m));
                        lineData.labels.forEach(l => layerGroup.removeLayer(l));
                        lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
                        drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);
                    });
                }

                // Remove the switch
                layerGroup.removeLayer(marker);
                drawnObjects.switches = drawnObjects.switches.filter(s => s.id !== switchData.id);
                saveToLocalStorage();
            }
        }
    });

    marker.on('contextmenu', (e) => {
        L.DomEvent.stopPropagation(e);
        L.DomEvent.preventDefault(e);

        if (currentTool === 'move') {
            // Move the switch with right-click in move mode
            marker.closePopup();
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                switchData.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else {
            // Delete the switch in other modes
            if (confirm('Are you sure you want to delete this switch?')) {
                layerGroup.removeLayer(marker);
                drawnObjects.switches = drawnObjects.switches.filter(s => s.id !== switchData.id);
                if (selectedObject && selectedObject.data.id === switchData.id) {
                    selectedObject = null;
                }
                saveToLocalStorage();
            }
        }
    });

    drawnObjects.switches.push(switchData);
    saveToLocalStorage();
}

function updateSwitchPopup(switchData) {
    const usedPorts = switchData.layer._attachedLines ? switchData.layer._attachedLines.length : 0;

    switchData.layer.bindPopup(`
        <div style="min-width: 150px;">
            <strong>Network Switch</strong><br>
            <div style="margin: 8px 0;">
                Ports: <input type="number" id="editPorts_${switchData.id}" value="${switchData.ports}" min="${usedPorts}" style="width: 50px;">
                <button onclick="window.updateSwitchPorts('${switchData.id}')"
                        style="padding: 2px 8px; cursor: pointer; background: #3498db; color: white; border: none; border-radius: 3px;">
                    Update
                </button>
            </div>
            Used: ${usedPorts} / ${switchData.ports}<br>
            Available: ${switchData.ports - usedPorts}<br>
            <button onclick="window.deleteObject('${switchData.id}', 'switch')"
                    style="width: 100%; padding: 8px; background: #e74c3c; color: white; border: none;
                           cursor: pointer; border-radius: 4px; margin-top: 8px;">
                Delete Switch
            </button>
        </div>
    `);
}

// Global function for popup buttons
window.changeNodeColorFromPopup = function(nodeId, color) {
    const node = drawnObjects.nodes.find(n => n.id == nodeId);
    if (node) {
        changeNodeColor(node, color);
    }
};

window.updateSwitchPorts = function(switchId) {
    const switchObj = drawnObjects.switches.find(s => s.id == switchId);
    if (!switchObj) return;

    const newPorts = parseInt(document.getElementById(`editPorts_${switchId}`).value);
    const usedPorts = switchObj.layer._attachedLines ? switchObj.layer._attachedLines.length : 0;

    if (isNaN(newPorts) || newPorts < 1) {
        alert("Switch must have at least 1 port.");
        return;
    }

    if (newPorts < usedPorts) {
        alert(`Cannot set ports to ${newPorts}. ${usedPorts} ports are currently in use.`);
        return;
    }

    switchObj.ports = newPorts;
    switchObj.layer._maxPorts = newPorts;

    // Update icon
    switchObj.layer.setIcon(L.divIcon({
        className: 'switch-marker',
        html: `<div style="width: 30px; height: 30px; background: #2c3e50; border: 3px solid #3498db; border-radius: 5px; display: flex; align-items: center; justify-content: center; color: white; font-weight: bold; font-size: 10px;">${newPorts}</div>`,
        iconSize: [30, 30],
        iconAnchor: [15, 15]
    }));

    updateSwitchPopup(switchObj);
    saveToLocalStorage();
};

window.deleteObject = function(objId, type) {
    if (!confirm(`Are you sure you want to delete this ${type}?`)) {
        return;
    }

    if (type === 'node') {
        const node = drawnObjects.nodes.find(n => n.id == objId);
        if (node) {
            layerGroup.removeLayer(node.layer);
            drawnObjects.nodes = drawnObjects.nodes.filter(n => n.id != objId);
            updateNodeCount();
            saveToLocalStorage();
        }
    } else if (type === 'switch') {
        const switchObj = drawnObjects.switches.find(s => s.id == objId);
        if (switchObj) {
            layerGroup.removeLayer(switchObj.layer);
            drawnObjects.switches = drawnObjects.switches.filter(s => s.id != objId);
            saveToLocalStorage();
        }
    } else if (type === 'line') {
        const line = drawnObjects.lines.find(l => l.id == objId);
        if (line) {
            layerGroup.removeLayer(line.layer);
            line.markers.forEach(m => {
                // Remove from attached nodes/switches
                if (m._attachedNode && m._attachedNode.layer) {
                    m._attachedNode.layer._attachedLines =
                        m._attachedNode.layer._attachedLines.filter(l => l.line.id !== line.id);
                    if (m._attachedNode.layer._switchData) {
                        updateSwitchPopup(m._attachedNode.layer._switchData);
                    }
                }
                layerGroup.removeLayer(m);
            });
            line.labels.forEach(l => layerGroup.removeLayer(l));
            line.segments.forEach(s => layerGroup.removeLayer(s.layer));
            drawnObjects.lines = drawnObjects.lines.filter(l => l.id != objId);
            saveToLocalStorage();
        }
    }

    if (selectedObject && selectedObject.data.id == objId) {
        selectedObject = null;
    }
};

// Handle multi-segment line drawing
let attachedNodes = []; // Track which points are attached to nodes

function handleLineDrawing(latlng, nodeData = null) {
    if (!drawingLine) {
        // Start new line
        drawingLine = true;
        linePoints = [latlng];
        attachedNodes = [nodeData]; // Store node reference (can be null)
        tempMarkers = [];

        // Create start marker
        const marker = L.circleMarker(latlng, {
            radius: 5,
            fillColor: '#3498db',
            color: '#fff',
            weight: 2,
            fillOpacity: 1,
            draggable: false
        }).addTo(layerGroup);

        tempMarkers.push(marker);
    } else {
        // Add point to line
        linePoints.push(latlng);
        attachedNodes.push(nodeData); // Store node reference for this point

        // Create marker for this point
        const marker = L.circleMarker(latlng, {
            radius: 5,
            fillColor: '#3498db',
            color: '#fff',
            weight: 2,
            fillOpacity: 1,
            draggable: false
        }).addTo(layerGroup);

        tempMarkers.push(marker);

        // Update or create temporary polyline
        if (window.tempPolyline) {
            layerGroup.removeLayer(window.tempPolyline);
        }

        window.tempPolyline = L.polyline(linePoints, {
            color: '#3498db',
            weight: 3,
            opacity: 0.7
        }).addTo(layerGroup);
    }
}

// Helper function to find nearby node or switch and snap to it
function findNearbyObject(latlng, snapDistance = 20) {
    const point = map.latLngToContainerPoint(latlng);

    // Check nodes
    for (const node of drawnObjects.nodes) {
        const nodePoint = map.latLngToContainerPoint(node.latlng);
        const distance = Math.sqrt(
            Math.pow(point.x - nodePoint.x, 2) +
            Math.pow(point.y - nodePoint.y, 2)
        );
        if (distance <= snapDistance) {
            return { type: 'node', object: node, latlng: node.latlng };
        }
    }

    // Check switches
    for (const switchObj of drawnObjects.switches) {
        const switchPoint = map.latLngToContainerPoint(switchObj.latlng);
        const distance = Math.sqrt(
            Math.pow(point.x - switchPoint.x, 2) +
            Math.pow(point.y - switchPoint.y, 2)
        );
        if (distance <= snapDistance) {
            return { type: 'switch', object: switchObj, latlng: switchObj.latlng };
        }
    }

    return null;
}

function finishLineDrawing() {
    if (!drawingLine || linePoints.length < 2) {
        // Clean up if we didn't make a valid line
        tempMarkers.forEach(m => layerGroup.removeLayer(m));
        if (window.tempPolyline) {
            layerGroup.removeLayer(window.tempPolyline);
            window.tempPolyline = null;
        }
        tempMarkers = [];
        linePoints = [];
        attachedNodes = [];
        drawingLine = false;
        window.extendingLine = null;
        return;
    }

    // Remove temporary markers and polyline
    tempMarkers.forEach(m => layerGroup.removeLayer(m));
    if (window.tempPolyline) {
        layerGroup.removeLayer(window.tempPolyline);
        window.tempPolyline = null;
    }
    tempMarkers = [];

    // Snap endpoints to nearby nodes or switches
    const firstPoint = linePoints[0];
    const lastPoint = linePoints[linePoints.length - 1];

    const nearbyFirst = findNearbyObject(firstPoint);
    if (nearbyFirst && !attachedNodes[0]) {
        linePoints[0] = nearbyFirst.latlng;
        if (nearbyFirst.type === 'node') {
            attachedNodes[0] = nearbyFirst.object;
        } else if (nearbyFirst.type === 'switch') {
            // Attach to switch
            const usedPorts = nearbyFirst.object.layer._attachedLines ? nearbyFirst.object.layer._attachedLines.length : 0;
            if (usedPorts < nearbyFirst.object.ports) {
                attachedNodes[0] = nearbyFirst.object;
            }
        }
    }

    const nearbyLast = findNearbyObject(lastPoint);
    const lastIndex = linePoints.length - 1;
    if (nearbyLast && !attachedNodes[lastIndex]) {
        linePoints[lastIndex] = nearbyLast.latlng;
        if (nearbyLast.type === 'node') {
            attachedNodes[lastIndex] = nearbyLast.object;
        } else if (nearbyLast.type === 'switch') {
            // Attach to switch
            const usedPorts = nearbyLast.object.layer._attachedLines ? nearbyLast.object.layer._attachedLines.length : 0;
            if (usedPorts < nearbyLast.ports) {
                attachedNodes[lastIndex] = nearbyLast.object;
            }
        }
    }

    // Check if we're extending an existing line
    if (window.extendingLine) {
        const { lineData, fromStart } = window.extendingLine;

        // Remove the old line
        layerGroup.removeLayer(lineData.layer);
        lineData.markers.forEach(m => {
            if (m._attachedNode && m._attachedNode.layer) {
                m._attachedNode.layer._attachedLines =
                    m._attachedNode.layer._attachedLines.filter(l => l.line.id !== lineData.id);
            }
            layerGroup.removeLayer(m);
        });
        lineData.labels.forEach(l => layerGroup.removeLayer(l));
        lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
        drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);

        // Create new line with extended points
        if (fromStart) {
            // Reverse the new points and prepend to existing
            linePoints.reverse();
            attachedNodes.reverse();
        }

        createEditableLine(linePoints, attachedNodes);
        window.extendingLine = null;
    } else {
        // Create the final line with draggable endpoints
        createEditableLine(linePoints, attachedNodes);
    }

    // Reset state
    linePoints = [];
    attachedNodes = [];
    drawingLine = false;
}

function createEditableLine(points, nodeAttachments = []) {
    const polyline = L.polyline(points, {
        color: '#3498db',
        weight: 3,
        pane: 'linesPane'
    }).addTo(layerGroup);

    // Calculate total distance
    let totalDistance = 0;
    for (let i = 0; i < points.length - 1; i++) {
        totalDistance += map.distance(points[i], points[i + 1]);
    }

    const distanceKm = (totalDistance / 1000).toFixed(2);
    const distanceM = totalDistance.toFixed(2);

    // Create draggable endpoint markers
    const markers = points.map((point, index) => {
        const attachedNode = nodeAttachments[index];

        const marker = L.marker(point, {
            draggable: false,
            icon: L.divIcon({
                className: 'endpoint-marker',
                html: '<div style="width: 12px; height: 12px; background: #3498db; border: 2px solid white; border-radius: 50%;"></div>',
                iconSize: [12, 12],
                iconAnchor: [6, 6]
            }),
            pane: 'linesPane'
        }).addTo(layerGroup);

        // Store reference to line data and attached node
        marker._lineData = null; // Will be set after lineData is created
        marker._attachedNode = attachedNode;

        marker.on('click', (e) => {
            L.DomEvent.stopPropagation(e);
            if (currentTool === 'move') {
                marker.dragging.enable();
            } else if (currentTool === 'line') {
                // Continue the line from this endpoint
                const lineData = marker._lineData;
                if (!lineData) return;

                // Check if this is the first or last point
                const isFirstPoint = index === 0;
                const isLastPoint = index === lineData.points.length - 1;

                if (isFirstPoint || isLastPoint) {
                    // Start drawing from this endpoint
                    drawingLine = true;
                    linePoints = [...lineData.points];
                    attachedNodes = lineData.markers.map(m => m._attachedNode);
                    tempMarkers = [];

                    // Store the line we're extending
                    window.extendingLine = {
                        lineData: lineData,
                        fromStart: isFirstPoint
                    };
                }
            } else if (currentTool === 'delete') {
                // Delete the line endpoint (and entire line if only 2 points)
                const lineData = marker._lineData;
                if (!lineData) return;

                if (lineData.points.length === 2) {
                    if (confirm('Are you sure you want to delete this line?')) {
                        layerGroup.removeLayer(lineData.layer);
                        lineData.markers.forEach(m => {
                            if (m._attachedNode && m._attachedNode.layer) {
                                m._attachedNode.layer._attachedLines =
                                    m._attachedNode.layer._attachedLines.filter(l => l.line.id !== lineData.id);
                            }
                            layerGroup.removeLayer(m);
                        });
                        lineData.labels.forEach(l => layerGroup.removeLayer(l));
                        lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
                        drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);
                        saveToLocalStorage();
                    }
                } else {
                    if (confirm('Are you sure you want to delete this line point?')) {
                        // Remove from attached node if exists
                        if (marker._attachedNode && marker._attachedNode.layer) {
                            marker._attachedNode.layer._attachedLines =
                                marker._attachedNode.layer._attachedLines.filter(
                                    info => info.line.id !== lineData.id || info.pointIndex !== index
                                );
                        }

                        lineData.points.splice(index, 1);
                        layerGroup.removeLayer(marker);
                        lineData.markers.splice(index, 1);
                        lineData.layer.setLatLngs(lineData.points);

                        // Update all marker references
                        lineData.markers.forEach((m, newIndex) => {
                            m._lineData = lineData;
                            if (m._attachedNode && m._attachedNode.layer) {
                                m._attachedNode.layer._attachedLines =
                                    m._attachedNode.layer._attachedLines.map(info => {
                                        if (info.line.id === lineData.id) {
                                            return { ...info, pointIndex: newIndex };
                                        }
                                        return info;
                                    });
                            }
                        });

                        updateLineLabels(lineData);
                        saveToLocalStorage();
                    }
                }
            }
        });

        marker.on('contextmenu', (e) => {
            L.DomEvent.stopPropagation(e);
            L.DomEvent.preventDefault(e);

            const lineData = marker._lineData;
            if (!lineData) return;

            // Remove from attached node if exists
            if (marker._attachedNode && marker._attachedNode.layer) {
                marker._attachedNode.layer._attachedLines =
                    marker._attachedNode.layer._attachedLines.filter(
                        info => info.line.id !== lineData.id || info.pointIndex !== index
                    );
            }

            // If line has only 2 points, delete the entire line
            if (lineData.points.length === 2) {
                if (confirm('Are you sure you want to delete this line?')) {
                    layerGroup.removeLayer(lineData.layer);
                    lineData.markers.forEach(m => {
                        if (m._attachedNode && m._attachedNode.layer) {
                            m._attachedNode.layer._attachedLines =
                                m._attachedNode.layer._attachedLines.filter(l => l.line.id !== lineData.id);
                        }
                        layerGroup.removeLayer(m);
                    });
                    lineData.labels.forEach(l => layerGroup.removeLayer(l));
                    lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
                    drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);
                    saveToLocalStorage();
                }
            } else {
                // Remove this point from the line
                if (confirm('Are you sure you want to delete this line point?')) {
                    lineData.points.splice(index, 1);

                    // Remove this marker
                    layerGroup.removeLayer(marker);
                    lineData.markers.splice(index, 1);

                    // Update the line
                    lineData.layer.setLatLngs(lineData.points);

                    // Update all marker references
                    lineData.markers.forEach((m, newIndex) => {
                        m._lineData = lineData;
                        // Update attached node references with new index
                        if (m._attachedNode && m._attachedNode.layer) {
                            m._attachedNode.layer._attachedLines =
                                m._attachedNode.layer._attachedLines.map(info => {
                                    if (info.line.id === lineData.id) {
                                        return { ...info, pointIndex: newIndex };
                                    }
                                    return info;
                                });
                        }
                    });

                    updateLineLabels(lineData);
                    saveToLocalStorage();
                }
            }
        });

        // Update line when marker is dragged
        marker.on('drag', () => {
            const lineData = marker._lineData;
            if (lineData) {
                lineData.points[index] = marker.getLatLng();
                polyline.setLatLngs(lineData.points);
                updateLineLabels(lineData);
            }
        });

        marker.on('dragend', () => {
            marker.dragging.disable();
            saveToLocalStorage();
        });

        return marker;
    });

    // Add distance labels for each segment
    const labels = [];
    const segments = [];
    for (let i = 0; i < points.length - 1; i++) {
        const segmentDistance = map.distance(points[i], points[i + 1]);
        const segmentM = segmentDistance.toFixed(2);

        const midpoint = L.latLng(
            (points[i].lat + points[i + 1].lat) / 2,
            (points[i].lng + points[i + 1].lng) / 2
        );

        const label = L.marker(midpoint, {
            icon: L.divIcon({
                className: 'distance-label',
                html: `${segmentM}m`,
                iconSize: null
            }),
            interactive: true,
            opacity: 0
        }).addTo(layerGroup);

        // Create segment polyline for sum tool (make it wider for easier hovering)
        const segmentLine = L.polyline([points[i], points[i + 1]], {
            color: '#3498db',
            weight: 10,
            opacity: 0,
            interactive: true,
            pane: 'segmentsPane'
        }).addTo(layerGroup);

        const segmentData = {
            index: i,
            distance: segmentDistance,
            layer: segmentLine,
            points: [points[i], points[i + 1]],
            label: label
        };

        // Add hover handlers to show/hide label
        segmentLine.on('mouseover', (e) => {
            label.setOpacity(1);
        });

        segmentLine.on('mouseout', (e) => {
            label.setOpacity(0);
        });

        // Add click handler to the segment line
        segmentLine.on('click', (e) => {
            L.DomEvent.stopPropagation(e);
            if (currentTool === 'sum') {
                toggleSegmentSelection(lineData, segmentData);
            } else if (currentTool === 'select') {
                selectObject(lineData, 'line');
            }
        });

        // Add click handler to the label as well
        label.on('click', (e) => {
            L.DomEvent.stopPropagation(e);
            if (currentTool === 'sum') {
                toggleSegmentSelection(lineData, segmentData);
            } else if (currentTool === 'select') {
                selectObject(lineData, 'line');
            }
        });

        // Keep label visible when hovering over it
        label.on('mouseover', (e) => {
            label.setOpacity(1);
        });

        label.on('mouseout', (e) => {
            label.setOpacity(0);
        });

        segments.push(segmentData);
        labels.push(label);
    }

    const lineData = {
        id: Date.now(),
        points: points,
        distance: totalDistance,
        layer: polyline,
        markers: markers,
        labels: labels,
        segments: segments
    };

    // Set line data reference in markers and register with attached nodes
    markers.forEach((m, i) => {
        m._lineData = lineData;

        // If this marker is attached to a node, register the line with that node
        if (m._attachedNode && m._attachedNode.layer) {
            if (!m._attachedNode.layer._attachedLines) {
                m._attachedNode.layer._attachedLines = [];
            }
            m._attachedNode.layer._attachedLines.push({
                line: lineData,
                pointIndex: i
            });
            // Update switch popup if attached to a switch
            if (m._attachedNode.layer._switchData) {
                updateSwitchPopup(m._attachedNode.layer._switchData);
            }
        }
    });

    polyline.bindPopup(`
        <div>
            <strong>Line</strong><br>
            Total Distance: ${distanceM}m<br>
            ${distanceKm}km<br>
            Segments: ${points.length - 1}<br>
            <button onclick="window.deleteObject('${lineData.id}', 'line')"
                    style="width: 100%; padding: 8px; background: #e74c3c; color: white; border: none;
                           cursor: pointer; border-radius: 4px; margin-top: 8px;">
                Delete Line
            </button>
        </div>
    `);

    polyline.on('click', (e) => {
        L.DomEvent.stopPropagation(e);
        if (currentTool === 'select') {
            selectObject(lineData, 'line');
        }
    });

    polyline.on('contextmenu', (e) => {
        L.DomEvent.stopPropagation(e);
        L.DomEvent.preventDefault(e);
        // Delete the line
        if (confirm('Are you sure you want to delete this line?')) {
            layerGroup.removeLayer(polyline);
            lineData.markers.forEach(m => {
                // Remove from attached nodes
                const attachedNode = m._attachedNode;
                if (attachedNode) {
                    attachedNode._attachedLines = attachedNode._attachedLines.filter(l => l !== lineData);
                }
                layerGroup.removeLayer(m);
            });
            lineData.labels.forEach(l => layerGroup.removeLayer(l));
            lineData.segments.forEach(s => layerGroup.removeLayer(s.layer));
            drawnObjects.lines = drawnObjects.lines.filter(l => l.id !== lineData.id);
            if (selectedObject && selectedObject.data.id === lineData.id) {
                selectedObject = null;
            }
            saveToLocalStorage();
        }
    });

    // Make polyline draggable in move mode
    polyline.on('mousedown', (e) => {
        if (currentTool === 'move') {
            L.DomEvent.stopPropagation(e);
            const startLatLng = e.latlng;
            const originalPoints = lineData.points.map(p => ({lat: p.lat, lng: p.lng}));

            const onMouseMove = (e) => {
                const latDiff = e.latlng.lat - startLatLng.lat;
                const lngDiff = e.latlng.lng - startLatLng.lng;

                lineData.points = originalPoints.map(p =>
                    L.latLng(p.lat + latDiff, p.lng + lngDiff)
                );

                polyline.setLatLngs(lineData.points);
                lineData.markers.forEach((marker, i) => {
                    marker.setLatLng(lineData.points[i]);
                });
                updateLineLabels(lineData);
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                saveToLocalStorage();
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        }
    });

    drawnObjects.lines.push(lineData);
    saveToLocalStorage();
}

function updateLineLabels(lineData) {
    const points = lineData.points;

    // Remove old labels and segments
    lineData.labels.forEach(label => layerGroup.removeLayer(label));
    lineData.labels = [];

    if (lineData.segments) {
        lineData.segments.forEach(seg => layerGroup.removeLayer(seg.layer));
    }
    lineData.segments = [];

    // Calculate new total distance
    let totalDistance = 0;
    for (let i = 0; i < points.length - 1; i++) {
        totalDistance += map.distance(points[i], points[i + 1]);
    }
    lineData.distance = totalDistance;

    // Create new labels and segments
    for (let i = 0; i < points.length - 1; i++) {
        const segmentDistance = map.distance(points[i], points[i + 1]);
        const segmentM = segmentDistance.toFixed(2);

        const midpoint = L.latLng(
            (points[i].lat + points[i + 1].lat) / 2,
            (points[i].lng + points[i + 1].lng) / 2
        );

        const label = L.marker(midpoint, {
            icon: L.divIcon({
                className: 'distance-label',
                html: `${segmentM}m`,
                iconSize: null
            }),
            interactive: true,
            opacity: 0
        }).addTo(layerGroup);

        // Recreate segment polyline for sum tool (make it wider for easier hovering)
        const segmentLine = L.polyline([points[i], points[i + 1]], {
            color: '#3498db',
            weight: 10,
            opacity: 0,
            interactive: true,
            pane: 'segmentsPane'
        }).addTo(layerGroup);

        const segmentData = {
            index: i,
            distance: segmentDistance,
            layer: segmentLine,
            points: [points[i], points[i + 1]],
            label: label
        };

        // Add hover handlers to show/hide label
        segmentLine.on('mouseover', (e) => {
            label.setOpacity(1);
        });

        segmentLine.on('mouseout', (e) => {
            label.setOpacity(0);
        });

        // Add click handler to the segment line
        segmentLine.on('click', (e) => {
            L.DomEvent.stopPropagation(e);
            if (currentTool === 'sum') {
                toggleSegmentSelection(lineData, segmentData);
            } else if (currentTool === 'select') {
                selectObject(lineData, 'line');
            }
        });

        // Add click handler to the label as well
        label.on('click', (e) => {
            L.DomEvent.stopPropagation(e);
            if (currentTool === 'sum') {
                toggleSegmentSelection(lineData, segmentData);
            } else if (currentTool === 'select') {
                selectObject(lineData, 'line');
            }
        });

        // Keep label visible when hovering over it
        label.on('mouseover', (e) => {
            label.setOpacity(1);
        });

        label.on('mouseout', (e) => {
            label.setOpacity(0);
        });

        lineData.segments.push(segmentData);
        lineData.labels.push(label);
    }

    // Update popup
    const distanceKm = (totalDistance / 1000).toFixed(2);
    const distanceM = totalDistance.toFixed(2);
    lineData.layer.setPopupContent(`
        <div>
            <strong>Line</strong><br>
            Total Distance: ${distanceM}m<br>
            ${distanceKm}km<br>
            Segments: ${points.length - 1}<br>
            <button onclick="window.deleteObject('${lineData.id}', 'line')"
                    style="width: 100%; padding: 8px; background: #e74c3c; color: white; border: none;
                           cursor: pointer; border-radius: 4px; margin-top: 8px;">
                Delete Line
            </button>
        </div>
    `);
}

// Handle ruler measurement
function handleRulerMeasurement(latlng) {
    rulerPoints.push(latlng);

    if (rulerPoints.length === 1) {
        document.getElementById('rulerDisplay').style.display = 'block';
        document.getElementById('rulerDisplay').textContent = 'Click second point...';
    } else if (rulerPoints.length === 2) {
        const distance = map.distance(rulerPoints[0], rulerPoints[1]);
        const distanceKm = (distance / 1000).toFixed(2);
        const distanceM = distance.toFixed(2);

        document.getElementById('rulerDisplay').innerHTML = `
            Distance: ${distanceM}m<br>
            ${distanceKm}km
        `;

        const tempLine = L.polyline(rulerPoints, {
            color: '#f39c12',
            weight: 2,
            dashArray: '5, 10'
        }).addTo(layerGroup);

        setTimeout(() => {
            layerGroup.removeLayer(tempLine);
            rulerPoints = [];
            document.getElementById('rulerDisplay').style.display = 'none';
        }, 3000);
    }
}

// Sum tool functions
function toggleSegmentSelection(lineData, segmentData) {
    const segmentId = `${lineData.id}-${segmentData.index}`;

    if (selectedSegments.has(segmentId)) {
        // Deselect
        selectedSegments.delete(segmentId);
        segmentData.layer.setStyle({ opacity: 0, weight: 3 });
    } else {
        // Select
        selectedSegments.add(segmentId);
        segmentData.layer.setStyle({ opacity: 0.5, weight: 5, color: '#27ae60' });
    }

    updateSumDisplay();
}

function updateSumDisplay() {
    if (!sumDisplay) return;

    let totalDistance = 0;
    drawnObjects.lines.forEach(line => {
        line.segments.forEach(segment => {
            const segmentId = `${line.id}-${segment.index}`;
            if (selectedSegments.has(segmentId)) {
                totalDistance += segment.distance;
            }
        });
    });

    const distanceKm = (totalDistance / 1000).toFixed(2);
    const distanceM = totalDistance.toFixed(2);
    const count = selectedSegments.size;

    sumDisplay.innerHTML = `
        <strong>Sum Tool</strong><br>
        Segments: ${count}<br>
        Total: ${distanceM}m<br>
        ${distanceKm}km
    `;
}

function showSumDisplay() {
    if (!sumDisplay) {
        sumDisplay = document.getElementById('rulerDisplay');
    }
    sumDisplay.style.display = 'block';
    updateSumDisplay();
}

function clearSumSelections() {
    drawnObjects.lines.forEach(line => {
        if (line.segments) {
            line.segments.forEach(segment => {
                segment.layer.setStyle({ opacity: 0, weight: 3, color: '#3498db' });
            });
        }
    });
    selectedSegments.clear();
    if (sumDisplay) {
        sumDisplay.style.display = 'none';
    }
}

// Select object
function selectObject(objData, type) {
    if (selectedObject) {
        // Deselect previous
        if (selectedObject.type === 'node') {
            selectedObject.data.layer.setStyle({ weight: 2 });
        } else if (selectedObject.type === 'line') {
            selectedObject.data.layer.setStyle({ weight: 3 });
        }
    }

    selectedObject = { data: objData, type: type };

    // Highlight selected
    if (type === 'node') {
        objData.layer.setStyle({ weight: 4 });
    } else if (type === 'line') {
        objData.layer.setStyle({ weight: 5 });
    }
}

// Clear all
// Export button - copy base64 encoded localStorage to clipboard
document.getElementById('exportBtn').addEventListener('click', () => {
    const data = localStorage.getItem('mapMakerData');
    console.log('Exporting data:', data);
    if (!data) {
        alert('No data to export!');
        return;
    }
    const base64 = btoa(data);
    console.log('Base64 encoded:', base64);

    // Try to copy to clipboard
    navigator.clipboard.writeText(base64).then(() => {
        // Show in prompt as well so user can manually copy if needed
        prompt('Data copied to clipboard! You can also copy it manually from here:', base64);
    }).catch((err) => {
        console.error('Clipboard error:', err);
        // Fallback if clipboard API fails
        prompt('Copy this export code:', base64);
    });
});

// Import button - prompt for base64 string and load it
document.getElementById('importBtn').addEventListener('click', () => {
    const base64 = prompt('Paste the export code here:');
    console.log('Import base64:', base64);
    if (!base64 || base64.trim() === '') {
        console.log('Import cancelled - no data');
        return;
    }

    try {
        const data = atob(base64.trim());
        console.log('Decoded data:', data);

        // Validate it's valid JSON
        const parsed = JSON.parse(data);
        console.log('Parsed data:', parsed);

        // Set a flag to prevent beforeunload from saving
        window.importingData = true;

        // Clear existing data and set new data
        console.log('Clearing localStorage...');
        localStorage.clear();
        console.log('Setting new data...');
        localStorage.setItem('mapMakerData', data);
        console.log('Data saved, current localStorage:', localStorage.getItem('mapMakerData'));

        alert('Data imported successfully! The page will now reload.');
        location.reload();
    } catch (e) {
        console.error('Import error:', e);
        alert('Invalid import code! Please make sure you copied it correctly. Error: ' + e.message);
    }
});

document.getElementById('clearBtn').addEventListener('click', () => {
    if (confirm('Are you sure you want to clear all objects?')) {
        layerGroup.clearLayers();
        drawnObjects = { lines: [], nodes: [], switches: [] };
        selectedObject = null;
        updateNodeCount();
        saveToLocalStorage();
    }
});

// Update node count display
function updateNodeCount() {
    const counts = {};
    drawnObjects.nodes.forEach(node => {
        counts[node.color] = (counts[node.color] || 0) + 1;
    });

    const countHtml = Object.entries(counts).map(([color, count]) => `
        <div class="node-count-item">
            <div class="node-count-color" style="background: ${color}"></div>
            <div class="node-count-text">${count}</div>
        </div>
    `).join('');

    document.getElementById('nodeCount').innerHTML = countHtml || '<div class="node-count-text">No nodes</div>';
}

// Update stats display
function updateStatsDisplay() {
    const statsDisplay = document.getElementById('statsDisplay');
    if (!statsDisplay) return;

    // Calculate total distance of all lines
    let totalDistance = 0;
    drawnObjects.lines.forEach(line => {
        totalDistance += line.distance || 0;
    });

    const distanceKm = (totalDistance / 1000).toFixed(2);
    const distanceM = totalDistance.toFixed(2);
    const switchCount = drawnObjects.switches.length;
    const lineCount = drawnObjects.lines.length;

    statsDisplay.innerHTML = `
        <strong>Total Stats</strong><br>
        Lines: ${lineCount}<br>
        Switches: ${switchCount}<br>
        Total Length: ${distanceM}m<br>
        (${distanceKm}km)
    `;
}

// LocalStorage functions
function saveToLocalStorage() {
    const data = {
        lines: drawnObjects.lines.map(l => ({
            id: l.id,
            points: l.points.map(p => [p.lat, p.lng]),
            distance: l.distance
        })),
        nodes: drawnObjects.nodes.map(n => ({
            id: n.id,
            latlng: [n.latlng.lat, n.latlng.lng],
            color: n.color
        })),
        switches: drawnObjects.switches.map(s => ({
            id: s.id,
            latlng: [s.latlng.lat, s.latlng.lng],
            ports: s.ports
        })),
        mapView: {
            center: map.getCenter(),
            zoom: map.getZoom(),
            isSatellite: isSatelliteView
        }
    };

    localStorage.setItem('mapMakerData', JSON.stringify(data));
    updateStatsDisplay();
}

// Helper function to restore a node
function restoreNode(nodeData, latlng) {
    const marker = L.circleMarker(latlng, {
        radius: 6,
        fillColor: nodeData.color,
        color: '#fff',
        weight: 2,
        fillOpacity: 1,
        pane: 'nodesPane',
        interactive: true
    }).addTo(layerGroup);

    const restoredNode = {
        id: nodeData.id,
        latlng: latlng,
        color: nodeData.color,
        layer: marker
    };

    marker._attachedLines = [];
    marker._nodeData = restoredNode;

    updateNodePopup(restoredNode);

    marker.on('click', (e) => {
        L.DomEvent.stopPropagation(e);
        if (currentTool === 'select') {
            selectObject(restoredNode, 'node');
        } else if (currentTool === 'move') {
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                restoredNode.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else if (currentTool === 'line') {
            handleLineDrawing(marker.getLatLng(), restoredNode);
            // Auto-finish the line when clicking on a node
            if (drawingLine && linePoints.length >= 2) {
                finishLineDrawing();
            }
        }
    });

    marker.on('contextmenu', (e) => {
        L.DomEvent.stopPropagation(e);
        L.DomEvent.preventDefault(e);

        if (currentTool === 'move') {
            // Move the node with right-click in move mode
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                restoredNode.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else {
            // Delete the node in other modes
            if (confirm('Are you sure you want to delete this node?')) {
                layerGroup.removeLayer(marker);
                drawnObjects.nodes = drawnObjects.nodes.filter(n => n.id !== restoredNode.id);
                if (selectedObject && selectedObject.data.id === restoredNode.id) {
                    selectedObject = null;
                }
                updateNodeCount();
                saveToLocalStorage();
            }
        }
    });

    drawnObjects.nodes.push(restoredNode);
}

// Helper function to restore a switch
function restoreSwitch(switchData, latlng) {
    const marker = L.marker(latlng, {
        icon: L.divIcon({
            className: 'switch-marker',
            html: `<div style="width: 30px; height: 30px; background: #2c3e50; border: 3px solid #3498db; border-radius: 5px; display: flex; align-items: center; justify-content: center; color: white; font-weight: bold; font-size: 10px;">${switchData.ports}</div>`,
            iconSize: [30, 30],
            iconAnchor: [15, 15]
        }),
        pane: 'switchesPane'
    }).addTo(layerGroup);

    const restoredSwitch = {
        id: switchData.id,
        latlng: latlng,
        ports: switchData.ports,
        layer: marker
    };

    marker._attachedLines = [];
    marker._switchData = restoredSwitch;
    marker._maxPorts = switchData.ports;

    updateSwitchPopup(restoredSwitch);

    marker.on('click', (e) => {
        L.DomEvent.stopPropagation(e);
        if (currentTool === 'select') {
            selectObject(restoredSwitch, 'switch');
        } else if (currentTool === 'move') {
            // Close popup for move mode
            marker.closePopup();

            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                restoredSwitch.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else if (currentTool === 'line') {
            // Close popup for line mode
            marker.closePopup();

            if (marker._attachedLines.length >= marker._maxPorts) {
                alert(`This switch only has ${marker._maxPorts} ports and they are all in use!`);
                return;
            }
            handleLineDrawing(marker.getLatLng(), restoredSwitch);
            // Auto-finish the line when clicking on a switch
            if (drawingLine && linePoints.length >= 2) {
                finishLineDrawing();
            }
        }
    });

    marker.on('contextmenu', (e) => {
        L.DomEvent.stopPropagation(e);
        L.DomEvent.preventDefault(e);

        if (currentTool === 'move') {
            // Move the switch with right-click in move mode
            marker.closePopup();
            let isDragging = false;

            const onMouseMove = (e) => {
                isDragging = true;
                const newLatLng = e.latlng;
                marker.setLatLng(newLatLng);
                restoredSwitch.latlng = newLatLng;

                if (marker._attachedLines) {
                    marker._attachedLines.forEach(lineInfo => {
                        const { line, pointIndex } = lineInfo;
                        line.points[pointIndex] = newLatLng;
                        line.layer.setLatLngs(line.points);
                        line.markers[pointIndex].setLatLng(newLatLng);
                        updateLineLabels(line);
                    });
                }
            };

            const onMouseUp = () => {
                map.off('mousemove', onMouseMove);
                map.off('mouseup', onMouseUp);
                if (isDragging) {
                    saveToLocalStorage();
                }
            };

            map.on('mousemove', onMouseMove);
            map.on('mouseup', onMouseUp);
        } else {
            // Delete the switch in other modes
            if (confirm('Are you sure you want to delete this switch?')) {
                layerGroup.removeLayer(marker);
                drawnObjects.switches = drawnObjects.switches.filter(s => s.id !== restoredSwitch.id);
                if (selectedObject && selectedObject.data.id === restoredSwitch.id) {
                    selectedObject = null;
                }
                saveToLocalStorage();
            }
        }
    });

    drawnObjects.switches.push(restoredSwitch);
}

function loadFromLocalStorage() {
    const saved = localStorage.getItem('mapMakerData');
    if (!saved) return;

    const data = JSON.parse(saved);

    // Map view already restored during initialization

    // Restore nodes
    if (data.nodes) {
        data.nodes.forEach(nodeData => {
            const latlng = L.latLng(nodeData.latlng[0], nodeData.latlng[1]);
            restoreNode(nodeData, latlng);
        });
        updateNodeCount();
    }

    // Restore switches
    if (data.switches) {
        data.switches.forEach(switchData => {
            const latlng = L.latLng(switchData.latlng[0], switchData.latlng[1]);
            restoreSwitch(switchData, latlng);
        });
    }

    // Restore lines
    if (data.lines) {
        data.lines.forEach(lineData => {
            const points = lineData.points.map(p => L.latLng(p[0], p[1]));
            createEditableLine(points);
        });
    }
}

// Auto-save on map move
map.on('moveend', saveToLocalStorage);
map.on('zoomend', saveToLocalStorage);

// Load saved data on startup
loadFromLocalStorage();

// Initialize stats display
updateStatsDisplay();

// Ensure data is saved before page unload (refresh, close, navigate away)
window.addEventListener('beforeunload', (e) => {
    // Don't save if we're importing data (would overwrite the imported data)
    if (window.importingData) {
        return;
    }
    saveToLocalStorage();
});

// Refresh on Escape key
document.addEventListener('keydown', (e) => {
    if (e.key === 'Escape') {
        e.preventDefault();
        e.stopPropagation();
        location.reload();
    }
}, true);