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,
    flow_direction: vec2<f32>,
    flow_strength: f32,
    _flow_padding: f32,
}

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 sea_octave(uv_in: vec2<f32>, choppy: f32) -> f32 {
    var uv = uv_in + noise(uv_in);
    var wv = 1.0 - abs(sin(uv));
    let swv = abs(cos(uv));
    wv = mix(wv, swv, wv);
    return pow(1.0 - pow(wv.x * wv.y, 0.65), choppy);
}

fn water_height(xz: vec2<f32>, wave_height: f32, choppy: f32, speed: f32, freq: f32, time: f32, flow_dir: vec2<f32>, flow_str: f32) -> f32 {
    var current_freq = freq;
    var amp = wave_height;
    var current_choppy = choppy;
    var uv = vec2<f32>(xz.x * 0.75, xz.y);
    let sea_time = 1.0 + time * speed;
    let octave_m = mat2x2<f32>(1.6, 1.2, -1.2, 1.6);

    let has_flow = flow_str > 0.001;
    let adjusted_flow_dir = vec2<f32>(flow_dir.x * 0.75, flow_dir.y);
    var current_flow_dir = adjusted_flow_dir;

    var h = 0.0;
    for (var index = 0; index < 5; index = index + 1) {
        let flow_offset = current_flow_dir * sea_time;
        let default_offset = vec2<f32>(sea_time, sea_time);
        let anim_offset = select(default_offset, flow_offset, has_flow);

        let d = sea_octave((uv + anim_offset) * current_freq, current_choppy);
        h = h + d * amp;

        uv = octave_m * uv;
        current_flow_dir = octave_m * current_flow_dir;
        current_freq = current_freq * 1.9;
        amp = amp * 0.22;
        current_choppy = current_choppy + (1.0 - current_choppy) * 0.2;
    }
    return h;
}

fn water_normal(world_pos: vec3<f32>, mat: WaterMaterial, epsilon: f32) -> vec3<f32> {
    let h_center = water_height(vec2<f32>(world_pos.x, world_pos.z), mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time, mat.flow_direction, mat.flow_strength);
    let h_x = water_height(vec2<f32>(world_pos.x + epsilon, world_pos.z), mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time, mat.flow_direction, mat.flow_strength);
    let h_z = water_height(vec2<f32>(world_pos.x, world_pos.z + epsilon), mat.wave_height, mat.choppy, mat.speed, mat.freq, uniforms.time, mat.flow_direction, mat.flow_strength);

    let dx = h_x - h_center;
    let dz = h_z - h_center;

    return normalize(vec3<f32>(-dx, epsilon, -dz));
}

@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 wave_h = water_height(
        vec2<f32>(world_pos.x, world_pos.z),
        mat.wave_height,
        mat.choppy,
        mat.speed,
        mat.freq,
        uniforms.time,
        mat.flow_direction,
        mat.flow_strength
    );
    world_pos.y = world_pos.y + wave_h;

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

    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);

    let dist = length(in.world_pos - camera_pos);
    let epsilon = max(0.1, dist * 0.002);

    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(in.world_pos, mat, epsilon);
    let R = reflect(-V, N);

    let fresnel = pow(clamp(1.0 - dot(N, V), 0.0, 1.0), mat.fresnel_power);

    let sky_y = max(R.y, 0.0) * 0.8 + 0.2;
    let sky_color = vec3<f32>(0.4, 0.6, 0.9) * sky_y + vec3<f32>(0.8, 0.85, 0.9) * (1.0 - sky_y);

    var water_color = mix(mat.base_color.rgb, mat.water_color.rgb, fresnel);
    water_color = water_color + sky_color * fresnel * 0.5;

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

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

    let subsurface_color = mat.water_color.rgb * 0.3;
    let scatter = pow(max(dot(V, -sun_direction), 0.0), 2.0) * 0.2;
    water_color = water_color + subsurface_color * scatter;

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