export const captureVisualization = async (): Promise<string | null> => {
    const container = document.getElementById("real-time-vis-canvas");
    if (!container) return null;

    const canvas = document.createElement("canvas");
    const ctx = canvas.getContext("2d");
    if (!ctx) return null;

    const rect = container.getBoundingClientRect();
    canvas.width = rect.width;
    canvas.height = rect.height;

    // Draw background
    ctx.fillStyle = getComputedStyle(container).backgroundColor;
    ctx.fillRect(0, 0, canvas.width, canvas.height);

    // Capture all canvas elements
    const canvases = container.querySelectorAll("canvas");
    canvases.forEach((childCanvas) => {
        const childRect = childCanvas.getBoundingClientRect();
        const x = childRect.left - rect.left;
        const y = childRect.top - rect.top;
        ctx.drawImage(childCanvas, x, y, childRect.width, childRect.height);
    });

    // Capture SVG elements
    const svgs = container.querySelectorAll("svg");
    for (const svg of Array.from(svgs)) {
        const svgString = new XMLSerializer().serializeToString(svg);
        const svgBlob = new Blob([svgString], { type: "image/svg+xml" });
        const url = URL.createObjectURL(svgBlob);

        try {
            const img = await new Promise<HTMLImageElement>((resolve, reject) => {
                const image = new Image();
                image.onload = () => resolve(image);
                image.onerror = reject;
                image.src = url;
            });

            const svgRect = svg.getBoundingClientRect();
            ctx.drawImage(
                img,
                svgRect.left - rect.left,
                svgRect.top - rect.top,
                svgRect.width,
                svgRect.height
            );
        } finally {
            URL.revokeObjectURL(url);
        }
    }

    return canvas.toDataURL("image/png");
};
