precision highp float;
precision highp int;

#include <utils>
#include <cylinder_utils>

uniform float underwater;
uniform samplerCube sky;
uniform float useSky;
uniform sampler2D sceneTexture;
uniform vec2 resolution;
uniform float useSceneRefraction;
uniform float surfaceTransmittance;
uniform float normalStrength;
uniform float refractionStrength;
uniform float hasWall;
uniform vec3 surfaceColor;
uniform vec3 volumeColor;

in vec3 eye;
in vec3 pos;

out vec4 fragColor;


vec3 getSurfaceRayColor(vec3 origin, vec3 ray, vec3 waterColor) {
  vec3 color;

  if (ray.y < 0.0) {
    vec2 t = intersectCylinder(origin, ray, poolRadius, -poolHeight, poolHeight * 2.0);
    color = getCylinderWallColor(origin + ray * t.y);
  } else {
    vec2 t = intersectCylinder(origin, ray, poolRadius, -poolHeight, poolHeight * 2.0);
    vec3 hit = origin + ray * t.y;
    if (hit.y < poolHeight * 7.0 / 12.0) {
      color = getCylinderWallColor(hit);
    } else {
      if (useSky > 0.5) {
        color = texture(sky, ray).rgb;
        color += 0.01 * vec3(pow(max(0.0, dot(light, ray)), 20.0)) * vec3(10.0, 8.0, 6.0);
      } else {
        color = waterColor;
      }
    }
  }

  if (ray.y < 0.0) color *= waterColor;

  return color;
}


void main() {
  if (pos.x * pos.x + pos.z * pos.z > poolRadius * poolRadius) discard;

  vec2 coord = pos.xz / (poolRadius * 2.0) + 0.5;
  vec4 info = texture(water, coord);

  for (int i = 0; i < 5; i++) {
    coord += info.ba * 0.005;
    info = texture(water, coord);
  }

  vec2 normalXY = info.ba * normalStrength;
  vec3 normal = vec3(normalXY.x, sqrt(max(0.0, 1.0 - dot(normalXY, normalXY))), normalXY.y);
  vec3 incomingRay = normalize(pos - eye);
  vec3 surfaceTint = surfaceColor;
  vec3 volumeTint = volumeColor;

  if (underwater == 1.) {
    normal = -normal;
    vec3 reflectedRay = reflect(incomingRay, normal);
    vec3 refractedRay = refract(incomingRay, normal, IOR_WATER / IOR_AIR);
    float fresnel = mix(0.5, 1.0, pow(1.0 - dot(normal, -incomingRay), 3.0));

    vec3 reflectedColor = getSurfaceRayColor(pos, reflectedRay, volumeTint);

    vec3 refractedColor = getSurfaceRayColor(pos, refractedRay, vec3(1.0)) * vec3(0.8, 1.0, 1.1);
    if (useSceneRefraction > 0.5) {
      vec2 screenUV = gl_FragCoord.xy / resolution;
      vec2 refractedUV = clamp(screenUV + normal.xz * refractionStrength, 0.0, 1.0);
      vec3 sceneColor = texture(sceneTexture, refractedUV).rgb;
      // 有墙体时降低屏幕空间折射占比，避免水面过于“薄”
      float mixFactor = surfaceTransmittance;
      if (hasWall > 0.5) {
        mixFactor *= 0.6;
      }
      refractedColor = mix(refractedColor, sceneColor, mixFactor);
    }

    fragColor = vec4(mix(reflectedColor, refractedColor, (1.0 - fresnel) * length(refractedRay)), 1.0);
  } else {
    vec3 reflectedRay = reflect(incomingRay, normal);
    vec3 refractedRay = refract(incomingRay, normal, IOR_AIR / IOR_WATER);
    float fresnel = mix(0.25, 1.0, pow(1.0 - dot(normal, -incomingRay), 3.0));

    vec3 reflectedColor = getSurfaceRayColor(pos, reflectedRay, surfaceTint);

    vec3 refractedColor = getSurfaceRayColor(pos, refractedRay, surfaceTint);
    if (useSceneRefraction > 0.5) {
      vec2 screenUV = gl_FragCoord.xy / resolution;
      vec2 refractedUV = clamp(screenUV + normal.xz * refractionStrength, 0.0, 1.0);
      vec3 sceneColor = texture(sceneTexture, refractedUV).rgb;
      // 有墙体时降低屏幕空间折射占比，避免水面过于“薄”
      float mixFactor = surfaceTransmittance;
      if (hasWall > 0.5) {
        mixFactor *= 0.6;
      }
      refractedColor = mix(refractedColor, sceneColor, mixFactor);
    }

    fragColor = vec4(mix(refractedColor, reflectedColor, fresnel), 1.0);
  }
}