#include <common/lighting/brdf>
#include <common/lighting/functions>

vec3 CalculateDirectDiffuse(in LightingIntermediateData lightingData, in vec4 lightSourceColorAndIntensity)
{
    vec3 irradiance = vec3(lightingData.NoLSat) * lightSourceColorAndIntensity.rgb * lightSourceColorAndIntensity.w;
    return irradiance * DiffuseCoefficient_EnergyConservation;
}

vec3 CalculateDirectSpecular(in LightingIntermediateData lightingData, in vec4 lightSourceColorAndIntensity)
{
    vec3 irradiance = vec3(lightingData.NoLSat) * lightSourceColorAndIntensity.rgb * lightSourceColorAndIntensity.w;

    float roughness = lightingData.specularParam;
  #if CC_SURFACES_LIGHTING_ANISOTROPIC
      float rT, rB;
      GetAnisotropicRoughness(roughness, lightingData.anisotropyShape, rT, rB);
      float calcSpec = D_GGXAniso(rT, rB, lightingData.NoHSat, lightingData.H, lightingData.T, lightingData.B);
  #else
    #if CC_SURFACES_USE_LEGACY_COMPATIBLE_LIGHTING
      float calcSpec = (roughness * 0.25 + 0.25) * D_GGXMobile(roughness, lightingData.NoHSat);
    #else
      float calcSpec = D_GGX(roughness, lightingData.NoHSat);
    #endif
  #endif
    return irradiance * calcSpec;
}


#if CC_SURFACES_LIGHTING_ANISOTROPIC && CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT
  vec3 EnvAnisotropicReflection(samplerCube tex, vec3 R, float roughness, float mipCount, float anisotropyShape, vec3 V, vec3 N, vec3 T, vec3 B) {
      R = normalize(R);
      float integratedBRDF = 0.0;
      vec3 envSpec = vec3(0.0);
      // One direction sample count
      const int SAMPLE_STEP_COUNT = CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT;
      float sampleAngleRange = PI * abs(anisotropyShape);
      vec3 anisoDirection = anisotropyShape < 0.0 ? T : B;
      vec3 ROnNormalPlane = normalize(R - anisoDirection * dot(R, anisoDirection));
      //ROnTangentPlane = R; //for example: cross-style
      vec3 stepOffset = normalize(ROnNormalPlane - N) * (sampleAngleRange / float(SAMPLE_STEP_COUNT * 2));

      for (int i = -SAMPLE_STEP_COUNT; i <= SAMPLE_STEP_COUNT; ++i)
      {
          float rT, rB;
          GetAnisotropicRoughness(roughness, anisotropyShape, rT, rB);
          #if CC_IBL_CONVOLUTED
            float coef = abs(float(i)) / float(SAMPLE_STEP_COUNT) * float(SAMPLE_STEP_COUNT);
          #else
            float coef = pow(abs(float(i)) / float(SAMPLE_STEP_COUNT), 1.3) * float(SAMPLE_STEP_COUNT);
          #endif
          vec3 H = normalize(N + stepOffset * sign(float(i)) * coef);
          vec3 L = reflect(-V, H);
          float NoHSat = saturate(dot(N, H));
          float calcSpec = D_GGXAniso(rT, rB, NoHSat, H, T, B);

          envSpec += calcSpec * EnvReflection(tex, L, roughness, mipCount);
          integratedBRDF += calcSpec;
      }
      envSpec /= integratedBRDF;
      return envSpec;
  }
#endif

// for skybox IBL
vec3 SampleEnvironmentSpecular(samplerCube tex, in LightingIntermediateData lightingData, float mipCount)
{
    vec3 envSpec = vec3(0.0);
    float roughness = lightingData.specularParam;
    #if CC_SURFACES_LIGHTING_ANISOTROPIC && !CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT
      vec3 R = GetAnisotropicReflect(roughness, lightingData.anisotropyShape, lightingData.V, lightingData.N, lightingData.T, lightingData.B);
    #else
      vec3 R = CalculateReflectDirection(lightingData.N, lightingData.V, lightingData.NoV);
    #endif

    #if CC_SURFACES_LIGHTING_ANISOTROPIC && CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT
      envSpec = EnvAnisotropicReflection(tex, R, roughness, mipCount, lightingData.anisotropyShape, lightingData.V, lightingData.N, lightingData.T, lightingData.B);
    #else
      #if CC_SURFACES_USE_REFLECTION_DENOISE && !CC_IBL_CONVOLUTED
        envSpec = EnvReflectionWithMipFiltering(normalize(R), roughness, mipCount, 0.6);
      #else
        envSpec = EnvReflection(tex, R, roughness, mipCount);
      #endif
    #endif
    return envSpec;
}

// for reflection probe
vec3 SampleEnvironmentSpecular(samplerCube tex, in LightingIntermediateData lightingData, float mipCount, vec3 worldPos, vec3 cubeCenterPos, vec3 boxHalfSize, bool isRGBE)
{
    vec3 envSpec = vec3(0.0);
    float roughness = lightingData.specularParam;
    #if CC_SURFACES_LIGHTING_ANISOTROPIC && !CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT
      vec3 R = GetAnisotropicReflect(roughness, lightingData.anisotropyShape, lightingData.V, lightingData.N, lightingData.T, lightingData.B);
    #else
      vec3 R = CalculateReflectDirection(lightingData.N, lightingData.V, lightingData.NoV);
    #endif

    vec4 fixedR = CalculateBoxProjectedDirection(R, worldPos, cubeCenterPos, boxHalfSize);
    R = fixedR.xyz;

    vec3 envmap = SampleEnvironmentSpecular(cc_environment, lightingData, cc_ambientGround.w).xyz * cc_ambientSky.w;
    #if CC_SURFACES_LIGHTING_ANISOTROPIC && CC_SURFACES_LIGHTING_ANISOTROPIC_ENVCONVOLUTION_COUNT
      envSpec = EnvAnisotropicReflection(tex, fixedR.xyz, roughness, mipCount, lightingData.anisotropyShape, lightingData.V, lightingData.N, lightingData.T, lightingData.B);
      #if CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_CUBE
        envSpec = mix(envmap, envSpec, fixedR.w);
      #endif
    #else
      #if CC_SURFACES_USE_REFLECTION_DENOISE && !CC_IBL_CONVOLUTED
        envSpec = EnvReflectionWithMipFiltering(normalize(R), roughness, mipCount, 0.6);
      #else
        envSpec = EnvReflectionOfReflectionProbe(tex, R, roughness, mipCount, isRGBE);
        #if CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_CUBE
          envSpec = mix(envmap, envSpec, fixedR.w);
        #endif
      #endif
    #endif
    return envSpec;
}


vec3 CalculateEnvironmentDiffuse(in LightingIntermediateData lightingData, float lightIntensity)
{
  // Hemisphere Lighting
  float fAmb = max(EPSILON, 0.5 - lightingData.N.y * 0.5);
  vec3 ambDiff = mix(cc_ambientSky.rgb, cc_ambientGround.rgb, fAmb);

  // Diffuse Map
  #if CC_USE_IBL
    #if CC_USE_DIFFUSEMAP && !CC_USE_LIGHT_PROBE
      // Diffuse irradiance
      vec3 rotationDir = RotationVecFromAxisY(lightingData.N, cc_surfaceTransform.z, cc_surfaceTransform.w);
      vec4 diffuseMap = texture(cc_diffuseMap, rotationDir);
      #if CC_USE_DIFFUSEMAP == IBL_RGBE
        ambDiff = unpackRGBE(diffuseMap);
      #else
        ambDiff = SRGBToLinear(diffuseMap.rgb);
      #endif
    #endif
  #endif

  ambDiff.rgb *= lightIntensity;

  // Probe
  #if CC_USE_LIGHT_PROBE
    ambDiff.rgb += SHEvaluate(lightingData.N);
  #endif

  return ambDiff.rgb;
}

vec3 CalculateEnvironmentSpecular(in LightingIntermediateData lightingData, float lightIntensity)
{
  vec3 envSpec = vec3(0.0);

#if CC_USE_REFLECTION_PROBE
    vec3 worldPos;
    HIGHP_VALUE_FROM_STRUCT_DEFINED(worldPos, lightingData.worldPosition);

  #if CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_CUBE
    if(FSInput_reflectionProbeId < 0.0){
        envSpec = SampleEnvironmentSpecular(cc_environment, lightingData, cc_ambientGround.w);
    }else{
      vec3 centerPos, boxHalfSize;
      float mipCount;
      GetCubeReflectionProbeData(centerPos, boxHalfSize, mipCount, FSInput_reflectionProbeId);
      envSpec = SampleEnvironmentSpecular(cc_reflectionProbeCubemap, lightingData, mipCount, worldPos, centerPos, boxHalfSize, isReflectProbeUsingRGBE(FSInput_reflectionProbeId));
    }
  #elif CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_PLANAR
    vec3 R = normalize(CalculateReflectDirection(lightingData.N, lightingData.V, lightingData.NoV));
    if(FSInput_reflectionProbeId < 0.0){
        vec2 screenUV = GetPlanarReflectScreenUV(worldPos, cc_matViewProj, cc_cameraPos.w, lightingData.V, R);
        envSpec = unpackRGBE(fragTextureLod(cc_reflectionProbePlanarMap, screenUV, 1.0)).xyz;
    }else{
        vec4 plane;
        float planarReflectionDepthScale, mipCount;
        GetPlanarReflectionProbeData(plane, planarReflectionDepthScale, mipCount, FSInput_reflectionProbeId);
        vec3 worldPosOffset = CalculatePlanarReflectPositionOnPlane(lightingData.N, lightingData.V, worldPos, plane, cc_cameraPos.xyz, planarReflectionDepthScale);
        vec2 screenUV = GetPlanarReflectScreenUV(worldPosOffset, cc_matViewProj, cc_cameraPos.w, lightingData.V, R);
        envSpec = unpackRGBE(fragTextureLod(cc_reflectionProbePlanarMap, screenUV, mipCount)).xyz;
    }
  #elif CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_BLEND || CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_BLEND_AND_SKYBOX
    if(FSInput_reflectionProbeId < 0.0){
        envSpec = SampleEnvironmentSpecular(cc_environment, lightingData, cc_ambientGround.w);
    }else{
      vec3 centerPos, boxHalfSize;
      float mipCount;
      GetCubeReflectionProbeData(centerPos, boxHalfSize, mipCount, FSInput_reflectionProbeId);
      envSpec = SampleEnvironmentSpecular(cc_reflectionProbeCubemap, lightingData, mipCount, worldPos, centerPos, boxHalfSize, isReflectProbeUsingRGBE(FSInput_reflectionProbeId));

      float blendFactor = 0.0;
      #if USE_INSTANCING
        blendFactor = FSInput_reflectionProbeData.x;
      #else
        blendFactor = cc_reflectionProbeBlendData1.w;
      #endif

      if(FSInput_reflectionProbeBlendId < 0.0)
      {
        vec3 skyBoxEnv = SampleEnvironmentSpecular(cc_environment, lightingData, cc_ambientGround.w).rgb * lightIntensity;
        #if CC_USE_REFLECTION_PROBE == REFLECTION_PROBE_TYPE_BLEND_AND_SKYBOX
          //blend with skybox
          envSpec = mix(envSpec, skyBoxEnv, blendFactor);
        #else
          vec3 R = normalize(CalculateReflectDirection(lightingData.N, lightingData.V, lightingData.NoV));
          vec4 fixedR = CalculateBoxProjectedDirection(R, worldPos, centerPos, boxHalfSize);
          envSpec = mix(skyBoxEnv, envSpec, fixedR.w);
        #endif
      }
      // Disable probe blend for WebGPU
      // else {
      //   vec3 centerPosBlend, boxHalfSizeBlend;
      //   float mipCountBlend;
      //   GetBlendCubeReflectionProbeData(centerPosBlend, boxHalfSizeBlend, mipCountBlend, FSInput_reflectionProbeBlendId);
      //   vec3 probeBlend = SampleEnvironmentSpecular(cc_reflectionProbeBlendCubemap, lightingData, mipCountBlend, worldPos, centerPosBlend, boxHalfSizeBlend, isBlendReflectProbeUsingRGBE(FSInput_reflectionProbeBlendId));
      //   envSpec = mix(envSpec, probeBlend, blendFactor);
      // }
    }
  #endif
#elif CC_USE_IBL
    envSpec = SampleEnvironmentSpecular(cc_environment, lightingData, cc_ambientGround.w);
#endif

  #if CC_USE_REFLECTION_PROBE
    //If using reflection probe, no need to multiply by the ambient light intensity.
    lightIntensity = FSInput_reflectionProbeId < 0.0 ? lightIntensity : 1.0;
  #endif

  return envSpec * lightIntensity;
}