import * as THREE from "three";
import { useEffect, useMemo, useRef, useState } from "react";
import { useFrame, useThree } from "@react-three/fiber";
import { useSceneContext } from "../../../modules/scene/sceneContext";
import * as CONSTANTS from "../../../types/world/worldConstants";
import { determineExecutionMachineSequences } from "../../../modules/simulation/simulator/functions/determineExecutionMachineSequences";
import { useAnimationPlaySpeed, usePauseButtonStore, usePlayButtonStore, useResetButtonStore } from "../../../store/ui/usePlayButtonStore";
import { useHeatMapStore } from "../../../store/simulation/useHeatMapStore";

const DECAY_RATE = 0.01;
const GROWTH_TIME_MULTIPLIER = 20;
const RADIUS = 0.005;
const OPACITY = 0.8;
const UPDATE_INTERVAL = 1;

interface HeatPoint {
    x: number;
    y: number;
    strength: number;
    lastUpdated: number;
}

const RealTimeHeatMap = () => {
    const isUsingBBOX = false;
    const height = CONSTANTS.gridConfig.size;
    const width = CONSTANTS.gridConfig.size;
    const debugMode: "solid" | "grayscale" | "normal" = "normal";
    const debugModeMap = { solid: 0, grayscale: 1, normal: 2 } as const;
    const { productStore } = useSceneContext();
    const { getProductById, products, selectedProduct } = productStore();
    const { hasHuman, hasVehicle, monitoringHuman, monitoringVehicle, addMonitoringHuman, addMonitoringVehicle, addBakedPoint } = useHeatMapStore();
    const { isPlaying } = usePlayButtonStore();
    const { isReset } = useResetButtonStore();
    const { isPaused } = usePauseButtonStore();
    const { speed } = useAnimationPlaySpeed();
    const { scene } = useThree();

    const [events, setEvents] = useState<EventsSchema[]>([]);
    const [points, setPoints] = useState<HeatPoint[]>([]);
    const materialRef = useRef<THREE.ShaderMaterial>(null);

    const lastFrameTime = useRef<number | null>(null);
    const lastUpdateTime = useRef<number>(0);

    const uniformsRef = useRef({
        u_points: { value: null as THREE.DataTexture | null },
        u_count: { value: 0 },
        u_radius: { value: RADIUS },
        u_opacity: { value: OPACITY },
        u_debugMode: { value: debugModeMap[debugMode] },
        u_growthRate: { value: GROWTH_TIME_MULTIPLIER }, 
    });

    useEffect(() => {
        if (isReset || !isPlaying) {
            setPoints([]);
            lastFrameTime.current = null;
            lastUpdateTime.current = 0;
        }
    }, [isReset, isPlaying]);

    useEffect(() => {
        addMonitoringVehicle("26770368-55e8-4d40-87f7-8eacb48dc236");
        addMonitoringHuman("264a51e7-d8b9-4093-95ac-fa7e2dc49cfa");
    }, []);

    useEffect(() => {
        const selectedProductData = getProductById(selectedProduct.productUuid);
        const newEvents: EventsSchema[] = [];

        if (selectedProductData) {
            determineExecutionMachineSequences([selectedProductData]).then((sequences) => {
                sequences.forEach((sequence) => {
                    sequence.forEach((event) => {
                        if (event.type === "human") {
                            if (hasHuman(event.modelUuid)) {
                                newEvents.push(event);
                            }
                        } else if (event.type === "vehicle") {
                            if (hasVehicle(event.modelUuid)) {
                                newEvents.push(event);
                            }
                        }
                    });
                });
                setEvents(newEvents);
            });
        }
    }, [selectedProduct, products, monitoringHuman, monitoringVehicle]);

    useFrame((state) => {
        if (!scene || !isPlaying || isReset) return;

        const now = state.clock.elapsedTime;
        if (isPaused || speed === 0) {
            lastFrameTime.current = now;
            return;
        }

        let delta = 0;
        if (lastFrameTime.current !== null) {
            delta = now - lastFrameTime.current;
        }
        lastFrameTime.current = now;

        if (delta <= 0) return;
        const updateInterval = UPDATE_INTERVAL / Math.max(speed, 0.1);

        if (now - lastUpdateTime.current < updateInterval) return;
        lastUpdateTime.current = now;
        const scaledDelta = delta * speed;
        let updatedPoints = [...points];

        events.forEach((event) => {
            const model = scene.getObjectByProperty("uuid", event.modelUuid) as THREE.Object3D;
            if (!model) return;

            const pos = isUsingBBOX
                ? (() => {
                      const box = new THREE.Box3().setFromObject(model);
                      const { min, max } = box;
                      return new THREE.Vector3((min.x + max.x) / 2, 0, (min.z + max.z) / 2);
                  })()
                : model.position;

            addBakedPoint({
                type: event.type,
                points: { x: pos.x, y: pos.z },
            });

            updatedPoints.push({ x: pos.x, y: pos.z, strength: 0.3, lastUpdated: now });
        });

        updatedPoints = updatedPoints.map((p) => ({ ...p, strength: Math.max(0, p.strength - DECAY_RATE * scaledDelta) })).filter((p) => p.strength > 0.01);

        setPoints(updatedPoints);
    });

    const pointTexture = useMemo(() => {
        if (points.length === 0) return null;
        const data = new Float32Array(points.length * 4);
        points.forEach((p, i) => {
            const index = i * 4;
            data[index] = (p.x + width / 2) / width;
            data[index + 1] = (p.y + height / 2) / height;
            data[index + 2] = p.strength;
            data[index + 3] = 0.0;
        });

        const texture = new THREE.DataTexture(data, points.length, 1, THREE.RGBAFormat, THREE.FloatType);
        texture.needsUpdate = true;
        return texture;
    }, [points, width, height]);

    useEffect(() => {
        uniformsRef.current.u_radius.value = RADIUS;
        uniformsRef.current.u_opacity.value = OPACITY;
        uniformsRef.current.u_debugMode.value = debugModeMap[debugMode];
        uniformsRef.current.u_growthRate.value = GROWTH_TIME_MULTIPLIER;
    }, [RADIUS, OPACITY, debugMode, GROWTH_TIME_MULTIPLIER]);

    useEffect(() => {
        uniformsRef.current.u_points.value = pointTexture;
        uniformsRef.current.u_count.value = points.length;
    }, [pointTexture, points.length]);

    return (
        <mesh rotation={[Math.PI / 2, 0, 0]} position={[0, 0.025, 0]}>
            <planeGeometry args={[width, height]} />
            <shaderMaterial
                ref={materialRef}
                transparent
                depthWrite={false}
                blending={THREE.AdditiveBlending}
                uniforms={uniformsRef.current}
                side={THREE.DoubleSide}
                vertexShader={`
                    varying vec2 vUv;
                    void main() {
                        vUv = uv;
                        gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
                    }
                `}
                fragmentShader={`
                    precision highp float;
                    uniform sampler2D u_points;
                    uniform int u_count;
                    uniform float u_radius;
                    uniform float u_opacity;
                    uniform int u_debugMode;
                    uniform float u_growthRate;
                    varying vec2 vUv;

                    float gauss(float dist, float radius) {
                        return exp(-pow(dist / radius, 2.0));
                    }

                    void main() {
                        if (u_debugMode == 0) {
                            gl_FragColor = vec4(1.0, 0.0, 0.0, 1.0);
                            return;
                        }

                        float intensity = 0.0;

                        for (int i = 0; i < 10000; i++) {
                            if (i >= u_count) break;
                            float fi = float(i) + 0.5;
                            float u = fi / float(u_count);

                            vec4 point = texture2D(u_points, vec2(u, 0.5));
                            vec2 pos = point.rg;
                            float strength = point.b;

                            float d = distance(vUv, pos);
                            intensity += strength * gauss(d, u_radius);
                        }

                        if (u_debugMode == 1) {
                            gl_FragColor = vec4(vec3(intensity), 1.0);
                            return;
                        }

                        // Apply growth rate BEFORE normalization
                        float adjustedIntensity = intensity * u_growthRate;

                        // Normalize intensity between 0-1
                         float normalized = clamp(intensity / max(u_growthRate, 0.0001), 0.0, 1.0);

                        vec3 color = vec3(0.0);
                        if (normalized < 0.33) {
                            color = mix(vec3(0.0, 0.0, 1.0), vec3(0.0, 1.0, 0.0), normalized / 0.33);
                        } else if (normalized < 0.66) {
                            color = mix(vec3(0.0, 1.0, 0.0), vec3(1.0, 1.0, 0.0), (normalized - 0.33) / 0.33);
                        } else {
                            color = mix(vec3(1.0, 1.0, 0.0), vec3(1.0, 0.0, 0.0), (normalized - 0.66) / 0.34);
                        }

                        gl_FragColor = vec4(color, normalized * u_opacity);
                    }

                `}
            />
        </mesh>
    );
};

export default RealTimeHeatMap;
