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 baseScale = grassSize + Math.random() * grassSize * 0.5;
			const heightMult = 0.5 + Math.random() * 1.0;
			dummy.scale.set(baseScale, baseScale * heightMult, baseScale);

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

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

			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;
			
			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>
			vGrassUv = uv;
			${grassVert}
			`
			);

			shader.fragmentShader = `
			uniform float uTime;
			varying vec2 vGrassUv;
			${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>
	);
}