import * as THREE from 'three';

export class StandardUVMaterial extends THREE.MeshPhysicalMaterial {
  public uvUniforms: {
    uMappingMode: { value: number }; // 0: imported, 1: planar, 2: box, 3: cylindrical
    uUVTransform: { value: THREE.Matrix3 };
  };

  constructor(parameters?: THREE.MeshPhysicalMaterialParameters) {
    super(parameters);

    this.uvUniforms = {
      uMappingMode: { value: 0 },
      uUVTransform: { value: new THREE.Matrix3() },
    };

    this.customProgramCacheKey = () => `standard_uv_v2`;

    this.onBeforeCompile = (shader) => {
      shader.uniforms.uMappingMode = this.uvUniforms.uMappingMode;
      shader.uniforms.uUVTransform = this.uvUniforms.uUVTransform;

      shader.vertexShader = `
        attribute vec2 uvPlanar;
        attribute vec2 uvBox;
        attribute vec2 uvCylinder;
        
        uniform int uMappingMode;
        uniform mat3 uUVTransform;
        
        varying vec2 vCustomUV;
      ` + shader.vertexShader;

      shader.vertexShader = shader.vertexShader.replace(
        `#include <uv_vertex>`,
        `
        #include <uv_vertex>
        
        vec2 selectedUV = vec2(0.0);
        #ifdef USE_UV
          selectedUV = uv;
        #endif
        
        if (uMappingMode == 1) {
          selectedUV = uvPlanar;
        } else if (uMappingMode == 2) {
          selectedUV = uvBox;
        } else if (uMappingMode == 3) {
          selectedUV = uvCylinder;
        }
        
        vCustomUV = (uUVTransform * vec3(selectedUV, 1.0)).xy;
        `
      );

      shader.fragmentShader = `
        varying vec2 vCustomUV;
      ` + shader.fragmentShader;

      const mainStartIndex = shader.fragmentShader.indexOf('void main() {');
      if (mainStartIndex !== -1) {
        const prefix = shader.fragmentShader.substring(0, mainStartIndex);
        let mainBody = shader.fragmentShader.substring(mainStartIndex);
        
        mainBody = mainBody.replace(/\bvUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvNormalMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvRoughnessMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvMetalnessMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvAlphaMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvEmissiveMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvAoMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvClearcoatMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvClearcoatNormalMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvClearcoatRoughnessMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvTransmissionMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvThicknessMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvSpecularMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvSpecularIntensityMapUv\b/g, 'vCustomUV');
        mainBody = mainBody.replace(/\bvSpecularColorMapUv\b/g, 'vCustomUV');
        
        shader.fragmentShader = prefix + mainBody;
      }
    };
  }

  public updateUVTransform(offset: [number, number], repeat: [number, number], rotation: number) {
    this.uvUniforms.uUVTransform.value.setUvTransform(
      offset[0], offset[1],
      repeat[0], repeat[1],
      rotation,
      0.5, 0.5
    );
  }
}