import React, { useMemo } from 'react';
import { useFrame } from '@react-three/fiber';
import { SceneData, SceneNode } from './Generator';
import * as THREE from 'three';

const ObjectInstances: React.FC<{ data: SceneData }> = ({ data }) => {
    // We'll separate nodes by type to feed into distinct InstancedMeshes
    const windNodes = useMemo(() => data.nodes.filter(n => n.type === 'wind'), [data]);
    const solarNodes = useMemo(() => data.nodes.filter(n => n.type === 'solar'), [data]);
    const cityNodes = useMemo(() => data.nodes.filter(n => n.type === 'city_building'), [data]);
    const treeNodes = useMemo(() => data.nodes.filter(n => n.type === 'tree'), [data]);
    const towerNodes = useMemo(() => data.nodes.filter(n => n.type === 'tower'), [data]);
    const subNodes = useMemo(() => data.nodes.filter(n => n.type === 'substation'), [data]);

    const setMatrixAt = (mesh: THREE.InstancedMesh, i: number, node: SceneNode) => {
        const obj = new THREE.Object3D();
        obj.position.set(node.position[0], node.position[1], node.position[2]);
        obj.rotation.set(node.rotation[0], node.rotation[1], node.rotation[2]);
        obj.scale.set(node.scale[0], node.scale[1], node.scale[2]);
        obj.updateMatrix();
        mesh.setMatrixAt(i, obj.matrix);
    };

    // Precompute instance matrices
    const windMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), windNodes.length);
        windNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [windNodes]);

    const solarMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), solarNodes.length);
        solarNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [solarNodes]);

    const treeMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), treeNodes.length);
        treeNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [treeNodes]);

    const towerMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), towerNodes.length);
        towerNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [towerNodes]);

    const subMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), subNodes.length);
        subNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [subNodes]);

    const cityMatrices = useMemo(() => {
        const mesh = new THREE.InstancedMesh(new THREE.BoxGeometry(), new THREE.MeshBasicMaterial(), cityNodes.length);
        cityNodes.forEach((node, i) => setMatrixAt(mesh, i, node));
        mesh.instanceMatrix.needsUpdate = true;
        return mesh.instanceMatrix;
    }, [cityNodes]);

    return (
        <group>
            {/* Trees */}
            {treeNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, treeNodes.length]} castShadow receiveShadow>
                    <coneGeometry args={[1.5, 5, 4]} />
                    <meshStandardMaterial color="#0b2e1b" roughness={0.9} />
                    <primitive object={treeMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}

            {/* Wind Turbines Towers */}
            {windNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, windNodes.length]} castShadow receiveShadow>
                    <cylinderGeometry args={[0.3, 0.5, 12, 8]} />
                    <meshStandardMaterial color="#ffffff" roughness={0.5} />
                    <primitive object={windMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}

            {/* Solar Panels */}
            {solarNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, solarNodes.length]} castShadow receiveShadow>
                    <boxGeometry args={[4, 0.2, 2]} />
                    <meshStandardMaterial color="#0055ff" roughness={0.2} metalness={0.8} />
                    <primitive object={solarMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}

            {/* City Buildings */}
            {cityNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, cityNodes.length]} castShadow receiveShadow>
                    <boxGeometry args={[2, 1, 2]} />
                    <meshStandardMaterial color="#1a2b4c" roughness={0.7} emissive="#002b49" emissiveIntensity={0.2} />
                    <primitive object={cityMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}

            {/* Substations */}
            {subNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, subNodes.length]} castShadow receiveShadow>
                    <boxGeometry args={[6, 3, 6]} />
                    <meshStandardMaterial color="#4a5568" roughness={0.8} metalness={0.2} />
                    <primitive object={subMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}

            {/* Transition Towers (Grid) */}
            {towerNodes.length > 0 && (
                <instancedMesh args={[undefined, undefined, towerNodes.length]} castShadow receiveShadow>
                    <cylinderGeometry args={[0.5, 1.5, 12, 4]} />
                    <meshStandardMaterial color="#a0aec0" wireframe={true} roughness={0.8} metalness={0.6} />
                    <primitive object={towerMatrices} attach="instanceMatrix" />
                </instancedMesh>
            )}
        </group>
    );
};

export default ObjectInstances;