const PI: f32 = 3.141592653589793;

struct Uniforms {
    view: mat4x4<f32>,
    projection: mat4x4<f32>,
    view_projection: mat4x4<f32>,
    camera_position: vec4<f32>,
    time: f32,
    _pad0: f32,
    _pad1: f32,
    _pad2: f32,
}

struct WaterMaterial {
    base_color: vec4<f32>,
    water_color: vec4<f32>,
    wave_height: f32,
    choppy: f32,
    speed: f32,
    freq: f32,
    specular_strength: f32,
    fresnel_power: f32,
    volume_shape: u32,
    volume_flow_type: u32,
    volume_size: vec3<f32>,
    is_volumetric: u32,
}

struct WaterInstance {
    model_0: vec4<f32>,
    model_1: vec4<f32>,
    model_2: vec4<f32>,
    model_3: vec4<f32>,
    material_index: u32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
}

struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) tex_coords: vec2<f32>,
}

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) world_pos: vec3<f32>,
    @location(1) normal: vec3<f32>,
    @location(2) @interpolate(flat) material_index: u32,
    @location(3) tex_coords: vec2<f32>,
}

@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage, read> materials: array<WaterMaterial>;
@group(0) @binding(2) var<storage, read> instances: array<WaterInstance>;
@group(0) @binding(3) var<storage, read> visible_indices: array<u32>;

fn hash(p: vec2<f32>) -> f32 {
    var p3 = fract(vec3<f32>(p.x, p.y, p.x) * 0.1031);
    p3 = p3 + dot(p3, p3.yzx + 33.33);
    return fract((p3.x + p3.y) * p3.z);
}

fn noise(p: vec2<f32>) -> f32 {
    let i = floor(p);
    let f = fract(p);
    let u = f * f * (3.0 - 2.0 * f);
    return -1.0 + 2.0 * mix(
        mix(hash(i + vec2<f32>(0.0, 0.0)), 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_vertical(uv: vec2<f32>, time: f32) -> f32 {
    var value = 0.0;
    var amplitude = 0.5;
    var frequency = 1.0;

    for (var index = 0; index < 5; index = index + 1) {
        let stretched_uv = vec2<f32>(uv.x * frequency * 2.0, uv.y * frequency * 0.5 + time);
        value = value + amplitude * noise(stretched_uv);
        amplitude = amplitude * 0.5;
        frequency = frequency * 2.0;
    }
    return value;
}

fn water_height_vertical(uv_in: vec2<f32>, wave_height: f32, choppy: f32, speed: f32, freq: f32, time: f32) -> f32 {
    let flow_speed = speed * 2.0;
    let flow_time = time * flow_speed;

    let uv_scaled = uv_in * freq * 8.0;

    let layer1 = fbm_vertical(uv_scaled, flow_time);
    let layer2 = fbm_vertical(uv_scaled * 1.7 + vec2<f32>(3.7, 0.0), flow_time * 1.3) * 0.5;
    let layer3 = fbm_vertical(uv_scaled * 3.1 + vec2<f32>(7.3, 0.0), flow_time * 1.7) * 0.25;

    let combined = layer1 + layer2 + layer3;

    let turbulence = abs(noise(uv_scaled * 4.0 + vec2<f32>(0.0, flow_time * 2.0))) * choppy * 0.1;

    return combined * wave_height + turbulence;
}

fn water_normal_vertical(tex_coords: vec2<f32>, base_normal: vec3<f32>, mat: WaterMaterial, epsilon: f32) -> vec3<f32> {
    let h_center = water_height_vertical(tex_coords, mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time);
    let h_u = water_height_vertical(tex_coords + vec2<f32>(epsilon, 0.0), mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time);
    let h_v = water_height_vertical(tex_coords + vec2<f32>(0.0, epsilon), mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time);

    let du = h_u - h_center;
    let dv = h_v - h_center;

    let tangent = normalize(vec3<f32>(1.0, 0.0, 0.0) - base_normal * dot(vec3<f32>(1.0, 0.0, 0.0), base_normal));
    let bitangent = normalize(cross(base_normal, tangent));

    let perturbed = base_normal + tangent * (-du * 2.0) + bitangent * (-dv * 2.0);
    return normalize(perturbed);
}

@vertex
fn vs_main(
    vertex: VertexInput,
    @builtin(instance_index) instance_index: u32
) -> VertexOutput {
    var out: VertexOutput;

    let visible_idx = visible_indices[instance_index];
    let instance = instances[visible_idx];

    let model = mat4x4<f32>(
        instance.model_0,
        instance.model_1,
        instance.model_2,
        instance.model_3
    );

    var world_pos = (model * vec4<f32>(vertex.position, 1.0)).xyz;

    let mat = materials[instance.material_index];

    let normal_matrix = mat3x3<f32>(
        model[0].xyz,
        model[1].xyz,
        model[2].xyz
    );
    let world_normal = normalize(normal_matrix * vertex.normal);

    let wave_h = water_height_vertical(
        vertex.tex_coords,
        mat.wave_height,
        mat.choppy,
        mat.speed,
        mat.freq,
        uniforms.time
    );
    world_pos = world_pos + world_normal * wave_h;

    out.world_pos = world_pos;
    out.clip_position = uniforms.view_projection * vec4<f32>(world_pos, 1.0);
    out.normal = world_normal;
    out.material_index = instance.material_index;
    out.tex_coords = vertex.tex_coords;

    return out;
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let mat = materials[in.material_index];
    let camera_pos = uniforms.camera_position.xyz;
    let V = normalize(camera_pos - in.world_pos);

    var sun_direction = normalize(vec3<f32>(0.5, 0.8, 0.3));
    let sun_color = vec3<f32>(1.0, 0.95, 0.85);

    let N = water_normal_vertical(in.tex_coords, in.normal, mat, 0.01);

    let flow_time = uniforms.time * mat.speed * 2.0;
    let uv_scaled = in.tex_coords * mat.freq * 8.0;

    let foam_noise1 = noise(uv_scaled * 3.0 + vec2<f32>(0.0, flow_time * 2.0));
    let foam_noise2 = noise(uv_scaled * 7.0 + vec2<f32>(1.3, flow_time * 3.0));
    let foam_noise3 = noise(uv_scaled * 15.0 + vec2<f32>(2.7, flow_time * 4.0));
    let foam = clamp((foam_noise1 * 0.5 + foam_noise2 * 0.3 + foam_noise3 * 0.2) * 0.5 + 0.5, 0.0, 1.0);

    let streak_noise = noise(vec2<f32>(in.tex_coords.x * 20.0, in.tex_coords.y * 2.0 + flow_time));
    let streaks = clamp(streak_noise * 0.5 + 0.5, 0.0, 1.0);

    let foam_color = vec3<f32>(0.9, 0.95, 1.0);
    let deep_color = mat.base_color.rgb;

    var water_color = mix(deep_color, foam_color, foam * 0.7 + streaks * 0.3);

    let fresnel = pow(clamp(1.0 - dot(N, V), 0.0, 1.0), mat.fresnel_power);
    water_color = water_color + vec3<f32>(0.1, 0.15, 0.2) * fresnel;

    let H = normalize(sun_direction + V);
    let spec = pow(max(dot(N, H), 0.0), 64.0) * mat.specular_strength * 0.5;
    water_color = water_color + sun_color * spec;

    let diffuse = max(dot(N, sun_direction), 0.0) * 0.3;
    water_color = water_color + sun_color * diffuse;

    return vec4<f32>(water_color, 0.85);
}
