import { BufferAttribute, BufferGeometry, Color, DynamicDrawUsage, Group, LineBasicMaterial, LineSegments, MathUtils, Points, ShaderMaterial } from 'three';
import type { Vector3 } from 'three';
import { PRECIPITATION as P } from '../constants/precipitation.ts';

/** Bounded, reusable particles around the viewing camera, in world coordinates. */
export function createPrecipitation(heightAt: (x: number, z: number) => number) {
  const group = new Group();
  group.name = 'weather-precipitation';
  const rainMaterial = new LineBasicMaterial({ color: '#b9cddd', transparent: true, depthWrite: false });
  const snowMaterial = new ShaderMaterial({
    transparent: true, depthWrite: false,
    uniforms: { tint: { value: new Color() }, opacity: { value: P.snow.opacity }, pointScale: { value: 1 } },
    vertexShader: `
      uniform float pointScale;
      void main() {
        vec4 viewPosition = modelViewMatrix * vec4(position, 1.0);
        gl_Position = projectionMatrix * viewPosition;
        gl_PointSize = clamp(${P.snow.size} * pointScale * projectionMatrix[1][1] / max(0.1, -viewPosition.z), 1.0, 24.0);
      }`,
    fragmentShader: `
      uniform vec3 tint;
      uniform float opacity;
      void main() {
        float radius = length(gl_PointCoord - 0.5) * 2.0;
        float alpha = (1.0 - smoothstep(0.3, 1.0, radius)) * opacity;
        if (alpha < 0.01) discard;
        gl_FragColor = vec4(tint, alpha);
        #include <tonemapping_fragment>
        #include <colorspace_fragment>
      }`,
  });
  const layers = ([false, true] as const).map(snow => {
    const settings = snow ? P.snow : P.rain;
    const vertices = snow ? 1 : 2;
    const geometry = new BufferGeometry();
    const attribute = new BufferAttribute(new Float32Array(settings.maxParticles * vertices * 3), 3).setUsage(DynamicDrawUsage);
    geometry.setAttribute('position', attribute);
    geometry.setDrawRange(0, 0);
    const mesh = snow ? new Points(geometry, snowMaterial) : new LineSegments(geometry, rainMaterial);
    // The buffer moves with the camera; avoid recomputing its bounds every frame.
    mesh.frustumCulled = false;
    mesh.visible = false;
    group.add(mesh);
    return { snow, settings, vertices, mesh, attribute,
      positions: new Float32Array(settings.maxParticles * 3),
      ground: new Float32Array(settings.maxParticles),
      seeds: Float32Array.from({ length: settings.maxParticles }, () => Math.random()), count: 0 };
  });
  let elapsed = 0;
  const span = P.radius * 2;
  const wrap = (value: number, center: number) => MathUtils.euclideanModulo(value - center + P.radius, span) + center - P.radius;
  return {
    group,
    update(delta: number, camera: Vector3, rain: number, snowfall: number, daylight: number, viewportHeight: number) {
      const dt = Math.min(Math.max(delta, 0), P.maxDelta);
      elapsed += dt;
      rainMaterial.color.setRGB(0.65, 0.75, 0.85).multiplyScalar(0.25 + daylight * 0.75);
      snowMaterial.uniforms.tint.value.setRGB(0.9, 0.95, 1).multiplyScalar(0.25 + daylight * 0.75);
      snowMaterial.uniforms.pointScale.value = viewportHeight / 2;
      for (const layer of layers) {
        const { snow, settings, vertices, mesh, attribute, positions, ground, seeds } = layer;
        const rate = snow ? snowfall : rain;
        const strength = Math.sqrt(MathUtils.clamp(rate / settings.fullRate, 0, 1));
        const count = Math.floor(settings.maxParticles * strength);
        mesh.visible = count > 0;
        mesh.geometry.setDrawRange(0, count * vertices);
        if (!snow) rainMaterial.opacity = P.rain.opacity + strength * 0.3;
        for (let i = 0; i < count; i++) {
          const offset = i * 3, seed = seeds[i];
          const fresh = i >= layer.count;
          let x = fresh ? camera.x + (Math.random() - 0.5) * span : positions[offset];
          let z = fresh ? camera.z + (Math.random() - 0.5) * span : positions[offset + 2];
          if (snow) {
            x += Math.sin(elapsed + seed * Math.PI * 2) * P.snow.drift * dt;
            z += Math.cos(elapsed * 0.7 + seed * Math.PI * 2) * P.snow.drift * dt;
          }
          x = wrap(x, camera.x); z = wrap(z, camera.z);
          if (fresh || snow || x !== positions[offset] || z !== positions[offset + 2]) ground[i] = heightAt(x, z);
          const bottom = Math.max(camera.y - P.belowCamera, ground[i] + 0.05);
          const top = Math.max(bottom + 0.1, camera.y + P.height - P.belowCamera);
          const speed = settings.speed * (0.8 + seed * 0.4) * (snow ? 1 : 1 + strength * 0.4);
          let y = fresh ? MathUtils.lerp(bottom, top, seed) : positions[offset + 1] - speed * dt;
          if (y < bottom || y > top) y = top;
          positions[offset] = x; positions[offset + 1] = y; positions[offset + 2] = z;
          attribute.setXYZ(i * vertices, x, y, z);
          if (!snow) attribute.setXYZ(i * vertices + 1, x, y + P.rain.streak * (0.6 + strength), z);
        }
        if (count) attribute.needsUpdate = true;
        layer.count = count;
      }
      return layers.some(layer => layer.mesh.visible);
    },
    dispose() {
      for (const layer of layers) layer.mesh.geometry.dispose();
      rainMaterial.dispose(); snowMaterial.dispose();
      group.removeFromParent();
    },
  };
}