import * as THREE from 'three'

const vertexShader = `
varying vec2 vUv;
varying vec4 vParticleColor;

void main() {
    vUv = uv;
    vParticleColor = vec4(1.0);
    gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
}
`

const fragmentShader = `
uniform float uTime;
uniform vec3 uColour;
uniform float uErosion;
uniform float uFlowSpeed;
uniform sampler2D uFoamTex;

uniform vec2 uCircleCenter;
uniform float uCircleRadius;
uniform float uCircleSoftness;
uniform float uCircleEnabled;

varying vec2 vUv;
varying vec4 vParticleColor;

float radialGradientExponential(vec2 uv, vec2 center, float radius, float density) {
    float dist = length(uv - center);
    dist = clamp(dist / radius, 0.0, 1.0);
    return pow(1.0 - dist, density);
}

vec2 panner(vec2 uv, float speedX, float speedY, float time) {
    return uv + vec2(speedX, speedY) * time;
}

float circleMask(vec2 uv, vec2 center, float radius, float softness) {
    float dist = length(uv - center);
    return 1.0 - smoothstep(radius - softness, radius + softness, dist);
}

void main() {
    vec2 animatedUV = panner(vUv, uFlowSpeed, 0.0, uTime);

    vec4 foamTex = texture2D(uFoamTex, animatedUV);

    float gradient = radialGradientExponential(vUv, vec2(0.5), 1.0, 1.0);

    float addResult = vParticleColor.r * foamTex.r + vParticleColor.g * foamTex.g;
    float subtractResult = addResult - uErosion;
    float multiplyResult = subtractResult * gradient;

    float opacityShape = mix(multiplyResult, addResult, gradient);
    opacityShape = clamp(opacityShape, 0.0, 1.0);
    opacityShape = pow(opacityShape, 2.0);

    float circle = circleMask(vUv, uCircleCenter, uCircleRadius, uCircleSoftness);
    opacityShape *= mix(1.0, circle, uCircleEnabled);

    vec3 emissive = uColour * vParticleColor.rgb;

    gl_FragColor = vec4(emissive, opacityShape);
}
`

interface WaterFoamMaterialOptions {
  colour?: THREE.Color
  erosion?: number
  flowSpeed?: number
  foamTexture?: THREE.Texture
  circleCenter?: THREE.Vector2
  circleRadius?: number
  circleSoftness?: number
  circleEnabled?: boolean
}

function createWaterFoamMaterial(options: WaterFoamMaterialOptions = {}) {
  const {
    colour = new THREE.Color(0xffffff),
    erosion = 1.0,
    flowSpeed = 0.2,
    foamTexture,
    circleCenter = new THREE.Vector2(0.5, 0.5),
    circleRadius = 0.34,
    circleSoftness = 0.195,
    circleEnabled = true,
  } = options

  return new THREE.ShaderMaterial({
    vertexShader,
    fragmentShader,
    transparent: true,
    depthWrite: false,
    uniforms: {
      uTime: { value: 0 },
      uColour: { value: colour },
      uErosion: { value: erosion },
      uFlowSpeed: { value: flowSpeed },
      uFoamTex: { value: foamTexture || new THREE.Texture() },
      uCircleCenter: { value: circleCenter },
      uCircleRadius: { value: circleRadius },
      uCircleSoftness: { value: circleSoftness },
      uCircleEnabled: { value: circleEnabled ? 1.0 : 0.0 },
    },
  })
}

export { createWaterFoamMaterial }
export type { WaterFoamMaterialOptions }