/**
 * File: cloth.shader.ts
 * Description: WGSL shaders for cloth simulation and rendering.
 */

// --- SHARED DATA STRUCTURES ---
export const CLOTH_SHARED_STRUCTS = `
    struct Params {
        dt: f32,
        gravity_y: f32,
        compliance: f32,
        numVertices: f32,
        wind_x: f32,
        wind_y: f32,
        wind_z: f32,
        time: f32,
        elongation: f32
    };
`;

// ==========================================
// VERTEX SHADER
// ==========================================
export const CLOTH_VERTEX_SHADER_WGSL = `
    attribute uv : vec2<f32>;
    var<storage, read> positions : array<vec4<f32>>;

    uniform viewProjection : mat4x4<f32>;

    varying vUV : vec2<f32>;
    varying vWorldPos : vec3<f32>;

    @vertex
    fn main(input : VertexInputs) -> FragmentInputs {
        var output : FragmentInputs;

        let worldPos = positions[input.vertexIndex].xyz;
        output.position = uniforms.viewProjection * vec4<f32>(worldPos, 1.0);

        output.vUV = input.uv;
        output.vWorldPos = worldPos;

        return output;
    }
`;

// ==========================================
// FRAGMENT SHADER
// ==========================================
export const CLOTH_FRAGMENT_SHADER_WGSL = `
    varying vUV : vec2<f32>;
    varying vWorldPos : vec3<f32>;

    @fragment
    fn main(input: FragmentInputs) -> FragmentOutputs {
        var output: FragmentOutputs;

        let dx = dpdx(input.vWorldPos);
        let dy = dpdy(input.vWorldPos);
        let normal = normalize(cross(dx, dy));
        let lightDir = normalize(vec3<f32>(1.0, 1.0, 0.5));
        let diffuse = max(0.0, abs(dot(normal, lightDir)));
        let ambient = 0.3;
        let lightIntensity = ambient + (diffuse * 0.7);
        let grid = (floor(input.vUV.x * 20.0) + floor(input.vUV.y * 20.0)) % 2.0;
        let baseColor = mix(vec3<f32>(0.8, 0.4, 0.15), vec3<f32>(0.9, 0.5, 0.2), grid);
        let finalColor = baseColor * lightIntensity;

        output.color = vec4<f32>(finalColor, 1.0);

        return output;
    }
`;

// =====================================================================
// PASS 1: INTEGRATION (Apply Forces & Predict Positions)
// =====================================================================
export const CLOTH_INTEGRATE_COMPUTE_WGSL = CLOTH_SHARED_STRUCTS + `
    @group(0) @binding(0) var<storage, read> p : Params;
    @group(0) @binding(1) var<storage, read_write> positions : array<vec4<f32>>;
    @group(0) @binding(2) var<storage, read_write> prev_positions : array<vec4<f32>>;
    @group(0) @binding(3) var<storage, read_write> velocities : array<vec4<f32>>;

    @compute @workgroup_size(64)
    fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
        let idx = global_id.x;
        if (f32(idx) >= p.numVertices) { return; }

        var pos = positions[idx];
        var vel = velocities[idx];
        let invMass = pos.w;

        if (invMass > 0.0) {
            vel.y = vel.y + (p.gravity_y * p.dt);

            let flutter = sin(pos.x * 2.0 + p.time * 5.0) * cos(pos.y * 2.0 + p.time * 3.0);

            let windForce = vec3<f32>(
                p.wind_x + (flutter * p.wind_x * 0.8),
                p.wind_y + (flutter * 2.0), // Leichter Auftrieb durchs Flattern
                p.wind_z + (flutter * p.wind_z * 0.8)
            );

            vel.x = vel.x + (windForce.x * p.dt);
            vel.y = vel.y + (windForce.y * p.dt);
            vel.z = vel.z + (windForce.z * p.dt);

            prev_positions[idx] = pos;

            pos.x = pos.x + vel.x * p.dt;
            pos.y = pos.y + vel.y * p.dt;
            pos.z = pos.z + vel.z * p.dt;

            positions[idx] = pos;
            velocities[idx] = vel;
        }
    }
`;

// =====================================================================
// PASS 2: SOLVE CONSTRAINTS (The core of XPBD)
// =====================================================================
export const CLOTH_SOLVE_COMPUTE_WGSL = CLOTH_SHARED_STRUCTS + `
    @group(0) @binding(0) var<storage, read> p : Params;
    @group(0) @binding(1) var<storage, read_write> positions : array<vec4<f32>>;
    @group(0) @binding(2) var<storage, read> constraints : array<vec4<f32>>; // <--- Read-only as we do not modify them here

    @compute @workgroup_size(64)
    fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
        let idx = global_id.x;

        // Query the GPU directly for the length of the passed array
        if (idx >= arrayLength(&constraints)) { return; }

        let constraint = constraints[idx];

        // constraint.w: 0.0 = inactive, 1.0 = horizontal/diagonal, 2.0 = vertical
        if (constraint.w < 0.5) { return; }

        let idA = u32(constraint.x);
        let idB = u32(constraint.y);

        // constraint.w encodes type: 1.0 = horizontal/diagonal, 2.0 = vertical (elongation applies)
        let restLength =constraint.z * p.elongation;

        var pA = positions[idA];
        var pB = positions[idB];

        let wA = pA.w;
        let wB = pB.w;
        let wSum = wA + wB;

        if (wSum <= 0.0) { return; }

        let dir = pA.xyz - pB.xyz;
        let dist = length(dir);

        if (dist < 0.0001) { return; }

        let n = dir / dist;
        let C = dist - restLength;

        let alpha = p.compliance / (p.dt * p.dt);
        let lambda = -C / (wSum + alpha);

        let corrA = n * (lambda * wA);
        let corrB = n * (-lambda * wB);

        // This is because we are using graph coloring to be thread safe
        if (wA > 0.0) {
            positions[idA].x = positions[idA].x + corrA.x;
            positions[idA].y = positions[idA].y + corrA.y;
            positions[idA].z = positions[idA].z + corrA.z;
        }
        if (wB > 0.0) {
            positions[idB].x = positions[idB].x + corrB.x;
            positions[idB].y = positions[idB].y + corrB.y;
            positions[idB].z = positions[idB].z + corrB.z;
        }
    }
`;

// =====================================================================
// PASS 3: VELOCITY UPDATE (Derive velocity from position changes)
// =====================================================================
export const CLOTH_VELOCITY_COMPUTE_WGSL = CLOTH_SHARED_STRUCTS + `
    @group(0) @binding(0) var<storage, read> p : Params;
    @group(0) @binding(1) var<storage, read_write> positions : array<vec4<f32>>;
    @group(0) @binding(2) var<storage, read_write> prev_positions : array<vec4<f32>>;
    @group(0) @binding(3) var<storage, read_write> velocities : array<vec4<f32>>;

    @compute @workgroup_size(64)
    fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
        let idx = global_id.x;
        if (f32(idx) >= p.numVertices) { return; }

        let pos = positions[idx];
        let prev = prev_positions[idx];
        let invMass = pos.w;

        if (invMass > 0.0) {
            var vel = velocities[idx];

            // v = (p - p_prev) / dt
            vel.x = (pos.x - prev.x) / p.dt;
            vel.y = (pos.y - prev.y) / p.dt;
            vel.z = (pos.z - prev.z) / p.dt;

            // Optional: Add simple damping here
            // vel = vel * 0.99;

            velocities[idx] = vel;
        }
    }
`;
