import * as THREE from 'three';

// Enhanced ripple simulation with multiple trailing ripples and lighting

const FluidSimShader = {
  uniforms: {
    tPrev: { value: null },
    iResolution: { value: new THREE.Vector2() },
    iTime: { value: 0.0 },
    mouse: { value: new THREE.Vector3(-1, -1, 0.0) },
    dissipation: { value: 0.950 }, // Slightly more persistent for trails
    tension: { value: 2.2 }, // Higher tension for stronger ripples
    radius: { value: 20.0 }, // Larger splat radius
    trailLength: { value: 5 }, // Number of trailing ripples
  },

  vertexShader: `
    varying vec2 vUv;
    void main() {
      vUv = uv;
      gl_Position = vec4(position.xy, 0.0, 1.0);
    }
  `,

  fragmentShader: `
    precision highp float;
    varying vec2 vUv;
    uniform sampler2D tPrev;
    uniform vec2 iResolution;
    uniform float iTime;
    uniform vec3 mouse;
    uniform float dissipation;
    uniform float tension;
    uniform float radius;
    uniform float trailLength;

    vec2 readRG(vec2 uv) {
      vec4 c = texture2D(tPrev, uv);
      return c.rg;
    }

    void main() {
      vec2 texel = 1.0 / iResolution;
      vec2 currPrev = readRG(vUv);
      float curr = currPrev.r;
      float prev = currPrev.g;

      // Enhanced 8-neighbor laplacian for stronger ripples
      float up = readRG(vUv + vec2(0.0, texel.y)).r;
      float down = readRG(vUv + vec2(0.0, -texel.y)).r;
      float right = readRG(vUv + vec2( texel.x, 0.0)).r;
      float left = readRG(vUv + vec2(-texel.x, 0.0)).r;

      // Diagonal neighbors for smoother ripples
      float upLeft = readRG(vUv + vec2(-texel.x, texel.y)).r;
      float upRight = readRG(vUv + vec2( texel.x, texel.y)).r;
      float downLeft = readRG(vUv + vec2(-texel.x, -texel.y)).r;
      float downRight = readRG(vUv + vec2( texel.x, -texel.y)).r;

      // Enhanced laplacian with diagonal weights
      float lap = (up + down + left + right) * 0.2 +
                  (upLeft + upRight + downLeft + downRight) * 0.05 - curr;

      // Wave equation with enhanced parameters
      float next = curr + (curr - prev) * dissipation + lap * tension;

      // Multiple trailing ripples from mouse movement
      if (mouse.z > 0.0001) {
        vec2 uvPx = vUv * iResolution;
        vec2 d = uvPx - mouse.xy;
        float dist = length(d);

        // Create multiple concentric ripples
        for (float i = 0.0; i < 4.0; i++) {
          if (i >= trailLength) break;
          
          float offset = i * radius * 0.4; // Spacing between ripples
          float r = radius + offset;
          float timeOffset = i * 0.3; // Temporal offset for trailing effect
          
          // Gaussian with sine wave for ripple pattern
          float g = exp(-pow(dist - offset, 2.0) / (r * r * 0.5));
          float ripple = sin(dist * 0.2 - iTime * 8.0 + timeOffset) * g;
          
          // Diminishing strength for trailing ripples
          float strength = mouse.z * (1.0 - i * 0.2) * 0.8;
          next += ripple * strength;
        }
      }

      gl_FragColor = vec4(next, curr, 0.0, 1.0);
    }
  `
};

// Enhanced distortion shader with dynamic lighting and ripple whiteness
export const FluidDistortionShader = {
  uniforms: {
    tDiffuse: { value: null },
    tSim: { value: null },
    iResolution: { value: new THREE.Vector2() },
    amount: { value: 0.12 }, // Stronger base distortion
    chromaticAmount: { value: 0.015 }, // Enhanced chromatic aberration
    lightPosition: { value: new THREE.Vector3(0.5, 0.5, 1.0) }, // Light position
    lightIntensity: { value: 1.5 }, // Light brightness
    lightColor: { value: new THREE.Color(0.8, 0.9, 1.0) }, // Cool light color
    normalStrength: { value: 2.0 }, // How pronounced the lighting effect is
    ambientLight: { value: 0.15 }, // Base ambient lighting
    rippleWhiteness: { value: 0.15 }, // Amount of white tint for ripples
    rippleBrightness: { value: 1.8 }, // Brightness multiplier for ripple areas
  },

  vertexShader: `
    varying vec2 vUv;
    void main() {
      vUv = uv;
      gl_Position = vec4(position.xy, 0.0, 1.0);
    }
  `,

  fragmentShader: `
    precision highp float;
    varying vec2 vUv;
    uniform sampler2D tDiffuse;
    uniform sampler2D tSim;
    uniform vec2 iResolution;
    uniform float amount;
    uniform float chromaticAmount;
    uniform vec3 lightPosition;
    uniform float lightIntensity;
    uniform vec3 lightColor;
    uniform float normalStrength;
    uniform float ambientLight;
    uniform float rippleWhiteness;
    uniform float rippleBrightness;

    void main() {
      vec2 texel = 1.0 / iResolution;

      // Sample height field for normal calculation
      float hC = texture2D(tSim, vUv).r;
      float hL = texture2D(tSim, vUv - vec2(texel.x, 0.0)).r;
      float hR = texture2D(tSim, vUv + vec2(texel.x, 0.0)).r;
      float hD = texture2D(tSim, vUv - vec2(0.0, texel.y)).r;
      float hU = texture2D(tSim, vUv + vec2(0.0, texel.y)).r;

      // Calculate gradient and normal
      vec2 grad = vec2(hR - hL, hU - hD) * normalStrength;
      vec3 normal = normalize(vec3(-grad.x, -grad.y, 1.0));

      // Enhanced distortion with trailing effect
      vec2 baseOffset = grad * amount;
      
      // Add subtle trailing distortion based on height
      vec2 trailOffset = grad * abs(hC) * amount * 0.3;
      vec2 totalOffset = baseOffset + trailOffset;

      // Chromatic aberration with enhanced separation
      vec2 chromaticOffset = grad * chromaticAmount;
      vec2 uvR = vUv + totalOffset + chromaticOffset;
      vec2 uvG = vUv + totalOffset;
      vec2 uvB = vUv + totalOffset - chromaticOffset;

      // Clamp UVs
      uvR = clamp(uvR, vec2(0.0), vec2(1.0));
      uvG = clamp(uvG, vec2(0.0), vec2(1.0));
      uvB = clamp(uvB, vec2(0.0), vec2(1.0));

      // Sample distorted colors
      float r = texture2D(tDiffuse, uvR).r;
      float g = texture2D(tDiffuse, uvG).g;
      float b = texture2D(tDiffuse, uvB).b;
      vec3 distortedColor = vec3(r, g, b);

      // Dynamic lighting calculation
      vec3 lightDir = normalize(vec3(lightPosition.xy - vUv, lightPosition.z));
      float NdotL = max(dot(normal, lightDir), 0.0);

      // Create rim lighting effect for ripples
      float rimLight = pow(1.0 - abs(dot(normal, vec3(0.0, 0.0, 1.0))), 2.0);

      // Combine lighting effects
      vec3 lighting = lightColor * (NdotL * lightIntensity + rimLight * 0.3) + ambientLight;

      // Calculate ripple intensity for both lighting and whiteness
      float rippleIntensity = abs(hC) + length(grad) * 0.5;
      rippleIntensity = clamp(rippleIntensity, 0.0, 1.0);

      // Apply lighting selectively - stronger where there are ripples
      vec3 litColor = mix(distortedColor, distortedColor * lighting, rippleIntensity);

      // Add white tint to ripples for visibility over black areas
      vec3 whiteColor = vec3(1.0, 1.0, 1.0);
      
      // Create a smooth falloff for the whiteness effect
      float whiteIntensity = smoothstep(0.0, 0.3, rippleIntensity) * rippleWhiteness;
      
      // Blend in the white tint
      vec3 rippleColor = mix(litColor, whiteColor, whiteIntensity);
      
      // Brighten areas with ripples
      rippleColor = mix(rippleColor, rippleColor * rippleBrightness, rippleIntensity * 0.5);

      gl_FragColor = vec4(rippleColor, 1.0);
    }
  `
};

// Enhanced fluid simulation factory
export function createFluidSimulation(renderer, dpr = 1) {
  const simScene = new THREE.Scene();
  const simCamera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0, 1);
  
  const quad = new THREE.Mesh(
    new THREE.PlaneGeometry(2, 2),
    new THREE.ShaderMaterial({
      uniforms: THREE.UniformsUtils.clone(FluidSimShader.uniforms),
      vertexShader: FluidSimShader.vertexShader,
      fragmentShader: FluidSimShader.fragmentShader,
      depthTest: false,
      depthWrite: false
    })
  );
  
  simScene.add(quad);

  // Higher precision for better ripple quality
  const params = {
    minFilter: THREE.LinearFilter,
    magFilter: THREE.LinearFilter,
    format: THREE.RGBAFormat,
    type: THREE.FloatType, // Use float for better precision
    depthBuffer: false,
    stencilBuffer: false
  };

  let width = Math.max(2, Math.floor(window.innerWidth * dpr));
  let height = Math.max(2, Math.floor(window.innerHeight * dpr));
  
  let rtA = new THREE.WebGLRenderTarget(width, height, params);
  let rtB = new THREE.WebGLRenderTarget(width, height, params);

  // Initialize
  renderer.setRenderTarget(rtA);
  renderer.clear();
  renderer.setRenderTarget(rtB);
  renderer.clear();
  renderer.setRenderTarget(null);

  quad.material.uniforms.iResolution.value.set(width, height);
  quad.material.uniforms.tPrev.value = rtA.texture;

  function swap() {
    const tmp = rtA; rtA = rtB; rtB = tmp;
  }

  function update(mouseX, mouseY, strength, timeSec) {
    quad.material.uniforms.iTime.value = timeSec;
    
    if (mouseX < 0.0 || mouseY < 0.0) {
      quad.material.uniforms.mouse.value.set(-1, -1, 0.0);
    } else {
      // Enhanced strength for better trailing effect
      const enhancedStrength = Math.max(0.0, Math.min(1.0, strength * 1.5));
      quad.material.uniforms.mouse.value.set(mouseX, mouseY, enhancedStrength);
    }

    quad.material.uniforms.tPrev.value = rtA.texture;
    renderer.setRenderTarget(rtB);
    renderer.render(simScene, simCamera);
    renderer.setRenderTarget(null);
    
    swap();
  }

  function getTexture() {
    return rtA.texture;
  }

  function resize(w, h, newDpr = dpr) {
    width = Math.max(2, Math.floor(w * newDpr));
    height = Math.max(2, Math.floor(h * newDpr));
    rtA.setSize(width, height);
    rtB.setSize(width, height);
    quad.material.uniforms.iResolution.value.set(width, height);
  }

  return { update, getTexture, resize };
}
