// WGSL Shaders for WebGPU

// Vertex shader for noise generation (fullscreen quad)
export const noiseVertexShader = `
@vertex
fn main(@builtin(vertex_index) vertexIndex: u32) -> @builtin(position) vec4<f32> {
    var pos = array<vec2<f32>, 6>(
        vec2<f32>(-1.0, -1.0),
        vec2<f32>(1.0, -1.0),
        vec2<f32>(-1.0, 1.0),
        vec2<f32>(1.0, -1.0),
        vec2<f32>(1.0, 1.0),
        vec2<f32>(-1.0, 1.0)
    );
    return vec4<f32>(pos[vertexIndex], 0.0, 1.0);
}
`;

// Fragment shader for Perlin noise
export const noiseFragmentShader = `
@group(0) @binding(0) var<uniform> uTime: f32;

fn permute(x: vec4<f32>) -> vec4<f32> {
    return ((x * 34.0 + 1.0) * x) % vec4<f32>(289.0);
}

fn taylorInvSqrt(r: vec4<f32>) -> vec4<f32> {
    return 1.79284291400159 - 0.85373472095314 * r;
}

fn fade(t: vec3<f32>) -> vec3<f32> {
    return t * t * t * (t * (t * 6.0 - 15.0) + 10.0);
}

fn cnoise(P: vec3<f32>) -> f32 {
    var Pi0: vec3<f32> = floor(P);
    var Pi1: vec3<f32> = Pi0 + vec3<f32>(1.0);
    Pi0 = Pi0 % vec3<f32>(289.0);
    Pi1 = Pi1 % vec3<f32>(289.0);
    let Pf0 = fract(P);
    let Pf1 = Pf0 - vec3<f32>(1.0);
    let ix = vec4<f32>(Pi0.x, Pi1.x, Pi0.x, Pi1.x);
    let iy = vec4<f32>(Pi0.yy, Pi1.yy);
    let iz0 = Pi0.zzzz;
    let iz1 = Pi1.zzzz;

    let ixy = permute(permute(ix) + iy);
    let ixy0 = permute(ixy + iz0);
    let ixy1 = permute(ixy + iz1);

    var gx0: vec4<f32> = ixy0 / 7.0;
    var gy0: vec4<f32> = fract(floor(gx0) / 7.0) - 0.5;
    gx0 = fract(gx0);
    let gz0 = vec4<f32>(0.5) - abs(gx0) - abs(gy0);
    let sz0 = step(gz0, vec4<f32>(0.0));
    gx0 = gx0 - sz0 * (step(vec4<f32>(0.0), gx0) - 0.5);
    gy0 = gy0 - sz0 * (step(vec4<f32>(0.0), gy0) - 0.5);

    var gx1: vec4<f32> = ixy1 / 7.0;
    var gy1: vec4<f32> = fract(floor(gx1) / 7.0) - 0.5;
    gx1 = fract(gx1);
    let gz1 = vec4<f32>(0.5) - abs(gx1) - abs(gy1);
    let sz1 = step(gz1, vec4<f32>(0.0));
    gx1 = gx1 - sz1 * (step(vec4<f32>(0.0), gx1) - 0.5);
    gy1 = gy1 - sz1 * (step(vec4<f32>(0.0), gy1) - 0.5);

    var g000: vec3<f32> = vec3<f32>(gx0.x, gy0.x, gz0.x);
    var g100: vec3<f32> = vec3<f32>(gx0.y, gy0.y, gz0.y);
    var g010: vec3<f32> = vec3<f32>(gx0.z, gy0.z, gz0.z);
    var g110: vec3<f32> = vec3<f32>(gx0.w, gy0.w, gz0.w);
    var g001: vec3<f32> = vec3<f32>(gx1.x, gy1.x, gz1.x);
    var g101: vec3<f32> = vec3<f32>(gx1.y, gy1.y, gz1.y);
    var g011: vec3<f32> = vec3<f32>(gx1.z, gy1.z, gz1.z);
    var g111: vec3<f32> = vec3<f32>(gx1.w, gy1.w, gz1.w);

    let norm0 = taylorInvSqrt(vec4<f32>(dot(g000, g000), dot(g100, g100), dot(g010, g010), dot(g110, g110)));
    g000 = g000 * norm0.x;
    g100 = g100 * norm0.y;
    g010 = g010 * norm0.z;
    g110 = g110 * norm0.w;
    let norm1 = taylorInvSqrt(vec4<f32>(dot(g001, g001), dot(g101, g101), dot(g011, g011), dot(g111, g111)));
    g001 = g001 * norm1.x;
    g101 = g101 * norm1.y;
    g011 = g011 * norm1.z;
    g111 = g111 * norm1.w;

    let n000 = dot(g000, Pf0);
    let n100 = dot(g100, vec3<f32>(Pf1.x, Pf0.yz));
    let n010 = dot(g010, vec3<f32>(Pf0.x, Pf1.y, Pf0.z));
    let n110 = dot(g110, vec3<f32>(Pf1.xy, Pf0.z));
    let n001 = dot(g001, vec3<f32>(Pf0.xy, Pf1.z));
    let n101 = dot(g101, vec3<f32>(Pf1.x, Pf0.y, Pf1.z));
    let n011 = dot(g011, vec3<f32>(Pf0.x, Pf1.yz));
    let n111 = dot(g111, Pf1);

    let fade_xyz = fade(Pf0);
    let n_z = mix(vec4<f32>(n000, n100, n010, n110), vec4<f32>(n001, n101, n011, n111), fade_xyz.z);
    let n_yz = mix(n_z.xy, n_z.zw, fade_xyz.y);
    let n_xyz = mix(n_yz.x, n_yz.y, fade_xyz.x);
    return 2.2 * n_xyz;
}

@fragment
fn main(@builtin(position) fragCoord: vec4<f32>) -> @location(0) vec4<f32> {
    let resolution = vec2<f32>(256.0, 256.0);
    let uv = fragCoord.xy / resolution;
    var sum = 0.0;
    var scale = 1.0;
    var amplitude = 0.5;
    
    for (var i = 0; i < 5; i = i + 1) {
        // Make noise tileable by wrapping coordinates
        let wrapped_uv = fract(uv * scale);
        let p = vec3<f32>(wrapped_uv * 10.0, uTime * 0.2);
        sum += cnoise(p) * amplitude;
        scale *= 2.0;
        amplitude *= 0.5;
    }
    
    return vec4<f32>(sum, 0.0, 0.0, 1.0);
}
`;

// Ocean vertex shader
export const oceanVertexShader = `
struct Uniforms {
    view: mat4x4<f32>,
    model: mat4x4<f32>,
    projection: mat4x4<f32>,
    eyePos: vec3<f32>,
};

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var displacementTexture: texture_2d<f32>;
@group(0) @binding(2) var displacementSampler: sampler;

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) fragPos: vec3<f32>,
    @location(1) uv: vec2<f32>,
};

@vertex
fn main(
    @location(0) position: vec3<f32>,
    @location(1) uv: vec2<f32>
) -> VertexOutput {
    var output: VertexOutput;
    var worldPos = uniforms.model * vec4<f32>(position, 1.0);
    
    output.uv = uv;
    
    // Sample displacement using textureSampleLevel (works in vertex shader)
    let displace = textureSampleLevel(displacementTexture, displacementSampler, uv, 0.0);
    worldPos.z = worldPos.z + displace.r * 0.15;
    
    output.position = uniforms.projection * uniforms.view * worldPos;
    output.fragPos = worldPos.xyz;
    
    return output;
}
`;

// Ocean fragment shader
export const oceanFragmentShader = `
struct Uniforms {
    view: mat4x4<f32>,
    model: mat4x4<f32>,
    projection: mat4x4<f32>,
    eyePos: vec3<f32>,
};

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var displacementTexture: texture_2d<f32>;
@group(0) @binding(2) var displacementSampler: sampler;

@fragment
fn main(
    @location(0) fragPos: vec3<f32>,
    @location(1) uv: vec2<f32>
) -> @location(0) vec4<f32> {
    // Sample displacement for normal calculation only
    let gridPointDelta = 1.0 / 256.0;
    let displacementScale = 0.15;
    
    let displace = textureSample(displacementTexture, displacementSampler, uv).r * displacementScale;
    let right = textureSample(displacementTexture, displacementSampler, vec2<f32>(uv.x + gridPointDelta, uv.y)).r * displacementScale;
    let left = textureSample(displacementTexture, displacementSampler, vec2<f32>(uv.x - gridPointDelta, uv.y)).r * displacementScale;
    let up = textureSample(displacementTexture, displacementSampler, vec2<f32>(uv.x, uv.y + gridPointDelta)).r * displacementScale;
    let down = textureSample(displacementTexture, displacementSampler, vec2<f32>(uv.x, uv.y - gridPointDelta)).r * displacementScale;
    
    // Calculate surface normal
    let dX = vec3<f32>(gridPointDelta * 2.0, 0.0, right - left);
    let dY = vec3<f32>(0.0, gridPointDelta * 2.0, up - down);
    var norm = normalize(cross(dX, dY));
    
    // Lighting
    let lightDir = normalize(vec3<f32>(0.3, 0.5, 0.8));
    let diff = max(dot(norm, lightDir), 0.0);
    let diffuse = diff * vec3<f32>(0.8, 0.9, 1.0);
    
    // Fresnel
    let toCameraVector = normalize(fragPos - uniforms.eyePos);
    let reflec = normalize(reflect(toCameraVector, norm));
    let n1 = 1.0;
    let n2 = 1.33333;
    let R0 = pow((n1 - n2) / (n1 + n2), 2.0);
    let fresnel = R0 + (1.0 - R0) * pow((1.0 - dot(norm, reflec)), 5.0);
    
    let oceanColor = vec3<f32>(0.0, 0.25, 0.35);
    let skyColor = vec3<f32>(0.4, 0.6, 0.8);
    
    // Subsurface scattering
    let sssSun = vec3<f32>(0.0, -5.0, -7.0);
    let tosssSunVec = normalize(sssSun - fragPos);
    let tosssSun = normalize(vec3<f32>(0.0, -100.0, 1.0));
    let ssDistortion = 0.1;
    let sssIntensity = 1.0;
    let halfWay = normalize(tosssSun + norm * ssDistortion);
    let ssScateringCoef = pow(clamp(dot(toCameraVector, -halfWay), 0.0, 1.0), 5.0) * sssIntensity;
    
    // Sun glittering
    var glitterFactor = max(0.0, dot(tosssSunVec, reflect(-toCameraVector, norm)));
    if (glitterFactor <= 0.98) {
        glitterFactor = 0.0;
    }
    
    let lightColor = vec3<f32>(1.0, 1.0, 1.0);
    let ambientColor = vec3<f32>(0.1, 0.15, 0.2);
    let finalColor = ambientColor + 
                     diffuse * 0.4 + 
                     mix(oceanColor * (1.0 + ssScateringCoef), skyColor * 0.5, fresnel * 0.7) + 
                     lightColor * glitterFactor * 0.5;
    
    return vec4<f32>(clamp(finalColor, vec3<f32>(0.0), vec3<f32>(1.0)), 1.0);
}
`;

// Skybox vertex shader
export const skyboxVertexShader = `
struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    sunDirection: vec3<f32>,
};

@group(0) @binding(0) var<uniform> uniforms: Uniforms;

struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) rayDir: vec3<f32>,
};

@vertex
fn main(@location(0) position: vec3<f32>) -> VertexOutput {
    var output: VertexOutput;
    output.rayDir = position;
    
    // Remove translation from view matrix
    var rotView = uniforms.view;
    rotView[3] = vec4<f32>(0.0, 0.0, 0.0, 1.0);
    
    let pos = uniforms.projection * rotView * vec4<f32>(position, 1.0);
    output.position = pos;
    
    return output;
}
`;

// Skybox fragment shader
export const skyboxFragmentShader = `
struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    sunDirection: vec3<f32>,
};

@group(0) @binding(0) var<uniform> uniforms: Uniforms;

@fragment
fn main(@location(0) rayDir: vec3<f32>) -> @location(0) vec4<f32> {
    let ray = normalize(rayDir);
    let upAmount = ray.z;
    
    // Sky gradient
    let horizonBlend = pow(1.0 - max(upAmount, 0.0), 2.0);
    let zenithColor = vec3<f32>(0.15, 0.35, 0.75);
    let horizonColor = vec3<f32>(0.55, 0.7, 0.9);
    var skyColor = mix(zenithColor, horizonColor, horizonBlend);
    
    // Horizon glow
    let horizonGlow = pow(max(1.0 - abs(upAmount), 0.0), 6.0);
    skyColor = skyColor + vec3<f32>(0.4, 0.25, 0.1) * horizonGlow * 0.4;
    
    // Sun
    let sunDir = normalize(uniforms.sunDirection);
    let sunAngle = max(dot(ray, sunDir), 0.0);
    
    let sunDisk = smoothstep(0.9993, 0.9998, sunAngle);
    let sunColor = vec3<f32>(1.0, 0.95, 0.85);
    let sunGlow = pow(sunAngle, 48.0) * 0.6;
    let sunHalo = pow(sunAngle, 6.0) * 0.25;
    
    skyColor = skyColor + sunColor * sunDisk * 3.0;
    skyColor = skyColor + vec3<f32>(1.0, 0.85, 0.5) * sunGlow;
    skyColor = skyColor + vec3<f32>(1.0, 0.9, 0.7) * sunHalo;
    
    // Below horizon
    if (upAmount < 0.0) {
        let depth = -upAmount;
        let deepColor = vec3<f32>(0.02, 0.08, 0.15);
        skyColor = mix(horizonColor * 0.7, deepColor, smoothstep(0.0, 0.5, depth));
    }
    
    return vec4<f32>(skyColor, 1.0);
}
`;