import React, { useRef, useMemo, useEffect } from 'react';
import { useFrame } from '@react-three/fiber';
import * as THREE from 'three';

// -----------------------------------------------------------------------------
// SHADERS
// -----------------------------------------------------------------------------

const vertexShader = `
  uniform float uTime;
  uniform float uScroll;
  uniform float uMorph; // 0.0 = Sphere, 1.0 = Triangle
  uniform float uExplode; // 0.0 = Normal, 1.0 = Exploded/Work Mode
  uniform float uShiftX; 
  uniform float uShiftY; 
  
  attribute float aRandom;
  attribute vec3 aOriginalPos;
  attribute vec3 aTrianglePos;

  varying float vDepth;
  varying float vRim;
  varying vec3 vPos;
  varying float vExplodeAlpha;

  // Simplex 3D Noise 
  vec3 mod289(vec3 x) { return x - floor(x * (1.0 / 289.0)) * 289.0; }
  vec4 mod289(vec4 x) { return x - floor(x * (1.0 / 289.0)) * 289.0; }
  vec4 permute(vec4 x) { return mod289(((x * 34.0) + 1.0) * x); }
  vec4 taylorInvSqrt(vec4 r) { return 1.79284291400159 - 0.85373472095314 * r; }

  float snoise(vec3 v) {
    const vec2 C = vec2(1.0/6.0, 1.0/3.0);
    const vec4 D = vec4(0.0, 0.5, 1.0, 2.0);

    // First corner
    vec3 i  = floor(v + dot(v, C.yyy));
    vec3 x0 = v - i + dot(i, C.xxx);

    // Other corners
    vec3 g = step(x0.yzx, x0.xyz);
    vec3 l = 1.0 - g;
    vec3 i1 = min( g.xyz, l.zxy );
    vec3 i2 = max( g.xyz, l.zxy );

    vec3 x1 = x0 - i1 + C.xxx;
    vec3 x2 = x0 - i2 + C.yyy; 
    vec3 x3 = x0 - D.yyy;      

    // Permutations
    i = mod289(i);
    vec4 p = permute( permute( permute(
              i.z + vec4(0.0, i1.z, i2.z, 1.0 ))
            + i.y + vec4(0.0, i1.y, i2.y, 1.0 ))
            + i.x + vec4(0.0, i1.x, i2.x, 1.0 ));

    // Gradients
    float n_ = 0.142857142857; 
    vec3  ns = n_ * D.wyz - D.xzx;

    vec4 j = p - 49.0 * floor(p * ns.z * ns.z); 

    vec4 x_ = floor(j * ns.z);
    vec4 y_ = floor(j - 7.0 * x_ );   

    vec4 x = x_ *ns.x + ns.yyyy;
    vec4 y = y_ *ns.x + ns.yyyy;
    vec4 h = 1.0 - abs(x) - abs(y);

    vec4 b0 = vec4( x.xy, y.xy );
    vec4 b1 = vec4( x.zw, y.zw );

    vec4 s0 = floor(b0)*2.0 + 1.0;
    vec4 s1 = floor(b1)*2.0 + 1.0;
    vec4 sh = -step(h, vec4(0.0));

    vec4 a0 = b0.xzyw + s0.xzyw*sh.xxyy ;
    vec4 a1 = b1.xzyw + s1.xzyw*sh.zzww ;

    vec3 p0 = vec3(a0.xy,h.x);
    vec3 p1 = vec3(a0.zw,h.y);
    vec3 p2 = vec3(a1.xy,h.z);
    vec3 p3 = vec3(a1.zw,h.w);

    vec4 norm = taylorInvSqrt(vec4(dot(p0,p0), dot(p1,p1), dot(p2, p2), dot(p3,p3)));
    p0 *= norm.x;
    p1 *= norm.y;
    p2 *= norm.z;
    p3 *= norm.w;

    vec4 m = max(0.6 - vec4(dot(x0,x0), dot(x1,x1), dot(x2,x2), dot(x3,x3)), 0.0);
    m = m * m;
    return 42.0 * dot( m*m, vec4( dot(p0,x0), dot(p1,x1),
                                  dot(p2,x2), dot(p3,x3) ) );
  }

  void main() {
    // 0. Rotation Logic
    // Slow down rotation significantly when exploded to stabilize the backdrop
    float rotSpeed = mix(0.02, 0.002, uExplode);
    float c = cos(uTime * rotSpeed);
    float s = sin(uTime * rotSpeed);
    mat3 rotate = mat3(c, 0.0, -s, 0.0, 1.0, 0.0, s, 0.0, c);
    
    vec3 spherePos = rotate * aOriginalPos;
    vec3 triPos = aTrianglePos;

    // 1. MORPH INTERPOLATION
    vec3 pos = mix(spherePos, triPos, uMorph);

    // 2. Noise Field (Internal Motion)
    float noiseScale = 0.8;
    float timeScale = 0.15;
    
    float noise1 = snoise(pos * noiseScale + vec3(uTime * timeScale));
    
    // Dampen noise when exploded so it's less chaotic background
    float displacement = noise1 * 0.6 * (1.0 - uExplode * 0.5);

    // 3. Apply Displacement
    vec3 normal = normalize(pos);
    float dispAmount = mix(0.5, 0.15, uMorph);
    vec3 newPos = pos + normal * displacement * dispAmount;

    // 4. Shift Logic (World Space)
    // When exploded, we remove lateral shifts so the tunnel is centered
    float effectiveShiftX = mix(uShiftX, 0.0, uExplode);
    float effectiveShiftY = mix(uShiftY, 0.0, uExplode);
    
    newPos.x += effectiveShiftX;
    newPos.y += effectiveShiftY;

    // 5. EXPLOSION / WIPE LOGIC (Enhanced)
    // Instead of just pushing away, we create a "Tunnel" or "Stargate" effect.
    // We push X and Y radially outward based on the uExplode factor, clearing the center.
    
    float radius2D = length(newPos.xy);
    
    // Radial push: The closer to the center, the harder we push out, ensuring a clear text area
    float pushFactor = smoothstep(0.0, 1.0, uExplode) * 20.0;
    
    // Direction from center
    vec2 dir2D = normalize(newPos.xy);
    
    // Apply Push
    // We add a nonlinear expansion so the middle clears fast, but the edges stay visible longer
    vec3 explodedPos = newPos;
    explodedPos.xy += dir2D * pushFactor * (1.0 + aRandom); // Add random to break uniformity
    explodedPos.z -= uExplode * 10.0; // Push deep into background

    newPos = mix(newPos, explodedPos, uExplode);

    // 6. Final Position
    vec4 mvPosition = modelViewMatrix * vec4(newPos, 1.0);
    gl_Position = projectionMatrix * mvPosition;
    
    // 7. Point Size
    // Make them smaller and sharper when they are background stars
    float sizeBase = mix(2.0, 1.5, uExplode);
    gl_PointSize = sizeBase * (15.0 / -mvPosition.z);

    // 8. Varyings
    vPos = newPos;
    
    // Alpha Logic for Explosion:
    // Fade out particles that are still too close to center-screen after explosion
    // This guarantees text legibility
    float screenCenterDist = length(newPos.xy);
    float centerClearMask = smoothstep(2.0, 8.0, screenCenterDist); // 0 at center, 1 at edges
    vExplodeAlpha = mix(1.0, centerClearMask, uExplode);
    
    vec3 viewDir = normalize(-mvPosition.xyz);
    vec3 viewNormal = normalize(normalMatrix * normal);
    float dotNV = dot(viewDir, viewNormal);
    vRim = 1.0 - max(0.0, abs(dotNV));
    vRim = pow(vRim, 2.5);

    vDepth = smoothstep(-2.0, 5.0, newPos.z); 
  }
`;

const fragmentShader = `
  varying float vDepth;
  varying float vRim;
  varying vec3 vPos;
  varying float vExplodeAlpha;

  uniform vec3 uColorCore;
  uniform vec3 uColorMid;

  void main() {
    vec2 coord = gl_PointCoord - vec2(0.5);
    float dist = length(coord);
    if (dist > 0.5) discard;

    float alpha = 1.0 - smoothstep(0.3, 0.5, dist);

    // Colors
    vec3 cCore = uColorCore;       
    vec3 cMid  = uColorMid;        
    vec3 cRim  = vec3(0.8, 0.1, 0.0); 

    vec3 finalColor;
    float midMix = smoothstep(0.0, 0.6, vRim);
    finalColor = mix(cCore, cMid, midMix);
    
    float rimMix = smoothstep(0.6, 1.0, vRim);
    finalColor = mix(finalColor, cRim, rimMix);
    
    float intensity = 1.0 + (midMix * 0.5); 
    
    // Adjust depth alpha calculation to be more forgiving in exploded state
    float depthAlpha = smoothstep(-10.0, 5.0, vPos.z) * 0.9 + 0.1;
    
    // Combine standard alpha with explosion alpha (hollows out the center)
    float finalAlpha = alpha * depthAlpha * vExplodeAlpha;

    gl_FragColor = vec4(finalColor * intensity, finalAlpha);
  }
`;

// -----------------------------------------------------------------------------
// COMPONENT
// -----------------------------------------------------------------------------

interface GenerativeSphereProps {
  scrollRef?: React.MutableRefObject<number>;
  shiftRef?: React.MutableRefObject<number>;
  mode: 'sphere' | 'triangle' | 'explode';
  isMobile: boolean;
}

export const GenerativeSphere: React.FC<GenerativeSphereProps> = ({ scrollRef, shiftRef, mode, isMobile }) => {
  const pointsRef = useRef<THREE.Points>(null);
  const materialRef = useRef<THREE.ShaderMaterial>(null);
  
  // Smoothing refs
  const smoothedScroll = useRef(0);
  const currentMorph = useRef(0); 
  const currentExplode = useRef(0);
  const currentShiftX = useRef(0);
  const currentShiftY = useRef(0);

  const COUNT = 32000; 
  // Adjusted radius for mobile to be even smaller
  const RADIUS = isMobile ? 1.4 : 3.5;

  const { positions, originalPositions, trianglePositions, randoms } = useMemo(() => {
    const pos = new Float32Array(COUNT * 3);
    const origPos = new Float32Array(COUNT * 3);
    const triPos = new Float32Array(COUNT * 3);
    const rnd = new Float32Array(COUNT);
    
    const phi = Math.PI * (3 - Math.sqrt(5)); 
    const triangleScale = RADIUS * 0.6; 

    for (let i = 0; i < COUNT; i++) {
      // 1. SPHERE
      const y = 1 - (i / (COUNT - 1)) * 2; 
      const radiusAtY = Math.sqrt(1 - y * y); 
      const theta = phi * i; 

      const x = Math.cos(theta) * radiusAtY;
      const z = Math.sin(theta) * radiusAtY;

      const sx = x * RADIUS;
      const sy = y * RADIUS;
      const sz = z * RADIUS;

      pos[i * 3] = sx;
      pos[i * 3 + 1] = sy;
      pos[i * 3 + 2] = sz;

      origPos[i * 3] = sx;
      origPos[i * 3 + 1] = sy;
      origPos[i * 3 + 2] = sz;

      // 2. TRIANGLE
      const angle = Math.atan2(sy, sx);
      const distXY = Math.sqrt(sx*sx + sy*sy);

      const segmentAngle = (2 * Math.PI) / 3;
      const offsetAngle = angle + Math.PI / 2;
      const constrainedAngle = offsetAngle - segmentAngle * Math.floor((offsetAngle + segmentAngle / 2) / segmentAngle);
      
      const r_sharp = 1.0 / Math.cos(constrainedAngle);
      
      const r_factor = r_sharp; 
      
      const normalizedRadialDist = distXY / RADIUS; 
      const finalR = normalizedRadialDist * r_factor * triangleScale;
      
      const tx = Math.cos(angle) * finalR;
      const ty = Math.sin(angle) * finalR;
      const tz = sz * 0.35;

      triPos[i * 3] = tx;
      triPos[i * 3 + 1] = ty;
      triPos[i * 3 + 2] = tz;

      rnd[i] = Math.random();
    }
    
    return {
      positions: pos,
      originalPositions: origPos,
      trianglePositions: triPos,
      randoms: rnd
    };
  }, [RADIUS]);

  const uniforms = useMemo(() => ({
    uTime: { value: 0 },
    uScroll: { value: 0 },
    uMorph: { value: 0 },
    uExplode: { value: 0 },
    uShiftX: { value: 0 },
    uShiftY: { value: 0 },
    uColorCore: { value: new THREE.Color('#ff5522') }, 
    uColorMid: { value: new THREE.Color('#ffddaa') }, 
  }), []);

  useFrame((state, delta) => {
    if (materialRef.current) {
      materialRef.current.uniforms.uTime.value = state.clock.getElapsedTime();
      
      // Scroll Physics
      if (scrollRef) {
        smoothedScroll.current = THREE.MathUtils.lerp(smoothedScroll.current, scrollRef.current, 0.05);
        materialRef.current.uniforms.uScroll.value = smoothedScroll.current;
      }

      // Physics State Logic
      const isTriangle = mode === 'triangle';
      const isExplode = mode === 'explode';

      const targetMorph = isTriangle ? 1.0 : 0.0;
      const targetExplode = isExplode ? 1.0 : 0.0;

      // Interolate values
      currentMorph.current = THREE.MathUtils.lerp(currentMorph.current, targetMorph, 2.0 * delta);
      // Explode needs to be slightly snappy but smooth
      currentExplode.current = THREE.MathUtils.lerp(currentExplode.current, targetExplode, 3.0 * delta); 

      materialRef.current.uniforms.uMorph.value = currentMorph.current;
      materialRef.current.uniforms.uExplode.value = currentExplode.current;

      // Position Physics (Shift)
      let targetX = 0;
      let targetY = 0;

      if (isMobile) {
        targetX = 0; 
        targetY = 0.5; // Adjusted Y shift for smaller sphere on mobile
      } else {
        if (shiftRef && !isExplode) {
           // Only shift if NOT exploded. If exploded, we want center screen.
          targetX = shiftRef.current;
        }
        targetY = 0;
      }

      currentShiftX.current = THREE.MathUtils.lerp(currentShiftX.current, targetX, 2.0 * delta);
      currentShiftY.current = THREE.MathUtils.lerp(currentShiftY.current, targetY, 2.0 * delta);

      materialRef.current.uniforms.uShiftX.value = currentShiftX.current;
      materialRef.current.uniforms.uShiftY.value = currentShiftY.current;
    }
  });

  return (
    <points ref={pointsRef}>
      <bufferGeometry>
        <bufferAttribute
          attach="attributes-position"
          count={positions.length / 3}
          array={positions}
          itemSize={3}
        />
        <bufferAttribute
          attach="attributes-aOriginalPos"
          count={originalPositions.length / 3}
          array={originalPositions}
          itemSize={3}
        />
         <bufferAttribute
          attach="attributes-aTrianglePos"
          count={trianglePositions.length / 3}
          array={trianglePositions}
          itemSize={3}
        />
        <bufferAttribute
          attach="attributes-aRandom"
          count={randoms.length}
          array={randoms}
          itemSize={1}
        />
      </bufferGeometry>
      <shaderMaterial
        ref={materialRef}
        vertexShader={vertexShader}
        fragmentShader={fragmentShader}
        uniforms={uniforms}
        transparent={true}
        depthWrite={false} 
        blending={THREE.AdditiveBlending}
      />
    </points>
  );
};