import * as THREE from 'three';
import type {Shader} from '../../utils/Types';

export class OcclusionUtils {
  private static webgpuMaterialHandler?: (material: THREE.Material) => Shader;
  private static pendingMaterials: Array<{
    material: THREE.Material;
    onShaderReady?: (shader: Shader) => void;
  }> = [];

  /**
   * Registers or clears the WebGPU TSL material occlusion handler.
   * Called internally by `Depth.init()` when `WebGPURenderer` is active.
   */
  static setWebGPUMaterialHandler(
    handler: ((material: THREE.Material) => Shader) | undefined
  ) {
    this.webgpuMaterialHandler = handler;
    if (handler && this.pendingMaterials.length > 0) {
      const pending = this.pendingMaterials.splice(0);
      for (const {material, onShaderReady} of pending) {
        const shader = handler(material);
        onShaderReady?.(shader);
      }
    } else if (!handler) {
      this.pendingMaterials.length = 0;
    }
  }

  /**
   * Configures a material for depth occlusion across both `WebGLRenderer` and
   * `WebGPURenderer`, invoking `onShaderReady` with the uniform handle to
   * register in `Depth.occludableShaders`.
   */
  static addOcclusionToMaterial(
    material: THREE.Material,
    onShaderReady?: (shader: Shader) => void
  ) {
    material.transparent = true;
    if (this.webgpuMaterialHandler) {
      const shader = this.webgpuMaterialHandler(material);
      onShaderReady?.(shader);
      return;
    }
    this.pendingMaterials.push({material, onShaderReady});
    const previous = material.onBeforeCompile;
    material.onBeforeCompile = (shader, renderer) => {
      previous.call(material, shader, renderer);
      OcclusionUtils.addOcclusionToShader(shader);
      onShaderReady?.(shader);
    };
    material.needsUpdate = true;
  }

  /**
   * Creates a simple material used for rendering objects into the occlusion
   * map. This material is intended to be used with `renderer.overrideMaterial`.
   * @returns A new instance of THREE.MeshBasicMaterial.
   */
  static createOcclusionMapOverrideMaterial() {
    return new THREE.MeshBasicMaterial();
  }

  /**
   * Modifies a material's shader in-place to incorporate distance-based
   * alpha occlusion. This is designed to be used with a material's
   * `onBeforeCompile` property. This only works with built-in three.js
   * materials.
   * @param shader - The shader object provided by onBeforeCompile.
   */
  static addOcclusionToShader(shader: Shader) {
    shader.uniforms.occlusionEnabled = {value: true};
    shader.uniforms.tOcclusionMap = {value: null};
    shader.uniforms.uOcclusionClipFromWorld = {value: new THREE.Matrix4()};
    shader.defines = {USE_UV: true, DISTANCE: true};
    shader.vertexShader = shader.vertexShader
      .replace(
        '#include <common>',
        [
          'uniform mat4 uOcclusionClipFromWorld;',
          'varying vec4 vOcclusionScreenCoord;',
          '#include <common>',
        ].join('\n')
      )
      .replace(
        '#include <fog_vertex>',
        [
          '#include <fog_vertex>',
          'vOcclusionScreenCoord = uOcclusionClipFromWorld * worldPosition;',
        ].join('\n')
      );
    shader.fragmentShader = shader.fragmentShader
      .replace(
        'uniform vec3 diffuse;',
        [
          'uniform vec3 diffuse;',
          'uniform bool occlusionEnabled;',
          'uniform sampler2D tOcclusionMap;',
          'varying vec4 vOcclusionScreenCoord;',
        ].join('\n')
      )
      .replace(
        'vec4 diffuseColor = vec4( diffuse, opacity );',
        [
          'vec4 diffuseColor = vec4( diffuse, opacity );',
          'vec2 occlusion_coordinates = 0.5 + 0.5 * vOcclusionScreenCoord.xy / vOcclusionScreenCoord.w;',
          'vec2 occlusion_sample = texture2D(tOcclusionMap, occlusion_coordinates.xy).rg;',
          'occlusion_sample = occlusion_sample / max(0.0001, occlusion_sample.g);',
          'float occlusion_value = clamp(occlusion_sample.r, 0.0, 1.0);',
          'diffuseColor.a *= occlusionEnabled ? occlusion_value : 1.0;',
        ].join('\n')
      );
  }
}