import { useRef, useEffect, useMemo } from 'react'
import { useFrame } from '@react-three/fiber'
import * as THREE from 'three'
import type { TransformedGraph, GraphNode3D } from './types'
import { NODE_TYPE_COLORS } from './constants'

const typeColorArrays = Object.values(NODE_TYPE_COLORS)

// Fresnel glow vertex shader — reads instanceColor from InstancedMesh
const vertexShader = /* glsl */ `
  varying vec3 vNormal;
  varying vec3 vViewDir;
  varying vec3 vColor;
  varying float vFlashSeed;

  void main() {
    #ifdef USE_INSTANCING
      vec4 worldPos = instanceMatrix * vec4(position, 1.0);
      vNormal = normalize(mat3(instanceMatrix) * normal);
      vColor = instanceColor;
      // Use instance position as unique seed for neural firing
      vFlashSeed = instanceMatrix[3][0] * 0.13 + instanceMatrix[3][1] * 0.17 + instanceMatrix[3][2] * 0.11;
    #else
      vec4 worldPos = vec4(position, 1.0);
      vNormal = normalize(normal);
      vColor = vec3(0.5);
      vFlashSeed = 0.0;
    #endif
    vec4 mvPos = modelViewMatrix * worldPos;
    vViewDir = normalize(-mvPos.xyz);
    gl_Position = projectionMatrix * mvPos;
  }
`

const fragmentShader = /* glsl */ `
  uniform float uTime;
  varying vec3 vNormal;
  varying vec3 vViewDir;
  varying vec3 vColor;
  varying float vFlashSeed;

  void main() {
    float fresnel = pow(1.0 - abs(dot(normalize(vNormal), normalize(vViewDir))), 2.5);
    vec3 core = vColor * 0.5;
    vec3 glow = vColor * fresnel * 1.4;
    vec3 color = core + glow;

    // Neural firing — nodes flash like they're receiving electrical signals
    float wave = sin(uTime * 1.8 + vFlashSeed * 6.2831);
    float flash = smoothstep(0.85, 1.0, wave) * 2.0;
    color += vColor * flash;

    float alpha = 0.7 + fresnel * 0.3 + flash * 0.2;
    gl_FragColor = vec4(color, min(alpha, 1.0));
  }
`

interface NodeCloudProps {
  graph: TransformedGraph
  positions: Float32Array
  highlightSet: Set<number> | null
  onHover: (index: number | null, node: GraphNode3D | null, position: [number, number, number] | null) => void
  onClick: (index: number | null) => void
}

export default function NodeCloud({ graph, positions, highlightSet, onHover, onClick }: NodeCloudProps) {
  const meshRef = useRef<THREE.InstancedMesh>(null!)
  const matRef = useRef<THREE.ShaderMaterial>(null!)
  const dummy = useMemo(() => new THREE.Object3D(), [])
  const sphereGeo = useMemo(() => new THREE.SphereGeometry(1, 16, 12), [])

  const uniforms = useMemo(() => ({
    uTime: { value: 0 },
  }), [])

  // Precompute degree-based brightness: more connections = brighter
  const maxDegree = useMemo(() => Math.max(1, ...graph.nodes.map(n => n.degree)), [graph])
  const degreeBrightness = useMemo(() =>
    graph.nodes.map(n => 0.4 + 0.6 * Math.pow(n.degree / maxDegree, 0.5)),
    [graph, maxDegree]
  )

  // Set instance colors from node types, scaled by degree brightness
  useEffect(() => {
    const mesh = meshRef.current
    if (!mesh) return
    const colors = new Float32Array(graph.nodes.length * 3)
    for (let i = 0; i < graph.nodes.length; i++) {
      const c = typeColorArrays[graph.nodes[i].colorIndex] || typeColorArrays[typeColorArrays.length - 1]
      const b = degreeBrightness[i]
      colors[i * 3] = c[0] * b
      colors[i * 3 + 1] = c[1] * b
      colors[i * 3 + 2] = c[2] * b
    }
    mesh.instanceColor = new THREE.InstancedBufferAttribute(colors, 3)
    mesh.instanceColor.needsUpdate = true
  }, [graph, degreeBrightness])

  // Update instance matrices from positions each frame
  useFrame((_, delta) => {
    const mesh = meshRef.current
    if (!mesh || !positions) return

    matRef.current.uniforms.uTime.value += delta

    for (let i = 0; i < graph.nodes.length; i++) {
      const inHighlight = highlightSet ? highlightSet.has(i) : false
      const degreeBoost = 0.8 + 0.4 * Math.pow(graph.nodes[i].degree / maxDegree, 0.5)
      const scale = graph.nodes[i].radius * degreeBoost * (highlightSet ? (inHighlight ? 1.4 : 0.6) : 1.0)
      dummy.position.set(positions[i * 3], positions[i * 3 + 1], positions[i * 3 + 2])
      dummy.scale.setScalar(scale)
      dummy.updateMatrix()
      mesh.setMatrixAt(i, dummy.matrix)
    }
    mesh.instanceMatrix.needsUpdate = true

    // Apply degree brightness + highlight dimming
    if (mesh.instanceColor) {
      for (let i = 0; i < graph.nodes.length; i++) {
        const c = typeColorArrays[graph.nodes[i].colorIndex] || typeColorArrays[typeColorArrays.length - 1]
        const b = degreeBrightness[i]
        const mult = highlightSet
          ? (highlightSet.has(i) ? Math.max(b, 1.0) * 1.4 : b * 0.2)
          : b
        ;(mesh.instanceColor as THREE.InstancedBufferAttribute).setXYZ(
          i, Math.min(c[0] * mult, 1), Math.min(c[1] * mult, 1), Math.min(c[2] * mult, 1)
        )
      }
      mesh.instanceColor.needsUpdate = true
    }
  })

  return (
    <instancedMesh
      ref={meshRef}
      args={[sphereGeo, undefined, graph.nodes.length]}
      frustumCulled={false}
      onPointerMove={(e) => {
        e.stopPropagation()
        if (e.instanceId != null) {
          const idx = e.instanceId
          onHover(idx, graph.nodes[idx], [positions[idx * 3], positions[idx * 3 + 1], positions[idx * 3 + 2]])
        }
      }}
      onPointerLeave={() => onHover(null, null, null)}
      onClick={(e) => {
        e.stopPropagation()
        onClick(e.instanceId ?? null)
      }}
    >
      <shaderMaterial
        ref={matRef}
        vertexShader={vertexShader}
        fragmentShader={fragmentShader}
        uniforms={uniforms}
        transparent
        depthWrite={false}
      />
    </instancedMesh>
  )
}
