// GPU heightfield raymarcher for the Direct3D backends (the D3D port of
// terrain.vert/terrain.frag).
//
// Textures: heightmap (R16), a per-cell blend record (R16G16B16A16_UNORM: base
// layer, blend layer, three-way layer, packed direction/flags; unpacked with
// *65535) and a texture array of the tile materials (RGBA8, REPEAT). The
// material is sampled continuously (`uv = cell / span`), as the retail
// `Terrain.fx` / OpenSAGE `Terrain.frag` do, so it never restarts at a cell
// edge; material boundaries cross-fade with the SAGE blend ramp.
//
// The water plane is shaded with a port of the SAGE water effect
// (`Ocean.fx` / OpenSAGE `Water.frag`); the camera-below-water tint stands in
// for `UnderwaterDeferred.fx`. See docs/REVERSE_ENGINEERING.md.
//
// The Vulkan push constants (20 floats) become a constant buffer.

cbuffer TerrainCB : register(b0) {
    float4 cam;     // x=target_x, y=target_y, z=yaw, w=height
    float4 params;  // x=pitch, y=fov, z=water_z, w=has_water
    float4 sun;     // xyz=sun dir, w=ambient
    float4 mapinfo; // x=W, y=H, z=unused, w=z_scale
    float4 misc;    // x=time, y=water kind (0=ocean,1=river), z=cells per repeat, w=aspect
};

Texture2D<float> heightmap : register(t0);
Texture2D<float4> celldata : register(t1);
Texture2DArray<float4> atlas : register(t2);
SamplerState height_smp : register(s0);
SamplerState cell_smp : register(s1);
SamplerState atlas_smp : register(s2);

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

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

struct VSOut {
    float4 pos : SV_Position;
    float2 uv : TEXCOORD0;
};

VSOut VSMain(uint vertex_id : SV_VertexID) {
    float2 p = float2((vertex_id << 1) & 2, vertex_id & 2);
    VSOut o;
    o.uv = p;
    o.pos = float4(p.x * 2.0 - 1.0, 1.0 - p.y * 2.0, 0.0, 1.0);
    return o;
}

float height_at(int2 c) {
    c = clamp(c, int2(0, 0), int2((int) mapinfo.x - 1, (int) mapinfo.y - 1));
    return heightmap.Load(int3(c, 0)) * 65535.0 * mapinfo.w;
}

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

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

// The retail SAGE blend ramp: 0 at one edge of the cell, 1 at the opposite.
// Direction: 1 right, 2 top, 4 top-right, 8 top-left; flag bit 0 flips,
// bit 1 marks a two-sided diagonal.
float blend_factor(uint direction, uint flags, float2 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;
}

// Sample one tile material layer at global cell coordinates. The texture repeats
// every `span` cells with REPEAT addressing, so it never restarts at a cell edge.
float3 sample_layer(uint layer, float wx, float wy) {
    float span = max(misc.z, 1.0);
    uint lw = 0;
    uint lh = 0;
    uint layer_count = 0;
    atlas.GetDimensions(lw, lh, layer_count);
    float l = (float) min(layer, layer_count > 0u ? layer_count - 1u : 0u);
    return atlas.Sample(atlas_smp, float3(float2(wx, wy) / span, l)).rgb;
}

// ---- 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).
float3 water_tex(int layer, float2 uv) {
    uint w = 0;
    uint h = 0;
    uint lc = 0;
    atlas.GetDimensions(w, h, lc);
    float l = (float) clamp(layer, 0, (int) lc - 1);
    return atlas.Sample(atlas_smp, float3(uv, l)).rgb;
}

// 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.
float3 water_normal(float2 world_xy, float time) {
    float2 q = world_xy * (WATER_SCALE * 6.0);
    float a1 = q.x * 0.80 + q.y * 0.30 + time * 1.10;
    float a2 = q.y * 0.95 - q.x * 0.45 - time * 1.30;
    float a3 = (q.x + q.y) * 1.60 + time * 2.10;
    float a4 = (q.x - q.y) * 2.30 - time * 1.70;
    float dx = 0.224 * cos(a1) - 0.099 * cos(a2) + 0.256 * cos(a3) + 0.230 * cos(a4);
    float dy = 0.084 * cos(a1) + 0.209 * cos(a2) + 0.256 * cos(a3) - 0.230 * cos(a4);
    uint w = 0;
    uint h = 0;
    uint lc = 0;
    atlas.GetDimensions(w, h, lc);
    float3 flow = water_tex((int) lc - 2, q - float2(time * 0.010, time * 0.014)) * 2.0 - 1.0;
    float3 bump = water_tex((int) lc - 1, q + flow.xy * 0.05 + float2(time * 0.006, time * 0.008)) * 2.0 - 1.0;
    float sx = -dx * 0.30 + bump.x * 0.45;
    float sy = -dy * 0.30 + bump.y * 0.45;
    return normalize(float3(sx, sy, 1.0));
}

// Water.frag distortionPower * the flow texture: a small scrolling UV offset.
float water_distortion(float2 world_xy, float time) {
    float2 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 low-frequency
// scroll stands in.
float3 water_cloud(float2 world_xy, float time) {
    return float3(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. Applied to
// every output while the camera is below the water plane.
float3 apply_underwater(float3 color, float dist, float cam_z, float water_z) {
    if (cam_z >= water_z - 0.5) return color;
    const float3 absorb = float3(0.35, 0.62, 0.75);
    float fog = clamp(1.0 - exp(-dist * 0.00022), 0.0, 0.9);
    return lerp(color * absorb, float3(0.02, 0.10, 0.16), fog);
}

// Water.frag: fresnel-mixed reflection (sky) and refraction (depth-graded
// seabed), SAGE diffuse + specular lighting, cloud term, and a depth-based
// transparency fade. `river` (>0.5) applies RiverTransparencyMultiplier.
float3 water_shade(float3 hitpos, float3 dir, float dist) {
    float time = misc.x;
    float river = misc.y;
    float seabed = world_height(hitpos.x, hitpos.y);
    float depth = max(0.0, params.z - seabed);

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

    // Schlick fresnel, water F0 = 0.02.
    float cos_theta = clamp(dot(-dir, n), 0.0, 1.0);
    float fresnel = 0.02 + 0.98 * pow(1.0 - cos_theta, 5.0);

    // Reflection: the sky the surface mirrors. Refraction: the seabed, graded
    // from shallow to deep and lit by the SAGE diffuse + specular model.
    float3 reflection = sky_color(reflect(dir, n));
    float3 shallow = float3(0.10, 0.34, 0.38);
    float3 deep = float3(0.02, 0.12, 0.22);
    float3 refraction = lerp(shallow, deep, clamp(depth / 40.0, 0.0, 1.0));

    float ndotl = max(dot(n, sun_dir), 0.0);
    float ambient = sun.w;
    float3 diffuse = float3(ambient + (1.0 - ambient) * ndotl);
    float3 half_v = normalize(sun_dir - dir);
    float spec = pow(max(dot(n, half_v), 0.0), 90.0);

    float3 color = lerp(refraction, reflection, clamp(fresnel, 0.0, 1.0)) * diffuse * water_cloud(hitpos.xy, time);
    color += float3(1.0, 0.97, 0.9) * spec * 0.45;

    // Depth-based transparency: shallow water shows the seabed, deep water goes
    // opaque toward the deep colour.
    float alpha = clamp(depth / WATER_TRANSPARENT_DEPTH, 0.0, 1.0) * WATER_MIN_OPACITY;
    if (river > 0.5) alpha *= WATER_RIVER_MULTIPLIER;
    color = lerp(refraction, color, clamp(alpha + 0.15, 0.0, 1.0));

    // Distance haze toward the horizon, as the terrain.
    float wfog = clamp(1.0 - exp(-dist * 0.00009), 0.0, 0.75);
    return lerp(color, sky_color(float3(dir.x, dir.y, 0.0)), wfog);
}

float4 PSMain(VSOut input) : SV_Target {
    float4 p = cam;
    float pitch = clamp(params.x, 0.15, 1.45);
    float fov = clamp(params.y, 0.3, 1.4);
    float world_w = mapinfo.x * CELL;
    float world_h = mapinfo.y * CELL;

    float cp = cos(pitch);
    float3 fwd = float3(cp * sin(p.z), cp * cos(p.z), -sin(pitch));
    float3 right = normalize(cross(fwd, float3(0, 0, 1)));
    float3 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);
    float3 cam_pos = float3(p.x, p.y, target_z + p.w) - fwd * dist;

    float2 ndc = float2(input.uv.x * 2.0 - 1.0, 1.0 - input.uv.y * 2.0);
    float aspect = misc.w;
    float th = tan(fov * 0.5);
    float3 dir = normalize(fwd + right * ndc.x * th * aspect + up * ndc.y * th);

    if (dir.z >= -1e-4) {
        return float4(apply_underwater(sky_color(dir), 0.0, cam_pos.z, params.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.
    float t_enter = 0.0;
    float t_exit = 1.0e30;
    bool inside = true;
    if (abs(dir.x) < 1e-6) {
        inside = (cam_pos.x >= 0.0 && cam_pos.x <= world_w);
    } else {
        float a = (0.0 - cam_pos.x) / dir.x;
        float b = (world_w - cam_pos.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_pos.y >= 0.0 && cam_pos.y <= world_h);
        } else {
            float a = (0.0 - cam_pos.y) / dir.y;
            float b = (world_h - cam_pos.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 float4(apply_underwater(sky_color(dir), 0.0, cam_pos.z, params.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.
    float horiz = max(abs(dir.x), abs(dir.y));
    float cell_step = min(CELL / max(horiz, 1e-4), CELL * 32.0);
    float t = max(t_enter, CELL * 0.5);
    float prev = t;
    bool hit = false;
    float hit_t = 0.0;
    for (int i = 0; i < 1024 && t <= t_exit; ++i) {
        float3 w = cam_pos + dir * t;
        float h = world_height(w.x, w.y);
        float surface = (params.w > 0.5) ? max(h, params.z) : h;
        if (w.z <= surface) {
            hit = true;
            hit_t = t;
            break;
        }
        float clearance = (w.z - surface) / max(-dir.z, 1e-4);
        prev = t;
        t += clamp(clearance, cell_step, cell_step * 8.0);
    }
    if (!hit) {
        return float4(apply_underwater(sky_color(dir), 0.0, cam_pos.z, params.z), 1.0);
    }

    // Refine the first crossing; with a sub-cell bracket this converges to the
    // exact surface point.
    float lo = prev;
    float hi = hit_t;
    for (int i = 0; i < 18; ++i) {
        float mid = 0.5 * (lo + hi);
        float3 w = cam_pos + dir * mid;
        float h = world_height(w.x, w.y);
        float surface = (params.w > 0.5) ? max(h, params.z) : h;
        if (w.z <= surface) {
            hi = mid;
        } else {
            lo = mid;
        }
    }
    float3 hitpos = cam_pos + dir * hi;

    float3 sun_dir = normalize(sun.xyz);
    float ambient = sun.w;

    if (params.w > 0.5 && hitpos.z <= params.z + 0.01) {
        return float4(apply_underwater(water_shade(hitpos, dir, hi), hi, cam_pos.z, params.z), 1.0);
    }

    // Terrain: read the per-cell blend record, sample the base/blend/three-way
    // material layers continuously and ramp between them across the cell.
    float wx = hitpos.x / CELL;
    float wy = (world_h - hitpos.y) / CELL;
    int cx = clamp((int) wx, 0, (int) mapinfo.x - 1);
    int cy = clamp((int) wy, 0, (int) mapinfo.y - 1);
    float fx = wx - floor(wx);
    float fy = wy - floor(wy);

    uint4 record = (uint4) (celldata.Load(int3(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;

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

    // Per-pixel normal from the heightfield.
    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);
    float3 n = normalize(float3(hl - hr, hd - hu, 2.0 * CELL));
    float lambert = max(0.0, dot(n, sun_dir));
    float3 lit = albedo * (ambient + (1.0 - ambient) * lambert);
    // Distance haze toward the horizon so the map edge blends into the sky.
    float fog = clamp(1.0 - exp(-hi * 0.00009), 0.0, 0.75);
    lit = lerp(lit, sky_color(float3(dir.x, dir.y, 0.0)), fog);
    return float4(apply_underwater(lit, hi, cam_pos.z, params.z), 1.0);
}