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: 1.5 },  // Higher tension for stronger ripples
    radius:      { value: 10.0 },  // Larger splat radius
    trailLength: { value: 3 },     // 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
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
  },
  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;

    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;
      
      // Apply lighting selectively - stronger where there are ripples
      float rippleIntensity = abs(hC) + length(grad) * 0.5;
      rippleIntensity = clamp(rippleIntensity, 0.0, 1.0);
      
      // Blend original color with lit color based on ripple presence
      vec3 finalColor = mix(distortedColor, distortedColor * lighting, rippleIntensity);
      
      gl_FragColor = vec4(finalColor, 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 };
}