import { useFrame, useLoader } from "@react-three/fiber";
import { useLayoutEffect, useMemo, useRef } from "react";
import { BufferAttribute, BufferGeometry, Color, DoubleSide, InstancedMesh, MeshStandardMaterial, Object3D, TextureLoader } from "three";
import { getTerrainHeight, Shader } from "./helpers";

import grassVert from './shaders/grass.vert';
import grassFrag from './shaders/grass.frag';

interface GrassProps {
	x: number;
	y: number;
	size: number;
	count: number;
	grassSize: number;
	scale: number;
	hillScale: number;
	hillHeight: number;
	detailScale: number;
	detailHeight: number;
	noise2D: (x: number, y: number) => number;
	grassLOD: number;
	dryColor: string;
	lushColor: string;
	grassBlades?: number;
	grassSegments?: number;
	grassLODStart?: number;
	grassLODExponent?: number;
}

export default function Grass({
	x,
	y,
	size,
	count,
	grassSize,
	scale,
	hillScale,
	hillHeight,
	detailScale,
	detailHeight,
	noise2D,
	grassLOD,
	dryColor = '#556b19',
	lushColor = '#348a34',
	grassBlades = 3,
	grassSegments = 4,
	grassLODStart = 0.5,
	grassLODExponent = 1.0
}: GrassProps) {
	const meshRef = useRef<InstancedMesh>(null);
	const dummyRef = useRef<Object3D>(new Object3D());

	const [alphaMap, normalMap] = useLoader(TextureLoader, [
		'niko/img/grass_alpha.png',
		'niko/img/grass_normal.png'
	]);

	const materialRef = useRef<
		MeshStandardMaterial & { userData: { shader: Shader } }
	>(null);

	useFrame((state) => {
		if (
			materialRef.current &&
			materialRef.current.userData &&
			materialRef.current.userData.shader
		) {
			(materialRef.current.userData.shader as Shader).uniforms.uTime.value =
				state.clock.getElapsedTime();
		}
	});

	const geometry = useMemo(() => {
		const geo = new BufferGeometry();

		const w = 0.5;
		const h = 2;
		const segments = grassSegments;
		const bladesCount = grassBlades;

		const positions: number[] = [];
		const uvs: number[] = [];
		const indices: number[] = [];

		for (let i = 0; i < bladesCount; i++) {
			const angle = (Math.PI / bladesCount) * i;
			const sinA = Math.sin(angle);
			const cosA = Math.cos(angle);

			const offset = positions.length / 3;

			for (let y = 0; y <= segments; y++) {
				const v = y / segments;
				const yPos = v * h;

				positions.push(-w * cosA, yPos, -w * sinA);
				uvs.push(0, v);

				positions.push(w * cosA, yPos, w * sinA);
				uvs.push(1, v);
			}

			for (let y = 0; y < segments; y++) {
				const row1 = offset + y * 2;
				const row2 = offset + (y + 1) * 2;

				indices.push(row1, row1 + 1, row2);
				indices.push(row1 + 1, row2 + 1, row2);
			}
		}

		geo.setAttribute(
			'position',
			new BufferAttribute(new Float32Array(positions), 3)
		);
		geo.setAttribute('uv', new BufferAttribute(new Float32Array(uvs), 2));
		geo.setIndex(indices);
		geo.computeVertexNormals();

		return geo;
	}, [grassBlades, grassSegments]);

	useLayoutEffect(() => {
		if (!meshRef.current) return;
		const dummy = dummyRef.current;

		const worldXBase = x * size;
		const worldZBase = y * size;

		const color = new Color();
		const lushColorObj = new Color(lushColor);
		const dryColorObj = new Color(dryColor);

		let instanceIndex = 0;

		for (let i = 0; i < count; i++) {
			const localX = (Math.random() - 0.5) * size;
			const localZ = (Math.random() - 0.5) * size;

			const globalX = worldXBase + localX;
			const globalZ = worldZBase + localZ;

			const dist = Math.sqrt(globalX * globalX + globalZ * globalZ);

			const maxDist = grassLOD;

			const falloffStart = maxDist * grassLODStart;
			const falloffEnd = maxDist;

			let fallofFactor = (dist - falloffStart) / (falloffEnd - falloffStart);
			fallofFactor = Math.max(0, Math.min(1, fallofFactor));

			const density = Math.pow(1.0 - fallofFactor, grassLODExponent);

			if (Math.random() > density) continue;

			const localY = getTerrainHeight(
				localX,
				localZ,
				worldXBase,
				worldZBase,
				scale,
				hillScale,
				hillHeight,
				detailScale,
				detailHeight,
				noise2D
			);

			dummy.position.set(localX, localY, localZ);

			dummy.rotation.y = Math.random() * Math.PI * 2;
			dummy.rotation.x = (Math.random() - 0.5) * 0.2;
			dummy.rotation.z = (Math.random() - 0.5) * 0.2;

			const noiseVal = noise2D(globalX * 0.02, globalZ * 0.02);
			const t = (noiseVal + 1) / 2;
			const randomInternal = (Math.random() - 0.5) * 0.2;
			const finalT = Math.max(0, Math.min(1, t + randomInternal));
			color.lerpColors(dryColorObj, lushColorObj, finalT);
			meshRef.current.setColorAt(instanceIndex, color);

			const heightNoise = noise2D(globalX * 0.08, globalZ * 0.08);
			const macroHeight = (heightNoise + 1.0) * 0.5; // 0..1
			const microNoise = noise2D(globalX * 0.3, globalZ * 0.3);
			const microHeight = (microNoise + 1.0) * 0.25; // 0..0.5
			const perBladeRandom = Math.random() * 0.4;

			const grassWidth = grassSize * (0.7 + Math.random() * 0.5);
			const grassHeight = grassSize * (0.4 + macroHeight * 0.8 + microHeight + perBladeRandom);

			dummy.scale.set(grassWidth, grassHeight, grassWidth);

			dummy.updateMatrix();
			meshRef.current.setMatrixAt(instanceIndex, dummy.matrix);

			instanceIndex++;
		}
		meshRef.current.count = instanceIndex;

		meshRef.current.instanceMatrix.needsUpdate = true;
		if (meshRef.current.instanceColor)
			meshRef.current.instanceColor.needsUpdate = true;
	}, [
		x,
		y,
		size,
		count,
		grassSize,
		scale,
		hillScale,
		hillHeight,
		detailScale,
		detailHeight,
		noise2D,
		grassLOD,
		dryColor,
		lushColor,
		grassLODStart,
		grassLODExponent
	]);

	const onBeforeCompile = useMemo(
		() => (shader: Shader) => {
			shader.uniforms.uTime = { value: 0 };

			shader.vertexShader = `
            uniform float uTime;
            varying vec2 vGrassUv;
            varying vec3 vWorldPos;
            
            float hash(vec2 p) {
                return fract(sin(dot(p, vec2(127.1, 311.7))) * 43758.5453);
            }
            
            float noise(vec2 p) {
                vec2 i = floor(p);
                vec2 f = fract(p);
                f = f * f * (3.0 - 2.0 * f);
                float a = hash(i);
                float b = hash(i + vec2(1.0, 0.0));
                float c = hash(i + vec2(0.0, 1.0));
                float d = hash(i + vec2(1.0, 1.0));
                return mix(mix(a, b, f.x), mix(c, d, f.x), f.y);
            }
            
            float fbm(vec2 p) {
                float value = 0.0;
                float amplitude = 0.5;
                float frequency = 1.0;
                for(int i = 0; i < 4; i++) {
                    value += amplitude * noise(p * frequency);
                    frequency *= 2.0;
                    amplitude *= 0.5;
                }
                return value;
            }
            
            ${shader.vertexShader}
        `;

			shader.vertexShader = shader.vertexShader.replace(
				'#include <begin_vertex>',
				`
            #include <begin_vertex>
            ${grassVert}
            `
			);

			shader.fragmentShader = `
            uniform float uTime;
            varying vec2 vGrassUv;
            varying vec3 vWorldPos;
            ${shader.fragmentShader}
        `;

			shader.fragmentShader = shader.fragmentShader.replace(
				'#include <color_fragment>',
				`
            #include <color_fragment>
            ${grassFrag}
            `
			);

			if (materialRef.current) {
				materialRef.current.userData.shader = shader;
			}
		},
		[]
	);

	return (
		<instancedMesh
			ref={meshRef}
			args={[undefined, undefined, count]}
			position={[x * size, 0, y * size]}
			geometry={geometry}
		>
			<meshStandardMaterial
				ref={materialRef}
				color='#ffffff'
				side={DoubleSide}
				alphaMap={alphaMap}
				alphaTest={0.5}
				normalMap={normalMap}
				roughness={0.8}
				metalness={0.1}
				onBeforeCompile={onBeforeCompile}
			/>
		</instancedMesh>
	);
}