import { useRef, useMemo } from "react";
import { useFrame, useThree, extend, ThreeElement } from "@react-three/fiber";
import * as THREE from "three";
import { shaderMaterial } from "@react-three/drei";

// Type declaration for JSX
declare global {
  namespace JSX {
    interface IntrinsicElements {
      liquidMaterialImpl: any;
    }
  }
}

const LiquidMaterialImpl = shaderMaterial(
  {
    uTime: 0,
    uColor: new THREE.Color("#ff9900"),
    uResolution: new THREE.Vector2(),
    uFresnelBias: 0.1,
    uFresnelScale: 1.0,
    uFresnelPower: 2.0,
    uRefractionRatio: 0.98,
    uIor: 1.33,
  },
  // Vertex Shader
  `
  #include <clipping_planes_pars_vertex>
  varying vec3 vNormal;
  varying vec3 vViewPosition;
  varying vec2 vUv;
  varying vec3 vWorldPosition;

  void main() {
    vUv = uv;
    vNormal = normalize(normalMatrix * normal);
    vec4 worldPosition = modelMatrix * vec4(position, 1.0);
    vWorldPosition = worldPosition.xyz;
    vec4 mvPosition = modelViewMatrix * vec4(position, 1.0);
    vViewPosition = -mvPosition.xyz;
    #include <clipping_planes_vertex>
    gl_Position = projectionMatrix * mvPosition;
  }
  `,
  // Fragment Shader
  `
  #include <clipping_planes_pars_fragment>
  
  uniform float uTime;
  uniform vec3 uColor;
  uniform vec2 uResolution;
  uniform float uFresnelBias;
  uniform float uFresnelScale;
  uniform float uFresnelPower;
  uniform float uIor;

  varying vec3 vNormal;
  varying vec3 vViewPosition;
  varying vec2 vUv;
  varying vec3 vWorldPosition;

  // Simple noise function for caustics and imperfections
  float hash(vec2 p) {
    return fract(sin(dot(p, vec2(127.1, 311.7))) * 43758.5453123);
  }

  float noise(vec2 p) {
    vec2 i = floor(p);
    vec2 f = fract(p);
    f = f * f * (3.0 - 2.0 * f);
    return mix(mix(hash(i + vec2(0.0, 0.0)), hash(i + vec2(1.0, 0.0)), f.x),
               mix(hash(i + vec2(0.0, 1.0)), hash(i + vec2(1.0, 1.0)), f.x), f.y);
  }

  void main() {
    #include <clipping_planes_fragment>
    vec3 normal = normalize(vNormal);
    vec3 viewDir = normalize(vViewPosition);

    // Fresnel effect
    float fresnel = uFresnelBias + uFresnelScale * pow(1.0 + dot(viewDir, normal), uFresnelPower);
    
    // Caustics simulation
    vec2 causticUv = vWorldPosition.xz * 2.0 + vWorldPosition.y * 0.5;
    float caustic = noise(causticUv * 5.0 + uTime * 0.5) * 0.5 + 0.5;
    caustic *= noise(causticUv * 10.0 - uTime * 0.3) * 0.5 + 0.5;
    caustic = pow(caustic, 3.0) * 1.5;

    // Surface imperfections
    float imperfection = noise(vUv * 20.0 + uTime * 0.1) * 0.05;
    
    vec3 finalColor = uColor;
    finalColor += caustic * 0.2 * uColor;
    finalColor -= imperfection;

    // Simple refraction-like look by blending with fresnel
    vec3 skyColor = vec3(0.8, 0.9, 1.0);
    finalColor = mix(finalColor, skyColor, fresnel * 0.3);

    gl_FragColor = vec4(finalColor, 0.9);
  }
  `
);

extend({ LiquidMaterialImpl });

export function Liquid({ color = "#ff9900", size = [0.9, 1.8, 0.5] }) {
  const materialRef = useRef<any>(null);
  const { size: winSize } = useThree();

  useFrame((state) => {
    if (materialRef.current) {
      const mat = materialRef.current;
      if (mat.uTime !== undefined) mat.uTime = state.clock.elapsedTime;
      if (mat.uResolution && mat.uResolution.set) {
        mat.uResolution.set(winSize.width || 800, winSize.height || 600);
      }
    }
  });

  return (
    <mesh>
      <boxGeometry args={size as any} />
      {/* @ts-ignore */}
      <liquidMaterialImpl 
        ref={materialRef} 
        uColor={new THREE.Color(color)} 
        transparent 
        side={THREE.DoubleSide}
        clipping={true}
      />
    </mesh>
  );
}