import * as THREE from "three";

const RADIUS = 0.0025;
const OPACITY = 0.8;
const GROWTH_RATE = 20.0;

export function exportHeatmapAsPNG({ bakedPoints, gl, width, height, mesh }: { bakedPoints: any[]; gl: THREE.WebGLRenderer; width: number; height: number; mesh: THREE.Mesh }) {
    const createPointTexture = (filteredPoints: typeof bakedPoints) => {
        if (filteredPoints.length === 0) return null;

        const data = new Float32Array(filteredPoints.length * 4);
        filteredPoints.forEach((p, i) => {
            const index = i * 4;
            data[index] = (p.points.x + width / 2) / width;
            data[index + 1] = (p.points.y + height / 2) / height;
            data[index + 2] = 0.3;
            data[index + 3] = 0.0;
        });

        const texture = new THREE.DataTexture(data, filteredPoints.length, 1, THREE.RGBAFormat, THREE.FloatType);
        texture.needsUpdate = true;
        return texture;
    };

    const downloadImage = (base64: string, filename: string) => {
        const link = document.createElement("a");
        link.href = base64;
        link.download = filename;
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);
    };

    const renderHeatmapToImage = (type: string) => {
        const filteredPoints = bakedPoints.filter((p) => p.type === type);
        if (filteredPoints.length === 0) return null;

        const pointTexture = createPointTexture(filteredPoints);
        if (!pointTexture) return null;

        const uniforms = {
            u_points: { value: pointTexture },
            u_count: { value: filteredPoints.length },
            u_radius: { value: RADIUS },
            u_opacity: { value: OPACITY },
            u_growthRate: { value: GROWTH_RATE },
        };

        const exportCamera = new THREE.OrthographicCamera(width / -2, width / 2, height / 2, height / -2, 0.1, 10);
        exportCamera.position.set(0, 1, 0);
        exportCamera.lookAt(0, 0, 0);

        const renderTarget = new THREE.WebGLRenderTarget(1024, 1024, {
            format: THREE.RGBAFormat,
            type: THREE.UnsignedByteType,
        });

        const tempScene = new THREE.Scene();
        tempScene.add(mesh);

        gl.setRenderTarget(renderTarget);
        gl.render(tempScene, exportCamera);
        gl.setRenderTarget(null);

        const pixels = new Uint8Array(1024 * 1024 * 4);
        gl.readRenderTargetPixels(renderTarget, 0, 0, 1024, 1024, pixels);

        const canvas = document.createElement("canvas");
        canvas.width = 1024;
        canvas.height = 1024;
        const ctx = canvas.getContext("2d");

        if (ctx) {
            const imageData = ctx.createImageData(1024, 1024);
            imageData.data.set(pixels);
            ctx.putImageData(imageData, 0, 0);
            return canvas.toDataURL("image/png");
        }

        return null;
    };

    const types = ["human", "vehicle"];
    const result = types.map((type) => {
        const image = renderHeatmapToImage(type);
        if (image) {
            downloadImage(image, `${type}-heatmap.png`);
        }
        return { type, image };
    });

    console.log("Exported Heatmaps:", result);
    return result;
}
