// WebGPU counterpart of the GPU heightfield raymarcher (the WGSL port of
// webgl_terrain_frag.glsl / terrain.frag). The Vulkan push constants become a
// uniform buffer of five vec4s, laid out exactly as `ra3::wasmgl::detail::terrain_uniforms`.
//
// Bind group 0: uniform data[5], the R16 heightmap (u32), the RGBA16 cell record
// (vec4<u32>), the RGBA8 tile atlas (array), and an atlas sampler.

struct TerrainUniforms {
    data: array<vec4<f32>, 5>,
};

@group(0) @binding(0) var<uniform> u: TerrainUniforms;
@group(0) @binding(1) var u_heightmap: texture_2d<u32>;    // R16Uint heights
@group(0) @binding(2) var u_celldata: texture_2d<u32>;     // RGBA16Uint blend record (texel is vec4<u32>)
@group(0) @binding(3) var u_atlas: texture_2d_array<f32>;  // RGBA8 tile materials
@group(0) @binding(4) var u_atlas_samp: sampler;

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

// SAGE water model constants (see docs/REVERSE_ENGINEERING.md).
const WATER_SCALE: f32 = 1.0 / 320.0;
const WATER_TRANSPARENT_DEPTH: f32 = 10.0;
const WATER_MIN_OPACITY: f32 = 0.70;
const WATER_RIVER_MULTIPLIER: f32 = 1.0;

fn cam_uniform() -> vec4<f32> { return u.data[0]; }    // x=target_x, y=target_y, z=yaw, w=height
fn params_uniform() -> vec4<f32> { return u.data[1]; } // x=pitch, y=fov, z=water_z, w=has_water
fn sun_uniform() -> vec4<f32> { return u.data[2]; }    // xyz=sun dir, w=ambient
fn mapinfo_uniform() -> vec4<f32> { return u.data[3]; }// x=W, y=H, z=unused, w=z_scale
fn misc_uniform() -> vec4<f32> { return u.data[4]; }   // x=time, y=unused, z=cells per repeat, w=aspect

struct TerrainOut {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32) -> TerrainOut {
    let p = vec2<f32>(f32((vertex_index << 1u) & 2u), f32(vertex_index & 2u));
    var out: TerrainOut;
    out.uv = p;
    out.position = vec4<f32>(p.x * 2.0 - 1.0, 1.0 - p.y * 2.0, 0.0, 1.0);
    return out;
}

fn height_at(cell: vec2<i32>) -> f32 {
    let limit = vec2<i32>(mapinfo_uniform().xy) - vec2<i32>(1);
    let c = clamp(cell, vec2<i32>(0), limit);
    return f32(textureLoad(u_heightmap, c, 0).r) * mapinfo_uniform().w;
}

fn world_height(wx: f32, wy: f32) -> f32 {
    let world_w = mapinfo_uniform().x * CELL;
    let world_h = mapinfo_uniform().y * CELL;
    if (wx < 0.0 || wy < 0.0 || wx >= world_w || wy >= world_h) {
        return -1.0e9;
    }
    let c = vec2<i32>(i32(wx / CELL), i32((world_h - wy) / CELL));
    return height_at(c);
}

fn sky_color(dir: vec3<f32>) -> vec3<f32> {
    let d = normalize(dir);
    let sun_dir = normalize(sun_uniform().xyz);
    let t = clamp(d.z, 0.0, 1.0);
    let horizon = vec3<f32>(0.70, 0.78, 0.85);
    let zenith = vec3<f32>(0.28, 0.48, 0.80);
    var col = mix(horizon, zenith, pow(t, 0.6));
    let sun = max(dot(d, sun_dir), 0.0);
    col += vec3<f32>(1.0, 0.95, 0.82) * pow(sun, 300.0) * 1.6;
    col += vec3<f32>(1.0, 0.90, 0.72) * pow(sun, 8.0) * 0.18;
    return col;
}

// The retail SAGE blend ramp (see terrain.frag).
fn blend_factor(direction: u32, flags: u32, f_in: vec2<f32>) -> f32 {
    var f = f_in;
    let flipped = (flags & 1u) != 0u;
    let 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) {
        let s = (1.0 - f.x) + (1.0 - f.y);
        return select(clamp(1.0 - s, 0.0, 1.0), 1.0 - clamp(s - 1.0, 0.0, 1.0), two_sided);
    }
    if (direction == 8u) {
        let s = f.x + (1.0 - f.y);
        return select(clamp(1.0 - s, 0.0, 1.0), 1.0 - clamp(s - 1.0, 0.0, 1.0), two_sided);
    }
    return 0.0;
}

fn sample_layer(layer: u32, wx: f32, wy: f32) -> vec3<f32> {
    let span = max(misc_uniform().z, 1.0);
    let layer_count = textureNumLayers(u_atlas);
    let l = f32(min(layer, layer_count - 1u));
    // The atlas is 0xAARRGGBB (BGRA in memory), uploaded as RGBA8.
    let c = textureSampleLevel(u_atlas, u_atlas_samp, vec2<f32>(wx, wy) / span, i32(l), 0.0);
    return vec3<f32>(c.b, c.g, c.r);
}

// ---- SAGE water (Ocean.fx / RiverWater.fx / Water.frag port) ----------------

// Sample an atlas layer by an explicit layer index (the water flow/bump maps
// are appended as the last two layers of the tile atlas). The atlas is
// 0xAARRGGBB, so the b/g/r swap restores RGB (as `sample_layer`).
fn water_tex(layer: i32, uv: vec2<f32>) -> vec3<f32> {
    let lc = i32(textureNumLayers(u_atlas));
    let l = clamp(layer, 0, max(lc - 1, 0));
    let c = textureSampleLevel(u_atlas, u_atlas_samp, uv, l, 0.0);
    return vec3<f32>(c.b, c.g, c.r);
}

// Scrolling wave normal on the water plane: the retail bump map (atlas's last
// layer), offset by the flow map (second-last layer) and combined with a
// de-gridded procedural wave so the sun glint is irregular and always moving.
fn water_normal(world_xy: vec2<f32>, time: f32) -> vec3<f32> {
    let q = world_xy * (WATER_SCALE * 6.0);
    let a1 = q.x * 0.80 + q.y * 0.30 + time * 1.10;
    let a2 = q.y * 0.95 - q.x * 0.45 - time * 1.30;
    let a3 = (q.x + q.y) * 1.60 + time * 2.10;
    let a4 = (q.x - q.y) * 2.30 - time * 1.70;
    let dx = 0.224 * cos(a1) - 0.099 * cos(a2) + 0.256 * cos(a3) + 0.230 * cos(a4);
    let dy = 0.084 * cos(a1) + 0.209 * cos(a2) + 0.256 * cos(a3) - 0.230 * cos(a4);
    let lc = i32(textureNumLayers(u_atlas));
    let flow = water_tex(lc - 2, q - vec2<f32>(time * 0.010, time * 0.014)) * 2.0 - 1.0;
    let bump = water_tex(lc - 1, q + flow.xy * 0.05 + vec2<f32>(time * 0.006, time * 0.008)) * 2.0 - 1.0;
    let sx = -dx * 0.30 + bump.x * 0.45;
    let sy = -dy * 0.30 + bump.y * 0.45;
    return normalize(vec3<f32>(sx, sy, 1.0));
}

// Water.frag distortionPower * the flow texture: a small scrolling UV offset.
fn water_distortion(world_xy: vec2<f32>, time: f32) -> f32 {
    let q = world_xy * (WATER_SCALE * 6.0);
    return 0.05 * (sin(q.x * 0.9 + time * 0.7) + sin(q.y * 1.1 - time * 0.5));
}

// Water.frag GetCloudColor: no cloud texture is bound, so a slow scroll stands in.
fn water_cloud(world_xy: vec2<f32>, time: f32) -> vec3<f32> {
    return vec3<f32>(0.85 + 0.15 * sin((world_xy.x + world_xy.y) * 0.0007 - time * 0.05));
}

// Retail UnderwaterDeferred.fx: absorbs red and fogs with distance.
fn apply_underwater(color: vec3<f32>, distance: f32, cam_z: f32, water_z: f32) -> vec3<f32> {
    if (cam_z >= water_z - 0.5) { return color; }
    let absorb = vec3<f32>(0.35, 0.62, 0.75);
    let fog = clamp(1.0 - exp(-distance * 0.00022), 0.0, 0.9);
    return mix(color * absorb, vec3<f32>(0.02, 0.10, 0.16), fog);
}

// Water.frag: fresnel-mixed sky reflection and depth-graded seabed refraction,
// SAGE diffuse + specular lighting, cloud term, depth-based transparency fade.
fn water_shade(hitpos: vec3<f32>, dir: vec3<f32>, distance: f32) -> vec3<f32> {
    let time = misc_uniform().x;
    let river = misc_uniform().y;
    let seabed = world_height(hitpos.x, hitpos.y);
    let depth = max(0.0, params_uniform().z - seabed);

    let n = water_normal(hitpos.xy, time + water_distortion(hitpos.xy, time));
    let sun_dir = normalize(sun_uniform().xyz);

    let cos_theta = clamp(dot(-dir, n), 0.0, 1.0);
    let fresnel = 0.02 + 0.98 * pow(1.0 - cos_theta, 5.0);

    let reflection = sky_color(reflect(dir, n));
    let shallow = vec3<f32>(0.10, 0.34, 0.38);
    let deep = vec3<f32>(0.02, 0.12, 0.22);
    let refraction = mix(shallow, deep, clamp(depth / 40.0, 0.0, 1.0));

    let ndotl = max(dot(n, sun_dir), 0.0);
    let ambient = sun_uniform().w;
    let diffuse = vec3<f32>(ambient + (1.0 - ambient) * ndotl);
    let half_v = normalize(sun_dir - dir);
    let spec = pow(max(dot(n, half_v), 0.0), 90.0);

    var color = mix(refraction, reflection, clamp(fresnel, 0.0, 1.0)) * diffuse * water_cloud(hitpos.xy, time);
    color += vec3<f32>(1.0, 0.97, 0.9) * spec * 0.45;

    var alpha = clamp(depth / WATER_TRANSPARENT_DEPTH, 0.0, 1.0) * WATER_MIN_OPACITY;
    if (river > 0.5) { alpha *= WATER_RIVER_MULTIPLIER; }
    color = mix(refraction, color, clamp(alpha + 0.15, 0.0, 1.0));

    let wfog = clamp(1.0 - exp(-distance * 0.00009), 0.0, 0.75);
    return mix(color, sky_color(vec3<f32>(dir.x, dir.y, 0.0)), wfog);
}

@fragment
fn fs_main(in: TerrainOut) -> @location(0) vec4<f32> {
    let p = cam_uniform();
    let pitch = clamp(params_uniform().x, 0.15, 1.45);
    let fov = clamp(params_uniform().y, 0.3, 1.4);
    let world_w = mapinfo_uniform().x * CELL;
    let world_h = mapinfo_uniform().y * CELL;

    let cp = cos(pitch);
    let fwd = vec3<f32>(cp * sin(p.z), cp * cos(p.z), -sin(pitch));
    let right = normalize(cross(fwd, vec3<f32>(0.0, 0.0, 1.0)));
    let up = cross(right, fwd);

    var target_z = world_height(p.x, p.y);
    if (target_z < -1.0e8) { target_z = 0.0; }
    let dist = p.w / sin(pitch);
    let cam = vec3<f32>(p.x, p.y, target_z + p.w) - fwd * dist;

    let ndc = vec2<f32>(in.uv.x * 2.0 - 1.0, 1.0 - in.uv.y * 2.0);
    let aspect = misc_uniform().w;
    let th = tan(fov * 0.5);
    let dir = normalize(fwd + right * ndc.x * th * aspect + up * ndc.y * th);

    if (dir.z >= -1e-4) {
        return vec4<f32>(apply_underwater(sky_color(dir), 0.0, cam.z, params_uniform().z), 1.0);
    }

    // Clip the ray to the map's XY rectangle: the boundary is an exact plane,
    // so the silhouette there stays razor-sharp instead of stair-stepping
    // across it. Outside the map is sky.
    var t_enter = 0.0;
    var t_exit = 1.0e30;
    var inside = true;
    if (abs(dir.x) < 1e-6) {
        inside = (cam.x >= 0.0 && cam.x <= world_w);
    } else {
        let a = (0.0 - cam.x) / dir.x;
        let b = (world_w - cam.x) / dir.x;
        t_enter = max(t_enter, min(a, b));
        t_exit = min(t_exit, max(a, b));
    }
    if (inside) {
        if (abs(dir.y) < 1e-6) {
            inside = (cam.y >= 0.0 && cam.y <= world_h);
        } else {
            let a = (0.0 - cam.y) / dir.y;
            let b = (world_h - cam.y) / dir.y;
            t_enter = max(t_enter, min(a, b));
            t_exit = min(t_exit, max(a, b));
        }
    }
    if (!inside || t_exit <= 0.0) {
        return vec4<f32>(apply_underwater(sky_color(dir), 0.0, cam.z, params_uniform().z), 1.0);
    }

    // March the heightfield cell by cell: the step is never longer than the
    // time to cross one cell (dominant horizontal axis), while a clearance term
    // lets the ray skip the empty air above the surface. Resolving every cell is
    // what keeps cliff and map-edge silhouettes from quantising into huge
    // stair-steps that crawl as the camera pans.
    let horiz = max(abs(dir.x), abs(dir.y));
    let cell_step = min(CELL / max(horiz, 1e-4), CELL * 32.0);
    var t = max(t_enter, CELL * 0.5);
    var prev = t;
    var hit = false;
    var hit_t = 0.0;
    for (var i = 0; i < 1024 && t <= t_exit; i = i + 1) {
        let w = cam + dir * t;
        let h = world_height(w.x, w.y);
        let surface = select(h, max(h, params_uniform().z), params_uniform().w > 0.5);
        if (w.z <= surface) {
            hit = true;
            hit_t = t;
            break;
        }
        let clearance = (w.z - surface) / max(-dir.z, 1e-4);
        prev = t;
        t += clamp(clearance, cell_step, cell_step * 8.0);
    }
    if (!hit) {
        return vec4<f32>(apply_underwater(sky_color(dir), 0.0, cam.z, params_uniform().z), 1.0);
    }

    // Refine the first crossing; with a sub-cell bracket this converges to the
    // exact surface point.
    var lo = prev;
    var hi = hit_t;
    for (var i = 0; i < 18; i = i + 1) {
        let mid = 0.5 * (lo + hi);
        let w = cam + dir * mid;
        let h = world_height(w.x, w.y);
        let surface = select(h, max(h, params_uniform().z), params_uniform().w > 0.5);
        if (w.z <= surface) {
            hi = mid;
        } else {
            lo = mid;
        }
    }
    let hitpos = cam + dir * hi;

    let sun = normalize(sun_uniform().xyz);
    let ambient = sun_uniform().w;

    if (params_uniform().w > 0.5 && hitpos.z <= params_uniform().z + 0.01) {
        return vec4<f32>(apply_underwater(water_shade(hitpos, dir, hi), hi, cam.z, params_uniform().z), 1.0);
    }

    let wx = hitpos.x / CELL;
    let wy = (world_h - hitpos.y) / CELL;
    let cx = clamp(i32(wx), 0, i32(mapinfo_uniform().x) - 1);
    let cy = clamp(i32(wy), 0, i32(mapinfo_uniform().y) - 1);
    let fx = wx - floor(wx);
    let fy = wy - floor(wy);

    let record = textureLoad(u_celldata, vec2<i32>(cx, cy), 0);
    let packed = record.w;
    let dir1 = packed & 0xFu;
    let flags1 = (packed >> 4u) & 0x3u;
    let dir2 = (packed >> 8u) & 0xFu;
    let flags2 = (packed >> 12u) & 0x3u;

    let frac_uv = vec2<f32>(fx, fy);
    let c0 = sample_layer(record.x, wx, wy);
    let c1 = sample_layer(record.y, wx, wy);
    let c2 = sample_layer(record.z, wx, wy);
    let f1 = blend_factor(dir1, flags1, frac_uv);
    let f2 = blend_factor(dir2, flags2, frac_uv);
    let albedo = mix(mix(c0, c1, f1), c2, f2);

    let hl = world_height(hitpos.x - CELL, hitpos.y);
    let hr = world_height(hitpos.x + CELL, hitpos.y);
    let hd = world_height(hitpos.x, hitpos.y - CELL);
    let hu = world_height(hitpos.x, hitpos.y + CELL);
    let n = normalize(vec3<f32>(hl - hr, hd - hu, 2.0 * CELL));
    let lambert = max(0.0, dot(n, sun));
    var lit = albedo * (ambient + (1.0 - ambient) * lambert);
    let fog = clamp(1.0 - exp(-hi * 0.00009), 0.0, 0.75);
    lit = mix(lit, sky_color(vec3<f32>(dir.x, dir.y, 0.0)), fog);
    return vec4<f32>(apply_underwater(lit, hi, cam.z, params_uniform().z), 1.0);
}