'use client';

import {
	Environment,
	OrbitControls,
	PerspectiveCamera,
	useProgress
} from '@react-three/drei';
import { Canvas, useLoader, useFrame } from '@react-three/fiber';
import {
	forwardRef,
	useImperativeHandle,
	useLayoutEffect,
	useMemo,
	useRef,
	useState
} from 'react';

import {
	BufferAttribute,
	BufferGeometry,
	Mesh,
	Object3D,
	InstancedMesh,
	DoubleSide,
	TextureLoader,
	Color,
	MeshStandardMaterial,
	FogExp2
} from 'three';
import { LUTCubeLoader } from 'three/examples/jsm/loaders/LUTCubeLoader';

import './page.css';
import { createNoise2D } from 'simplex-noise';
import { useControls } from 'leva';

import grassVert from './shaders/grass.vert';
import grassFrag from './shaders/grass.frag';
import {
	Bloom,
	EffectComposer,
	Noise,
	Vignette,
	DepthOfField,
	GodRays,
	SMAA,
	HueSaturation,
	BrightnessContrast,
	LUT
} from '@react-three/postprocessing';

interface Shader {
	uniforms: { [key: string]: { value: unknown } };
	vertexShader: string;
	fragmentShader: string;
}

function getTerrainHeight(
	localX: number,
	localZ: number,
	worldXBase: number,
	worldZBase: number,
	scale: number,
	hillScale: number,
	hillHeight: number,
	detailScale: number,
	detailHeight: number,
	noise2D: (x: number, y: number) => number
) {
	const worldX = (worldXBase + localX) * 0.1;
	const worldZ = (worldZBase + localZ) * 0.1;

	const noiseHill =
		noise2D(worldX * hillScale, worldZ * hillScale) * hillHeight;
	const noiseDetail =
		noise2D(worldX * detailScale, worldZ * detailScale) * detailHeight;

	return (noiseHill + noiseDetail) * scale;
}

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;
	enableShadows?: boolean;
}

function Grass({
	x,
	y,
	size,
	count,
	grassSize,
	scale,
	hillScale,
	hillHeight,
	detailScale,
	detailHeight,
	noise2D,
	grassLOD
}: GrassProps) {
	const meshRef = useRef<InstancedMesh>(null);
	const dummyRef = useRef<Object3D>(new Object3D());

	const [alphaMap, normalMap] = useLoader(TextureLoader, [
		'/img/grass_alpha.png',
		'/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 = 4;

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

		for (let i = 0; i < 3; i++) {
			const angle = (Math.PI / 3) * 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;
	}, []);

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

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

		const color = new Color();
		const lushColor = new Color('#348a34');
		const dryColor = new Color('#556b19');

		const chunkRadius = size * Math.sqrt(2);

		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 * 0.1;
			const falloffEnd = maxDist;

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

			const density = (1.0 - fallofFactor) * (1.0 - fallofFactor);

			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(dryColor, lushColor, 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
	]);

	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;
			}
		},
		[grassLOD]
	);

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

interface TerrainChunkProps {
	x: number;
	y: number;
	size: number;
	resolution: number;
	scale: number;
	hillScale: number;
	hillHeight: number;
	detailScale: number;
	detailHeight: number;
	noise2D: (x: number, y: number) => number;
	wireframe?: boolean;
	grassCount: number;
	grassSize: number;
	grassLOD: number;
}
function TerrainChunk({
	x,
	y,
	size,
	resolution,
	scale,
	hillScale,
	hillHeight,
	detailScale,
	detailHeight,
	noise2D,
	wireframe = false,
	grassCount,
	grassSize,
	grassLOD
}: TerrainChunkProps) {
	const chunkDist = Math.sqrt((x * size) ** 2 + (y * size) ** 2);
	const shouldRenderGrass = chunkDist < grassLOD + size;

	const meshRef = useRef<Mesh>(null);

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

		const vertices: Array<number> = [];
		const indices: Array<number> = [];

		const step = size / (resolution - 1);
		const halfSize = size / 2;
		const worldXBase = x * size;
		const worldZBase = y * size;

		/*
			vtx gen
		*/
		for (let iz = 0; iz < resolution; iz++) {
			for (let ix = 0; ix < resolution; ix++) {
				const localX = ix * step - halfSize;
				const localZ = iz * step - halfSize;

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

				vertices.push(localX, localY, localZ);
			}
		}

		/*
			idx gen
		*/
		for (let iz = 0; iz < resolution - 1; iz++) {
			for (let ix = 0; ix < resolution - 1; ix++) {
				const topLeft = iz * resolution + ix;
				const topRight = topLeft + 1;
				const bottomLeft = (iz + 1) * resolution + ix;
				const bottomRight = bottomLeft + 1;

				indices.push(topLeft, bottomLeft, topRight);
				indices.push(topRight, bottomLeft, bottomRight);
			}
		}
		geo.setAttribute(
			'position',
			new BufferAttribute(new Float32Array(vertices), 3)
		);

		const colors: Array<number> = [];
		for (let iz = 0; iz < resolution; iz++) {
			for (let ix = 0; ix < resolution; ix++) {
				const localX = ix * step - halfSize;
				const localZ = iz * step - halfSize;
				const globalX = worldXBase + localX;
				const globalZ = worldZBase + localZ;
				const colorNoise = noise2D(globalX * 0.02, globalZ * 0.02);
				const t = (colorNoise + 1) / 2;

				const dryGreen = { r: 0.01, g: 0.01, b: 0.0 };
				const lushGreen = { r: 0.0, g: 0.05, b: 0.0 };

				const r = dryGreen.r + (lushGreen.r - dryGreen.r) * t;
				const g = dryGreen.g + (lushGreen.g - dryGreen.g) * t;
				const b = dryGreen.b + (lushGreen.b - dryGreen.b) * t;

				colors.push(r, g, b);
			}
		}
		geo.setAttribute('color', new BufferAttribute(new Float32Array(colors), 3));

		geo.setIndex(indices);
		geo.computeVertexNormals();

		return geo;
	}, [
		x,
		y,
		size,
		resolution,
		scale,
		hillScale,
		hillHeight,
		detailScale,
		detailHeight,
		noise2D
	]);

	return (
		<group>
			<mesh
				ref={meshRef}
				geometry={geometry}
				position={[x * size, 0, y * size]}
			>
				<meshStandardMaterial
					vertexColors
					roughness={1.0}
					metalness={0.0}
					wireframe={wireframe}
				/>
			</mesh>
			{!wireframe && shouldRenderGrass && (
				<Grass
					x={x}
					y={y}
					size={size}
					count={grassCount}
					grassSize={grassSize}
					scale={scale}
					hillScale={hillScale}
					hillHeight={hillHeight}
					detailScale={detailScale}
					detailHeight={detailHeight}
					noise2D={noise2D}
					grassLOD={grassLOD}
				/>
			)}
		</group>
	);
}

interface TerrainProps {
	chunks?: number;
	chunkSize?: number;
	resolution?: number;
	scale?: number;
	hillScale: number;
	hillHeight: number;
	detailScale: number;
	detailHeight: number;
	wireframe?: boolean;
	grassCount?: number;
	grassSize?: number;
	grassLOD?: number;
}
function Terrain({
	chunks = 5,
	chunkSize = 10,
	resolution = 8,
	scale = 1,
	hillScale = 0.2,
	hillHeight = 3,
	detailScale = 3,
	detailHeight = 0.1,
	wireframe = false,
	grassCount = 6000,
	grassSize = 0.6,
	grassLOD = 40
}: TerrainProps) {
	const noise2D = useMemo(() => createNoise2D(), []);
	const offset = -Math.floor(chunks / 2);

	const chunkPositions = useMemo(() => {
		const positions: [number, number][] = [];
		for (let x = 0; x < chunks; x++)
			for (let y = 0; y < chunks; y++) positions.push([x + offset, y + offset]);

		return positions;
	}, [chunks, offset]);

	return (
		<group>
			{chunkPositions.map(([x, y], index) => (
				<TerrainChunk
					key={`${x}-${y}-${index}`}
					x={x}
					y={y}
					size={chunkSize}
					resolution={resolution}
					scale={scale}
					hillScale={hillScale}
					hillHeight={hillHeight}
					detailScale={detailScale}
					detailHeight={detailHeight}
					noise2D={noise2D}
					wireframe={wireframe}
					grassCount={grassCount}
					grassSize={grassSize}
					grassLOD={grassLOD}
				/>
			))}
		</group>
	);
}

const SealCube = forwardRef<Mesh>((props, ref) => {
	const texture = useLoader(TextureLoader, '/img/niko.jpg');
	const meshRef = useRef<Mesh>(null);

	useImperativeHandle(ref, () => meshRef.current!, []);

	useFrame((state, delta) => {
		if (meshRef.current) {
			meshRef.current.rotation.x += delta * 0.5;
			meshRef.current.rotation.y += delta * 0.5;
		}
	});

	return (
		<mesh ref={meshRef} position={[0, 3, 0]} castShadow receiveShadow>
			<boxGeometry args={[0.85, 0.85, 0.85]} />
			<meshBasicMaterial map={texture} depthWrite={true} />
		</mesh>
	);
});
SealCube.displayName = 'SealCube';

function Loader() {
	const { progress, active } = useProgress();
	const [visible, setVisible] = useState(true);

	useLayoutEffect(() => {
		if (!active && progress === 100) {
			const timeout = setTimeout(() => setVisible(false), 500);
			return () => clearTimeout(timeout);
		}
	}, [progress, active]);

	return (
		<div className={`loader ${!visible ? 'loader hidden' : ''}`}>
			<img src='/img/niko.jpg' className='niko-spin' />
		</div>
	);
}

function LutEffect() {
	const lutTexture = useLoader(LUTCubeLoader, '/lut/Moody3.cube');
	return <LUT lut={lutTexture.texture3D} />;
}

export default function Home() {
	const [sealMesh, setSealMesh] = useState<Mesh | null>(null);

	return (
		<>
			<Loader />
			<Canvas
				shadows
				camera={{ position: [0, 5, 15], fov: 50, far: 100 }}
				gl={{ antialias: true }}
				className='canvas'
			>
				<EffectComposer>
					<DepthOfField target={[0, 3, 0]} focalLength={10} bokehScale={5} />
					<Vignette />
					<Noise opacity={0.05} />
					<Bloom
						intensity={2}
						luminanceThreshold={0.5}
						luminanceSmoothing={0.1}
					/>
					<SMAA />
					<HueSaturation saturation={0.3} />
					<BrightnessContrast brightness={0.1} contrast={-0.1} />
					<LutEffect />
				</EffectComposer>

				<Environment
					files={'hdr/sky.hdr'}
					environmentIntensity={1}
					background
				/>
				<fogExp2 attach='fog' args={[0x9a9a9a, 0.01]} />

				<Terrain
					chunks={16}
					chunkSize={10.0}
					resolution={8.0}
					scale={1}
					hillScale={0.1}
					hillHeight={6.0}
					detailScale={1.0}
					detailHeight={0.2}
					wireframe={false}
					grassCount={6000}
					grassSize={0.6}
					grassLOD={60}
				/>
				<SealCube
					ref={(mesh) => {
						if (mesh && !sealMesh) setSealMesh(mesh);
					}}
				/>

				<OrbitControls
					target={[0, 3, 0]}
					enablePan={false}
					makeDefault
					minDistance={2}
					maxDistance={10}
				/>
			</Canvas>
		</>
	);
}