#version 300 es
// GPU heightfield raymarcher for the real RA3 terrain (the WebGL port of
// terrain.frag). The Vulkan push constants become a set of vec4 uniforms.
precision highp float;
precision highp int;
precision highp sampler2D;
precision highp sampler2DArray;

uniform sampler2D u_heightmap;      // R16 heights
uniform sampler2D u_celldata;       // per-cell blend record (RGBA16)
uniform sampler2DArray u_atlas;     // tile material array (RGBA8)

// One contiguous array so the host can upload all five vec4s with a single
// glUniform4fv; the names keep the shader body identical to terrain.frag.
uniform vec4 u_data[5];
#define u_cam u_data[0]     // x=target_x, y=target_y, z=yaw, w=height
#define u_params u_data[1]  // x=pitch, y=fov, z=water_z, w=has_water
#define u_sun u_data[2]     // xyz=sun dir, w=ambient
#define u_mapinfo u_data[3] // x=W, y=H, z=unused, w=z_scale
#define u_misc u_data[4]    // x=time, y=unused, z=cells per texture repeat, w=aspect

in vec2 v_uv;
out vec4 frag_color;

const float CELL = 10.0; // must match ra3::terrain::cell_size

float height_at(ivec2 c) {
    c = clamp(c, ivec2(0), ivec2(u_mapinfo.xy) - 1);
    return texelFetch(u_heightmap, c, 0).r * 65535.0 * u_mapinfo.w;
}

float world_height(float wx, float wy) {
    float world_w = u_mapinfo.x * CELL;
    float world_h = u_mapinfo.y * CELL;
    if (wx < 0.0 || wy < 0.0 || wx >= world_w || wy >= world_h) return -1.0e9;
    ivec2 c = ivec2(int(wx / CELL), int((world_h - wy) / CELL));
    return height_at(c);
}

vec3 sky_color(vec3 dir) {
    vec3 d = normalize(dir);
    vec3 sun_dir = normalize(u_sun.xyz);
    float t = clamp(d.z, 0.0, 1.0);
    vec3 horizon = vec3(0.70, 0.78, 0.85);
    vec3 zenith = vec3(0.28, 0.48, 0.80);
    vec3 col = mix(horizon, zenith, pow(t, 0.6));
    float sun = max(dot(d, sun_dir), 0.0);
    col += vec3(1.0, 0.95, 0.82) * pow(sun, 300.0) * 1.6; // sun disk
    col += vec3(1.0, 0.90, 0.72) * pow(sun, 8.0) * 0.18;  // glow
    return col;
}

// The retail SAGE blend ramp (see terrain.frag).
float blend_factor(uint direction, uint flags, vec2 f) {
    bool flipped = (flags & 1u) != 0u;
    bool two_sided = (flags & 2u) != 0u;
    if (flipped) {
        if (direction == 1u) {
            f.x = 1.0 - f.x;
        } else if (direction == 2u || direction == 4u || direction == 8u) {
            f.y = 1.0 - f.y;
        }
    }
    if (direction == 1u) return f.x;
    if (direction == 2u) return f.y;
    if (direction == 4u) {
        float s = (1.0 - f.x) + (1.0 - f.y);
        return two_sided ? 1.0 - clamp(s - 1.0, 0.0, 1.0) : clamp(1.0 - s, 0.0, 1.0);
    }
    if (direction == 8u) {
        float s = f.x + (1.0 - f.y);
        return two_sided ? 1.0 - clamp(s - 1.0, 0.0, 1.0) : clamp(1.0 - s, 0.0, 1.0);
    }
    return 0.0;
}

vec3 sample_layer(uint layer, float wx, float wy) {
    float span = max(u_misc.z, 1.0);
    int layer_count = textureSize(u_atlas, 0).z;
    float l = float(min(layer, uint(layer_count - 1)));
    // The atlas is 0xAARRGGBB (BGRA in memory), uploaded as GL_RGBA.
    return texture(u_atlas, vec3(vec2(wx, wy) / span, l)).bgr;
}

void main() {
    vec4 p = u_cam;
    float pitch = clamp(u_params.x, 0.15, 1.45);
    float fov = clamp(u_params.y, 0.3, 1.4);
    float world_w = u_mapinfo.x * CELL;
    float world_h = u_mapinfo.y * CELL;

    float cp = cos(pitch);
    vec3 fwd = vec3(cp * sin(p.z), cp * cos(p.z), -sin(pitch));
    vec3 right = normalize(cross(fwd, vec3(0, 0, 1)));
    vec3 up = cross(right, fwd);

    float target_z = world_height(p.x, p.y);
    if (target_z < -1.0e8) target_z = 0.0;
    float dist = p.w / sin(pitch);
    vec3 cam = vec3(p.x, p.y, target_z + p.w) - fwd * dist;

    vec2 ndc = vec2(v_uv.x * 2.0 - 1.0, 1.0 - v_uv.y * 2.0);
    float aspect = u_misc.w;
    float th = tan(fov * 0.5);
    vec3 dir = normalize(fwd + right * ndc.x * th * aspect + up * ndc.y * th);

    if (dir.z >= -1e-4) {
        frag_color = vec4(sky_color(dir), 1.0);
        return;
    }

    float t = CELL * 0.5;
    float dt = CELL * 0.5;
    float prev = t;
    bool hit = false;
    float hit_t = 0.0;
    for (int i = 0; i < 512 && t < 60000.0; ++i) {
        vec3 w = cam + dir * t;
        if (w.x < 0.0 || w.y < 0.0 || w.x >= world_w || w.y >= world_h) {
            prev = t;
            dt *= 1.06;
            t += dt;
            continue;
        }
        if (u_params.w > 0.5 && w.z <= u_params.z) {
            hit = true;
            hit_t = t;
            break;
        }
        if (w.z <= world_height(w.x, w.y)) {
            hit = true;
            hit_t = t;
            break;
        }
        prev = t;
        dt *= 1.06;
        t += dt;
    }
    if (!hit) {
        frag_color = vec4(sky_color(dir), 1.0);
        return;
    }

    float lo = prev;
    float hi = hit_t;
    for (int i = 0; i < 6; ++i) {
        float mid = 0.5 * (lo + hi);
        vec3 w = cam + dir * mid;
        bool water = u_params.w > 0.5 && w.z <= u_params.z;
        if (water || w.z <= world_height(w.x, w.y)) {
            hi = mid;
        } else {
            lo = mid;
        }
    }
    vec3 hitpos = cam + dir * hi;

    vec3 sun = normalize(u_sun.xyz);
    float ambient = u_sun.w;

    if (u_params.w > 0.5 && hitpos.z <= u_params.z + 0.01) {
        float time = u_misc.x;
        vec2 q = hitpos.xy * 0.015;
        float nx = sin(q.x * 1.3 + time * 1.7) + 0.5 * sin(q.x * 3.1 - time * 2.3);
        float ny = sin(q.y * 1.1 - time * 1.3) + 0.5 * sin(q.y * 2.7 + time * 1.9);
        vec3 n = normalize(vec3(nx * 0.06, ny * 0.06, 1.0));
        float fres = pow(1.0 - clamp(-dir.z, 0.0, 1.0), 3.0);
        vec3 deep = vec3(0.03, 0.16, 0.28);
        vec3 refl = sky_color(reflect(dir, n));
        float lam = max(0.0, dot(n, sun));
        vec3 water = mix(deep, refl, clamp(0.25 + 0.55 * fres, 0.0, 0.9));
        water += vec3(1.0, 0.98, 0.9) * pow(lam, 64.0) * 0.6;
        float wfog = clamp(1.0 - exp(-hi * 0.00009), 0.0, 0.75);
        water = mix(water, sky_color(vec3(dir.x, dir.y, 0.0)), wfog);
        frag_color = vec4(water, 1.0);
        return;
    }

    float wx = hitpos.x / CELL;
    float wy = (world_h - hitpos.y) / CELL;
    int cx = clamp(int(wx), 0, int(u_mapinfo.x) - 1);
    int cy = clamp(int(wy), 0, int(u_mapinfo.y) - 1);
    float fx = wx - floor(wx);
    float fy = wy - floor(wy);

    uvec4 record = uvec4(texelFetch(u_celldata, ivec2(cx, cy), 0) * 65535.0 + 0.5);
    uint packed = record.w;
    uint dir1 = packed & 0xFu;
    uint flags1 = (packed >> 4u) & 0x3u;
    uint dir2 = (packed >> 8u) & 0xFu;
    uint flags2 = (packed >> 12u) & 0x3u;

    vec2 fracUV = vec2(fx, fy);
    vec3 c0 = sample_layer(record.x, wx, wy);
    vec3 c1 = sample_layer(record.y, wx, wy);
    vec3 c2 = sample_layer(record.z, wx, wy);
    float f1 = blend_factor(dir1, flags1, fracUV);
    float f2 = blend_factor(dir2, flags2, fracUV);
    vec3 albedo = mix(mix(c0, c1, f1), c2, f2);

    float hl = world_height(hitpos.x - CELL, hitpos.y);
    float hr = world_height(hitpos.x + CELL, hitpos.y);
    float hd = world_height(hitpos.x, hitpos.y - CELL);
    float hu = world_height(hitpos.x, hitpos.y + CELL);
    vec3 n = normalize(vec3(hl - hr, hd - hu, 2.0 * CELL));
    float lambert = max(0.0, dot(n, sun));
    vec3 lit = albedo * (ambient + (1.0 - ambient) * lambert);
    float fog = clamp(1.0 - exp(-hi * 0.00009), 0.0, 0.75);
    lit = mix(lit, sky_color(vec3(dir.x, dir.y, 0.0)), fog);
    frag_color = vec4(lit, 1.0);
}