#import bevy_sprite::mesh2d_vertex_output::VertexOutput
#import bevy_sprite::mesh2d_view_bindings::{globals, view}

// Single-scattering atmosphere (Rayleigh + Mie + ozone) with the sun fixed
// just above the horizon, plus a thin layer of drifting clouds.

const PI: f32 = 3.14159265;

const PLANET_RADIUS: f32 = 6371000.0;
const ATMOSPHERE_RADIUS: f32 = 6471000.0;
const RAYLEIGH_HEIGHT: f32 = 8000.0;
const MIE_HEIGHT: f32 = 1200.0;
const CAMERA_ALTITUDE: f32 = 200.0;

const C_RAYLEIGH: vec3<f32> = vec3<f32>(5.802, 13.558, 33.100) * 1e-6;
const C_MIE: vec3<f32> = vec3<f32>(3.996, 3.996, 3.996) * 1e-6;
const C_OZONE: vec3<f32> = vec3<f32>(0.650, 1.881, 0.085) * 1e-6;
const MIE_G: f32 = 0.76;

const VIEW_SAMPLES: i32 = 16;
const LIGHT_SAMPLES: i32 = 4;

const SUN_INTENSITY: f32 = 20.0;
const SUN_ANGULAR_RADIUS: f32 = 0.012;
// Vertical field of view is 2 * atan(FOV_SCALE)
const FOV_SCALE: f32 = 0.6;
// Screen height (in NDC) of the horizon line
const HORIZON_NDC: f32 = -0.55;

fn sun_direction() -> vec3<f32> {
    // Slightly right of center, ~2 degrees above the horizon
    return normalize(vec3<f32>(0.25, 0.035, -1.0));
}

// Near and far distances to a sphere centered at the planet center, or -1
// when the ray misses.
fn sphere_hit(origin: vec3<f32>, dir: vec3<f32>, radius: f32) -> vec2<f32> {
    let b = dot(origin, dir);
    let c = dot(origin, origin) - radius * radius;
    let d = b * b - c;
    if d < 0.0 {
        return vec2<f32>(-1.0);
    }
    let s = sqrt(d);
    return vec2<f32>(-b - s, -b + s);
}

fn density(p: vec3<f32>) -> vec3<f32> {
    let h = max(length(p) - PLANET_RADIUS, 0.0);
    let ozone = max(0.0, 1.0 - abs(h - 25000.0) / 15000.0);
    return vec3<f32>(exp(-h / RAYLEIGH_HEIGHT), exp(-h / MIE_HEIGHT), ozone);
}

fn absorb(optical_depth: vec3<f32>) -> vec3<f32> {
    return exp(-(optical_depth.x * C_RAYLEIGH
        + optical_depth.y * C_MIE * 1.1
        + optical_depth.z * C_OZONE));
}

fn light_optical_depth(p: vec3<f32>, dir: vec3<f32>) -> vec3<f32> {
    let len = sphere_hit(p, dir, ATMOSPHERE_RADIUS).y;
    let step = len / f32(LIGHT_SAMPLES);
    var depth = vec3<f32>(0.0);
    for (var i = 0; i < LIGHT_SAMPLES; i++) {
        depth += density(p + dir * (f32(i) + 0.5) * step) * step;
    }
    return depth;
}

fn phase_rayleigh(c: f32) -> f32 {
    return 3.0 * (1.0 + c * c) / (16.0 * PI);
}

fn phase_hg(c: f32, g: f32) -> f32 {
    let g2 = g * g;
    return (1.0 - g2) / (4.0 * PI * pow(1.0 + g2 - 2.0 * g * c, 1.5));
}

// In-scattered light along the view ray; `transmittance` receives the
// attenuation of whatever lies behind it.
fn scatter(origin: vec3<f32>, dir: vec3<f32>, sun: vec3<f32>, transmittance: ptr<function, vec3<f32>>) -> vec3<f32> {
    var len = sphere_hit(origin, dir, ATMOSPHERE_RADIUS).y;
    let ground = sphere_hit(origin, dir, PLANET_RADIUS);
    if ground.x > 0.0 {
        len = ground.x;
    }

    let c = dot(dir, sun);
    let phase_r = phase_rayleigh(c);
    let phase_m = phase_hg(c, MIE_G);

    var depth = vec3<f32>(0.0);
    var rayleigh = vec3<f32>(0.0);
    var mie = vec3<f32>(0.0);
    var prev_t = 0.0;
    for (var i = 1; i <= VIEW_SAMPLES; i++) {
        // Quadratic distribution: dense samples near the camera, where the air is thick
        let f = f32(i) / f32(VIEW_SAMPLES);
        let t = f * f * len;
        let step = t - prev_t;
        let p = origin + dir * (t - 0.5 * step);
        let local = density(p);
        depth += local * step;

        let light = absorb(depth + light_optical_depth(p, sun));
        rayleigh += light * local.x * step;
        mie += light * local.y * step;
        prev_t = t;
    }

    *transmittance = absorb(depth);
    return (rayleigh * C_RAYLEIGH * phase_r + mie * C_MIE * phase_m) * SUN_INTENSITY;
}

fn hash(p: vec2<f32>) -> f32 {
    return fract(sin(dot(p, vec2<f32>(127.1, 311.7))) * 43758.5453);
}

fn noise(p: vec2<f32>) -> f32 {
    let i = floor(p);
    let f = fract(p);
    let u = f * f * (3.0 - 2.0 * f);
    return mix(
        mix(hash(i), hash(i + vec2<f32>(1.0, 0.0)), u.x),
        mix(hash(i + vec2<f32>(0.0, 1.0)), hash(i + vec2<f32>(1.0, 1.0)), u.x),
        u.y,
    );
}

fn fbm(p: vec2<f32>) -> f32 {
    var value = 0.0;
    var amplitude = 0.5;
    var q = p;
    for (var i = 0; i < 5; i++) {
        value += amplitude * noise(q);
        q *= 2.0;
        amplitude *= 0.5;
    }
    return value;
}

@fragment
fn fragment(mesh: VertexOutput) -> @location(0) vec4<f32> {
    let aspect = view.viewport.z / view.viewport.w;
    // uv.y is 0 at the top of the quad
    let ndc = vec2<f32>(mesh.uv.x * 2.0 - 1.0, 1.0 - mesh.uv.y * 2.0);
    let dir = normalize(vec3<f32>(
        ndc.x * aspect * FOV_SCALE,
        (ndc.y - HORIZON_NDC) * FOV_SCALE,
        -1.0,
    ));

    let sun = sun_direction();
    let origin = vec3<f32>(0.0, PLANET_RADIUS + CAMERA_ALTITUDE, 0.0);

    var transmittance: vec3<f32>;
    var color = scatter(origin, dir, sun, &transmittance);

    let cos_sun = dot(dir, sun);
    if dir.y > 0.0 {
        // Sun disc, reddened by the air in front of it
        let disc = smoothstep(cos(SUN_ANGULAR_RADIUS * 1.2), cos(SUN_ANGULAR_RADIUS * 0.8), cos_sun);
        color += disc * transmittance * SUN_INTENSITY * 20.0;

        // Clouds on a flat plane, lit by the sunlight that reaches them
        let p = dir.xz / dir.y * 2.0;
        let t = globals.time;
        let clouds = fbm(p + vec2<f32>(t * 0.02, 0.0)) * 0.6
            + fbm(p * 2.0 + vec2<f32>(t * 0.035, t * 0.008)) * 0.4;
        let cover = smoothstep(0.5, 0.8, clouds) * smoothstep(0.0, 0.15, dir.y);
        let sunlight = absorb(light_optical_depth(origin, sun)) * SUN_INTENSITY;
        let lit = sunlight * (0.04 + phase_hg(cos_sun, 0.6) * 0.3) + color * 0.5;
        color = mix(color, lit, cover * 0.85);
    }

    // Exposure tonemap, then dither to hide banding in the gradients
    color = 1.0 - exp(-color * 2.0);
    color += (hash(mesh.position.xy) - 0.5) / 255.0;

    return vec4<f32>(color, 1.0);
}