#version 330 core

in  vec3 v_world_normal;
in  vec2 v_tex;
in  vec3 v_world_pos;

out vec4 frag_color;

uniform sampler2D tex;
uniform vec3 color;        // base colour factor
uniform vec3 camera_pos;   // camera world position
uniform vec3 light_dir;    // directional light direction (from light towards scene)
uniform vec3 light_color;  // directional light color

// -------------------------------------
// Constants and helpers (ported from atmosphere.cginc)
const float PI = 3.14159265359;
const float PLANET_RADIUS = 6371000.0;
const float ATMOSPHERE_HEIGHT = 100000.0;
const float RAYLEIGH_HEIGHT = (ATMOSPHERE_HEIGHT * 0.08);
const float MIE_HEIGHT = (ATMOSPHERE_HEIGHT * 0.012);
const float EXPOSURE = 20.0;

// Coefficients (1e-6 scaling kept)
const vec3 C_RAYLEIGH = vec3(5.802, 13.558, 33.100) * 1e-6;
const vec3 C_MIE      = vec3(3.996,  3.996,  3.996) * 1e-6;
const vec3 C_OZONE    = vec3(0.650,  1.881,  0.085) * 1e-6;

float saturate(float x) { return clamp(x, 0.0, 1.0); }
vec2  saturate(vec2  x) { return clamp(x, vec2(0.0), vec2(1.0)); }
vec3  saturate(vec3  x) { return clamp(x, vec3(0.0), vec3(1.0)); }

float AtmosphereHeight(vec3 positionWS) {
    return length(positionWS - vec3(0.0, -PLANET_RADIUS, 0.0)) - PLANET_RADIUS;
}
float DensityRayleigh(float h) { return exp(-max(0.0, h / RAYLEIGH_HEIGHT)); }
float DensityMie     (float h) { return exp(-max(0.0, h / MIE_HEIGHT)); }
float DensityOzone   (float h) { return max(0.0, 1.0 - abs(h - 25000.0) / 15000.0); }
vec3  AtmosphereDensity(float h) { return vec3(DensityRayleigh(h), DensityMie(h), DensityOzone(h)); }

// Sphere intersection with atmosphere shell
vec2 SphereIntersection(vec3 rayStart, vec3 rayDir, vec3 sphereCenter, float sphereRadius) {
    vec3 o = rayStart - sphereCenter;
    float a = dot(rayDir, rayDir);
    float b = 2.0 * dot(o, rayDir);
    float c = dot(o, o) - (sphereRadius * sphereRadius);
    float d = b*b - 4.0*a*c;
    if (d < 0.0) return vec2(-1.0);
    d = sqrt(d);
    return vec2(-b - d, -b + d) / (2.0 * a);
}
vec2 AtmosphereIntersection(vec3 rayStart, vec3 rayDir) {
    return SphereIntersection(rayStart, rayDir, vec3(0.0, -PLANET_RADIUS, 0.0), PLANET_RADIUS + ATMOSPHERE_HEIGHT);
}

float PhaseRayleigh(float costh) {
    return 3.0 * (1.0 + costh*costh) / (16.0 * PI);
}
float PhaseMie(float costh, float g) {
    g = min(g, 0.9381);
    float k = 1.55*g - 0.55*g*g*g;
    float kcosth = k*costh;
    return (1.0 - k*k) / ((4.0*PI) * (1.0 - kcosth) * (1.0 - kcosth));
}

vec3 IntegrateOpticalDepth(vec3 rayStart, vec3 rayDir) {
    vec2 intersection = AtmosphereIntersection(rayStart, rayDir);
    float rayLength = intersection.y;
    int sampleCount = 8;
    float stepSize = rayLength / float(sampleCount);
    vec3 opticalDepth = vec3(0.0);
    for (int i = 0; i < sampleCount; ++i) {
        vec3 localPosition = rayStart + rayDir * (float(i) + 0.5) * stepSize;
        float localHeight  = AtmosphereHeight(localPosition);
        vec3 localDensity  = AtmosphereDensity(localHeight);
        opticalDepth += localDensity * stepSize;
    }
    return opticalDepth;
}

vec3 Absorb(vec3 opticalDepth) {
    // Mie absorbs slightly more than it scatters (~10%)
    return exp(-(opticalDepth.x * C_RAYLEIGH + opticalDepth.y * C_MIE * 1.1 + opticalDepth.z * C_OZONE));
}

vec3 IntegrateScattering(vec3 rayStart, vec3 rayDir, float rayLength, vec3 lightDir, vec3 lightColor, out vec3 transmittance) {
    float rayHeight = AtmosphereHeight(rayStart);
    float sampleDistributionExponent = 1.0 + saturate(1.0 - rayHeight / ATMOSPHERE_HEIGHT) * 8.0;

    vec2 intersection = AtmosphereIntersection(rayStart, rayDir);
    rayLength = min(rayLength, intersection.y);
    if (intersection.x > 0.0) {
        rayStart += rayDir * intersection.x;
        rayLength -= intersection.x;
    }

    float costh = dot(rayDir, lightDir);
    float phaseR = PhaseRayleigh(costh);
    float phaseM = PhaseMie(costh, 0.85);

    int sampleCount = 64;

    vec3 opticalDepth = vec3(0.0);
    vec3 rayleigh = vec3(0.0);
    vec3 mie = vec3(0.0);
    float prevRayTime = 0.0;

    for (int i = 0; i < sampleCount; ++i) {
        float t = pow(float(i) / float(sampleCount), sampleDistributionExponent) * rayLength;
        float stepSize = (t - prevRayTime);
        vec3 localPosition = rayStart + rayDir * t;
        float localHeight = AtmosphereHeight(localPosition);
        vec3 localDensity = AtmosphereDensity(localHeight);
        opticalDepth += localDensity * stepSize;

        vec3 viewTransmittance = Absorb(opticalDepth);
        vec3 opticalDepthLight = IntegrateOpticalDepth(localPosition, lightDir);
        vec3 lightTransmittance = Absorb(opticalDepthLight);

        rayleigh += viewTransmittance * lightTransmittance * phaseR * localDensity.x * stepSize;
        mie      += viewTransmittance * lightTransmittance * phaseM * localDensity.y * stepSize;
        prevRayTime = t;
    }

    transmittance = Absorb(opticalDepth);
    return (rayleigh * C_RAYLEIGH + mie * C_MIE) * lightColor * EXPOSURE;
}

void main() {
    // Base albedo
    vec3 base_col = texture(tex, v_tex).rgb * color;

    // Lighting setup
    vec3 L = normalize(light_dir);
    vec3 N = normalize(v_world_normal);

    // Directional light transmittance (planet shadow)
    vec3 lightTransmittance = Absorb(IntegrateOpticalDepth(v_world_pos, L));

    // Rough ambient by sampling sky upwards
    vec3 tmp;
    vec3 ambient = IntegrateScattering(v_world_pos, vec3(0.0, 1.0, 0.0), 1.0/0.0, L, light_color, tmp);

    // Lambert + atmospheric directional lighting
    float NdotL = max(0.0, dot(N, L));
    vec3 lit = base_col * NdotL * (ambient + light_color * lightTransmittance);

    // View-ray scattering and transmittance between surface and camera
    vec3 V = camera_pos - v_world_pos;
    float rayLength = length(V);
    V = V / max(rayLength, 1e-6);
    vec3 transmittance;
    vec3 scattering = IntegrateScattering(camera_pos, -V, rayLength, L, light_color, transmittance);

    vec3 result = lit * transmittance + scattering;
    frag_color = vec4(result, 1.0);
}