import React, { Suspense, useRef, useMemo, useEffect } from "react";
import { useGLTF } from "@react-three/drei";
import * as THREE from "three";
import { useStore, type Model3D as Model3DType } from "../store/store.ts";
import {
  hueShiftVertexShader,
  hueShiftFragmentShader,
} from "../shaders/hueShift.ts";

interface Model3DProps {
  model: Model3DType;
}

interface ShaderUniforms {
  [uniform: string]: THREE.IUniform;
  map: { value: THREE.Texture | null };
  hasMap: { value: boolean };
  baseColor: { value: THREE.Vector3 };
  hueShiftRed: { value: number };
  hueShiftGreen: { value: number };
  hueShiftBlue: { value: number };
}

interface MeshWithUniforms extends THREE.Mesh {
  userData: {
    uniforms?: ShaderUniforms;
  };
}

const Model3DInner: React.FC<Model3DProps> = ({ model }) => {
  const groupRef = useRef<THREE.Group>(null);
  const meshRefsRef = useRef<MeshWithUniforms[]>([]);
  const { scene } = useGLTF(model.url);
  const selectModel = useStore((state) => state.selectModel);
  const selectedModelId = useStore((state) => state.selectedModelId);

  const isSelected = selectedModelId === model.id;

  const { clonedScene, boundingBox, meshes } = useMemo(() => {
    const clone = scene.clone(true);
    const collectedMeshes: MeshWithUniforms[] = [];

    clone.traverse((child) => {
      if (child instanceof THREE.Mesh) {
        const originalMaterial = child.material;

        let textureMap: THREE.Texture | null = null;
        let baseColor = new THREE.Vector3(0.8, 0.8, 0.8);

        if (originalMaterial instanceof THREE.MeshStandardMaterial) {
          textureMap = originalMaterial.map;
          if (originalMaterial.color) {
            baseColor = new THREE.Vector3(
              originalMaterial.color.r,
              originalMaterial.color.g,
              originalMaterial.color.b,
            );
          }
        } else if (originalMaterial instanceof THREE.MeshBasicMaterial) {
          textureMap = originalMaterial.map;
          if (originalMaterial.color) {
            baseColor = new THREE.Vector3(
              originalMaterial.color.r,
              originalMaterial.color.g,
              originalMaterial.color.b,
            );
          }
        }

        const uniforms: ShaderUniforms = {
          map: { value: textureMap },
          hasMap: { value: !!textureMap },
          baseColor: { value: baseColor },
          hueShiftRed: { value: 0 },
          hueShiftGreen: { value: 0 },
          hueShiftBlue: { value: 0 },
        };

        const shaderMaterial = new THREE.ShaderMaterial({
          uniforms,
          vertexShader: hueShiftVertexShader,
          fragmentShader: hueShiftFragmentShader,
        });

        child.material = shaderMaterial;
        (child as MeshWithUniforms).userData.uniforms = uniforms;
        collectedMeshes.push(child as MeshWithUniforms);
      }
    });

    const box = new THREE.Box3().setFromObject(clone);
    const size = new THREE.Vector3();
    const center = new THREE.Vector3();
    box.getSize(size);
    box.getCenter(center);

    return {
      clonedScene: clone,
      boundingBox: { size, center },
      meshes: collectedMeshes,
    };
  }, [scene]);

  useEffect(() => {
    meshRefsRef.current = meshes;
  }, [meshes]);

  useEffect(() => {
    for (const mesh of meshRefsRef.current) {
      const uniforms = mesh.userData.uniforms;
      if (uniforms) {
        uniforms.hueShiftRed.value = model.hueShift.red;
        uniforms.hueShiftGreen.value = model.hueShift.green;
        uniforms.hueShiftBlue.value = model.hueShift.blue;
      }
    }
  }, [model.hueShift]);

  const handleClick = (event: { stopPropagation: () => void }) => {
    event.stopPropagation();
    selectModel(model.id);
  };

  const selectionBoxGeometry = useMemo(() => {
    if (!boundingBox) return null;
    return new THREE.BoxGeometry(
      boundingBox.size.x * 1.1,
      boundingBox.size.y * 1.1,
      boundingBox.size.z * 1.1,
    );
  }, [boundingBox]);

  return (
    <group
      ref={groupRef}
      position={[model.position.x, model.position.y, model.position.z]}
      rotation={[model.rotation.x, model.rotation.y, model.rotation.z]}
      onClick={handleClick}
    >
      <primitive object={clonedScene} />
      {isSelected && boundingBox && selectionBoxGeometry && (
        <lineSegments
          position={[
            boundingBox.center.x,
            boundingBox.center.y,
            boundingBox.center.z,
          ]}
        >
          <edgesGeometry args={[selectionBoxGeometry]} />
          <lineBasicMaterial color="#00ff00" linewidth={2} />
        </lineSegments>
      )}
    </group>
  );
};

const MemoizedModel3DInner = React.memo(Model3DInner);
MemoizedModel3DInner.displayName = "Model3DInner";

export const Model3DComponent: React.FC<Model3DProps> = ({ model }) => {
  return (
    <Suspense fallback={null}>
      <MemoizedModel3DInner model={model} />
    </Suspense>
  );
};