import * as THREE from 'three';

// Scene setup
const scene = new THREE.Scene();
const camera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0, 1); // Orthographic for 2D
const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setSize(window.innerWidth, window.innerHeight);
document.body.appendChild(renderer.domElement);

// Simulation parameters
const resolution = 512; // Texture resolution
const damping = 0.995; // Slightly stronger damping for smoother waves
const waveSpeed = 0.55; // Faster waves for dynamic feel

// Frame buffers for height and normal maps
const heightTarget1 = new THREE.WebGLRenderTarget(resolution, resolution, {
    type: THREE.FloatType,
    minFilter: THREE.NearestFilter,
    magFilter: THREE.NearestFilter,
});
const heightTarget2 = new THREE.WebGLRenderTarget(resolution, resolution, {
    type: THREE.FloatType,
    minFilter: THREE.NearestFilter,
    magFilter: THREE.NearestFilter,
});
const normalTarget = new THREE.WebGLRenderTarget(resolution, resolution, {
    minFilter: THREE.LinearFilter,
    magFilter: THREE.LinearFilter,
});

// Simulation uniforms
const simUniforms = {
    u_time: { value: 0.0 },
    u_resolution: { value: new THREE.Vector2(resolution, resolution) },
    u_mouse: { value: new THREE.Vector2(-1, -1) },
    u_mouseForce: { value: 0.0 },
    u_prevHeight: { value: heightTarget1.texture },
    u_damping: { value: damping },
    u_waveSpeed: { value: waveSpeed },
};

// Simulation vertex shader
const simVertexShader = `
    varying vec2 vUv;
    void main() {
        vUv = uv;
        gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
    }
`;

// Simulation fragment shader (wave equation with background motion)
const simFragmentShader = `
    uniform float u_time;
    uniform vec2 u_resolution;
    uniform vec2 u_mouse;
    uniform float u_mouseForce;
    uniform sampler2D u_prevHeight;
    uniform float u_damping;
    uniform float u_waveSpeed;
    varying vec2 vUv;

    // Noise for randomness and background motion
    float random(vec2 st) {
        return fract(sin(dot(st, vec2(127.1, 311.7))) * 43758.5453123);
    }

    float noise(vec2 st) {
        vec2 i = floor(st);
        vec2 f = fract(st);
        float a = random(i);
        float b = random(i + vec2(1.0, 0.0));
        float c = random(i + vec2(0.0, 1.0));
        float d = random(i + vec2(1.0, 1.0));
        vec2 u = f * f * (3.0 - 2.0 * f);
        return mix(a, b, u.x) + (c - a) * u.y * (1.0 - u.x) + (d - b) * u.x * u.y;
    }

    void main() {
        vec2 texel = 1.0 / u_resolution;
        float h = texture2D(u_prevHeight, vUv).r; // Current height
        float h_prev = texture2D(u_prevHeight, vUv).g; // Previous height

        // Sample neighbors
        float h_n = texture2D(u_prevHeight, vUv + vec2(0.0, texel.y)).r;
        float h_s = texture2D(u_prevHeight, vUv - vec2(0.0, texel.y)).r;
        float h_e = texture2D(u_prevHeight, vUv + vec2(texel.x, 0.0)).r;
        float h_w = texture2D(u_prevHeight, vUv - vec2(texel.x, 0.0)).r;

        // Wave equation
        float c2 = u_waveSpeed * u_waveSpeed;
        float newHeight = 2.0 * h - h_prev + c2 * (h_n + h_s + h_e + h_w - 4.0 * h);
        newHeight *= u_damping;

        // Mouse force with randomness
        float dist = length(vUv - u_mouse);
        if (dist < 0.05 && u_mouseForce > 0.0) {
            float rand = noise(vUv + u_time) * 0.1 + 0.9; // Random strength variation
            newHeight += 0.15 * rand * u_mouseForce * exp(-dist * 100.0);
        }

        // Background motion (subtle waves)
        float t = u_time * 2.0;
        float background = noise(vUv * 2.0 + t) * 0.00002;
        background += sin(u_time * 0.5 + vUv.x * 5.0) * 0.0000002; // Periodic pulse
        newHeight += background;

        gl_FragColor = vec4(newHeight, h, 0.0, 1.0);
    }
`;

// Simulation material
const simMaterial = new THREE.ShaderMaterial({
    uniforms: simUniforms,
    vertexShader: simVertexShader,
    fragmentShader: simFragmentShader,
});

// Normal map shader (computes normals from height field)
const normalVertexShader = `
    varying vec2 vUv;
    void main() {
        vUv = uv;
        gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
    }
`;

const normalFragmentShader = `
    uniform sampler2D u_heightMap;
    uniform vec2 u_resolution;
    varying vec2 vUv;

    void main() {
        vec2 texel = 1.0 / u_resolution;
        float h = texture2D(u_heightMap, vUv).r;
        float h_n = texture2D(u_heightMap, vUv + vec2(0.0, texel.y)).r;
        float h_s = texture2D(u_heightMap, vUv - vec2(0.0, texel.y)).r;
        float h_e = texture2D(u_heightMap, vUv + vec2(texel.x, 0.0)).r;
        float h_w = texture2D(u_heightMap, vUv - vec2(texel.x, 0.0)).r;

        // Compute normal using height gradients
        vec3 normal = normalize(vec3(h_w - h_e, h_s - h_n, 0.05));
        gl_FragColor = vec4(normal * 0.5 + 0.5, 1.0); // Pack to [0,1]
    }
`;

const normalMaterial = new THREE.ShaderMaterial({
    uniforms: {
        u_heightMap: { value: heightTarget1.texture },
        u_resolution: { value: new THREE.Vector2(resolution, resolution) },
    },
    vertexShader: normalVertexShader,
    fragmentShader: normalFragmentShader,
});

// Simulation and normal scenes
const simGeometry = new THREE.PlaneGeometry(2, 2);
const simMesh = new THREE.Mesh(simGeometry, simMaterial);
const simScene = new THREE.Scene();
simScene.add(simMesh);

const normalMesh = new THREE.Mesh(simGeometry, normalMaterial);
const normalScene = new THREE.Scene();
normalScene.add(normalMesh);

// Rendering uniforms
const renderUniforms = {
    u_heightMap: { value: heightTarget1.texture },
    u_normalMap: { value: normalTarget.texture },
    u_resolution: { value: new THREE.Vector2(window.innerWidth, window.innerHeight) },
    u_time: { value: 0.0 },
    u_lightPos: { value: new THREE.Vector2(0.5, 0.5) }, // Dynamic light position
};

// Rendering vertex shader
const renderVertexShader = `
    varying vec2 vUv;
    void main() {
        vUv = uv;
        gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
    }
`;

// Rendering fragment shader (with lighting and glowing effects)
const renderFragmentShader = `
    uniform sampler2D u_heightMap;
    uniform sampler2D u_normalMap;
    uniform vec2 u_resolution;
    uniform float u_time;
    uniform vec2 u_lightPos;
    varying vec2 vUv;

    // Noise for granularity
    float random(vec2 st) {
        return fract(sin(dot(st, vec2(127.1, 311.7))) * 43758.5453123);
    }

    float noise(vec2 st) {
        vec2 i = floor(st);
        vec2 f = fract(st);
        float a = random(i);
        float b = random(i + vec2(1.0, 0.0));
        float c = random(i + vec2(0.0, 1.0));
        float d = random(i + vec2(1.0, 1.0));
        vec2 u = f * f * (3.0 - 2.0 * f);
        return mix(a, b, u.x) + (c - a) * u.y * (1.0 - u.x) + (d - b) * u.x * u.y;
    }

    void main() {
        float height = texture2D(u_heightMap, vUv).r;
        vec3 normal = texture2D(u_normalMap, vUv).xyz * 2.0 - 1.0;
        vec2 p = vUv * u_resolution.xy / min(u_resolution.x, u_resolution.y);

        // Dynamic lighting
        vec3 lightDir = normalize(vec3(u_lightPos - vUv, 0.5));
        float diffuse = max(dot(normal, lightDir), 0.0);
        float specular = pow(diffuse, 32.0) * 0.5;

        // Base color with gradient
        vec3 baseColor = mix(vec3(0.0, 0.2, 0.5), vec3(0.1, 0.7, 1.0), 0.5 + height * 3.0);
        
        // Add granularity
        float n = noise(p * 15.0 + u_time * 0.2) * 0.1;
        baseColor += vec3(n);

        // Glowing wave crests
        float glow = abs(height) > 0.02 ? sin(height * 100.0 + u_time * 2.0) * 0.3 : 0.0;
        vec3 glowColor = vec3(0.2, 0.8, 1.0) * glow;
        vec3 glowColorB = vec3(0.2, 0.8, 1.0);

        // Combine lighting and effects
        vec3 color = baseColor * (0.5 + diffuse) + vec3(specular) + glowColor;
        color = clamp(color, 0.0, 1.0);

        gl_FragColor = vec4(color, 1);
    }
`;

// Rendering material
const renderMaterial = new THREE.ShaderMaterial({
    uniforms: renderUniforms,
    vertexShader: renderVertexShader,
    fragmentShader: renderFragmentShader,
});

// Rendering plane
const renderGeometry = new THREE.PlaneGeometry(2, 2);
const renderMesh = new THREE.Mesh(renderGeometry, renderMaterial);
scene.add(renderMesh);

// Handle window resize
window.addEventListener('resize', () => {
    renderer.setSize(window.innerWidth, window.innerHeight);
    renderUniforms.u_resolution.value.set(window.innerWidth, window.innerHeight);
});

// Mouse interaction (click and drag)
const mouse = new THREE.Vector2();
let isMouseDown = false;
window.addEventListener('mousedown', (event) => {
    isMouseDown = true;
    updateMouse(event);
});
window.addEventListener('mousemove', (event) => {
    if (isMouseDown) updateMouse(event);
});
window.addEventListener('mouseup', () => {
    isMouseDown = false;
    simUniforms.u_mouseForce.value = 0.0;
});

function updateMouse(event) {
    mouse.x = event.clientX / window.innerWidth;
    mouse.y = 1.0 - event.clientY / window.innerHeight;
    simUniforms.u_mouse.value.set(mouse.x, mouse.y);
    simUniforms.u_mouseForce.value = 1.0;
}

// Animation loop
const clock = new THREE.Clock();
let currentHeightTarget = heightTarget1;
let nextHeightTarget = heightTarget2;
function animate() {
    requestAnimationFrame(animate);
    const t = clock.getElapsedTime();
    simUniforms.u_time.value = t;
    renderUniforms.u_time.value = t;

    // Update light position (follows mouse or oscillates)
    renderUniforms.u_lightPos.value.set(
        isMouseDown ? mouse.x : 0.5 + Math.sin(t * 0.5) * 0.3,
        isMouseDown ? mouse.y : 0.5 + Math.cos(t * 0.5) * 0.3
    );

    // Update simulation
    simUniforms.u_prevHeight.value = currentHeightTarget.texture;
    renderer.setRenderTarget(nextHeightTarget);
    renderer.render(simScene, camera);

    // Update normal map
    normalMaterial.uniforms.u_heightMap.value = nextHeightTarget.texture;
    renderer.setRenderTarget(normalTarget);
    renderer.render(normalScene, camera);

    // Render to screen
    renderUniforms.u_heightMap.value = nextHeightTarget.texture;
    renderer.setRenderTarget(null);
    renderer.render(scene, camera);

    // Swap buffers
    const temp = currentHeightTarget;
    currentHeightTarget = nextHeightTarget;
    nextHeightTarget = temp;
}
animate();