using Sandbox;
using Sandbox.Rendering;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;

public sealed class GrassRenderObject : SceneCustomObject
{
	public struct GrassData
	{
		public Vector4 Position;
		public Vector2 Uv;
	}

	public struct ChunkCoord : IEquatable<ChunkCoord>
	{
		public int X;
		public int Y;

		public ChunkCoord( int x, int y )
		{
			X = x;
			Y = y;
		}

		public bool Equals( ChunkCoord other )
		{
			return X == other.X && Y == other.Y;
		}

		public override bool Equals( object obj )
		{
			return obj is ChunkCoord other && Equals( other );
		}

		public override int GetHashCode()
		{
			return HashCode.Combine( X, Y );
		}
	}

	public sealed class GrassChunk
	{
		public ChunkCoord Coord;

		public GpuBuffer ArgsBuffer;
		public GpuBuffer ArgsBufferLOD;

		public GpuBuffer<GrassData> PositionsBuffer;

		public Material Material;

		public BBox Bounds;
	}

	// =========================================================

	public readonly Dictionary<ChunkCoord, GrassChunk> ActiveChunks = new();
	public readonly Dictionary<ChunkCoord, GrassChunkDensityMap> DensityMap = new();
	public readonly Dictionary<ChunkCoord, GrassChunkInteractionMap> InteractionMap = new();
	public GrassDensityMapResource DensityResource { get; private set; }
	public CameraComponent CullingCamera { get; set; }

	private readonly Queue<GrassChunk> PendingDestroy = new();

	// =========================================================

	public float ChunkSize { get; set; } = 256.0f;

	public int ChunkResolution { get; set; } = 64;

	public int RenderRadius { get; set; } = 6;

	public float GrassScale { get; set; } = 1.0f;

	public float DisplacementStrength { get; set; } = 200.0f;

	public float LodCutoff { get; set; } = 1000.0f;

	public float DistanceCutoff { get; set; } = 2000.0f;

	// How wide (in world units) the transition/fade is around each cutoff
	public float LodTransitionRange { get; set; } = 200.0f;

	public float DistanceTransitionRange { get; set; } = 200.0f;

	public float TerrainProbeTop { get; set; } = 4096.0f;

	public float TerrainProbeBottom { get; set; } = -4096.0f;

	public float TerrainHeightOffset { get; set; } = 0.0f;

	public float GrassHeightPadding { get; set; } = 128.0f;

	public float FallbackHeight { get; set; } = 0.0f;

	public int HeightProbeResolution { get; set; } = 8;

	public float InteractionDecayRate { get; set; } = 1.5f;
	public float InteractionStrength { get; set; } = 8.0f;
	public float InteractionStampRate { get; set; } = 36.0f;
	public float InteractionDecayUpdateInterval { get; set; } = 0.05f;
	public float InteractionBendHoldDuration { get; set; } = 0.5f;
	public float CutDuration { get; set; } = 8.0f;

	// =========================================================

	private readonly Material GrassMaterial;

	private readonly Model GrassModel;
	private readonly Model GrassModelLOD;
	private float _interactionDecayAccumulator;

	public Texture HeightMap;

	// =========================================================
	// ==== WIND ===============================================
	// =========================================================

	public ComputeShader GenerateWindShader { get; private set; }

	public Texture Wind { get; private set; }
	public float WindSpeed { get; set; } = 65.0f;
	public float WindFrequency { get; set; } = 45.0f;
	public float WindAmplitude { get; set; } = 4f;
	public Vector2 WindDirection { get; set; } = new Vector2( 0, 1 );

	public GrassRenderObject( SceneWorld sceneWorld ) : base( sceneWorld )
	{
		GrassMaterial = Material.Load( "models/grass/materials/grass_blade01.vmat" );

		GrassModel = Model.Load( "models/grass/grass_strand01.vmdl" );
		GrassModelLOD = Model.Load( "models/grass/grass_strand01_lod1.vmdl" );

		HeightMap = Texture.Load( "models/grass/materials/heightmap_temp.vtex" );

		GenerateWindShader = new ComputeShader( "shaders/windnoise.shader" );

		Wind = Texture.CreateRenderTarget()
			.WithSize( 1024, 1024 )
			.WithDepth( 0 )
			.WithFormat( ImageFormat.RGBA8888 )
			.WithDynamicUsage()
			.WithUAVBinding()
			.Create( "texture_grass_wind" );

	}

	// =========================================================
	// MAIN THREAD
	// =========================================================

	public void UpdateStreaming( Vector3 cameraPos )
	{
		HashSet<ChunkCoord> needed = new();

		ChunkCoord camChunk = WorldToChunk( cameraPos );

		for ( int x = -RenderRadius; x <= RenderRadius; x++ )
		{
			for ( int y = -RenderRadius; y <= RenderRadius; y++ )
			{
				ChunkCoord coord = new(
					camChunk.X + x,
					camChunk.Y + y
				);

				if ( !DensityMap.TryGetValue( coord, out GrassChunkDensityMap densityMap ) || densityMap == null || !densityMap.HasAnyDensity() )
					continue;

				needed.Add( coord );

				if ( !ActiveChunks.ContainsKey( coord ) )
				{
					var chunk = CreateChunk( coord );

					InitializeChunkHeightData( chunk );

					ActiveChunks.Add( coord, chunk );
				}
			}
		}

		List<ChunkCoord> remove = new();

		foreach ( var pair in ActiveChunks )
		{
			if ( !needed.Contains( pair.Key ) )
			{
				remove.Add( pair.Key );
			}
		}

		foreach ( var coord in remove )
		{
			PendingDestroy.Enqueue( ActiveChunks[coord] );

			ActiveChunks.Remove( coord );
		}
	}

	public bool TryGetChunkDensity( ChunkCoord coord, out float density )
	{
		if ( DensityMap.TryGetValue( coord, out GrassChunkDensityMap chunkDensity ) && chunkDensity != null )
		{
			density = chunkDensity.HasAnyDensity() ? 1.0f : 0.0f;
			return true;
		}

		density = 0.0f;
		return false;
	}

	public GrassChunkDensityMap GetOrCreateChunkDensityMap( ChunkCoord coord, int resolution )
	{
		if ( !DensityMap.TryGetValue( coord, out GrassChunkDensityMap chunkDensity ) || chunkDensity == null )
		{
			chunkDensity = new GrassChunkDensityMap();
			DensityMap[coord] = chunkDensity;
		}

		chunkDensity.EnsureResolution( resolution );
		return chunkDensity;
	}

	public void SetChunkDensityMap( Vector2Int coord, GrassChunkDensityMap densityMap )
	{
		SetChunkDensityMap( new ChunkCoord( coord.x, coord.y ), densityMap );
	}

	public void SetChunkDensityMap( ChunkCoord coord, GrassChunkDensityMap densityMap )
	{
		if ( densityMap == null )
		{
			DensityMap.Remove( coord );
			return;
		}

		densityMap.EnsureResolution( Math.Max( 1, densityMap.Resolution ) );
		DensityMap[coord] = densityMap;
	}

	public void ReplaceDensityMap( Dictionary<Vector2Int, GrassChunkDensityMap> map )
	{
		DensityMap.Clear();

		foreach ( var pair in map )
		{
			SetChunkDensityMap( pair.Key, pair.Value );
		}
	}

	public GrassChunkInteractionMap GetOrCreateChunkInteractionMap( ChunkCoord coord, int resolution )
	{
		if ( !InteractionMap.TryGetValue( coord, out GrassChunkInteractionMap chunkInteraction ) || chunkInteraction == null )
		{
			chunkInteraction = new GrassChunkInteractionMap();
			InteractionMap[coord] = chunkInteraction;
		}

		chunkInteraction.EnsureResolution( resolution );
		return chunkInteraction;
	}

	public void UpdateInteractionField( IEnumerable<GrassInteractionSourceComponent> sources, float deltaTime )
	{
		if ( deltaTime > 0.0f && InteractionMap.Count > 0 )
		{
			List<ChunkCoord> remove = new();

			foreach ( var pair in InteractionMap )
			{
				if ( pair.Value == null || !pair.Value.DecayAndHasAny( deltaTime, InteractionBendHoldDuration, CutDuration ) )
					remove.Add( pair.Key );
			}

			foreach ( var coord in remove )
			{
				InteractionMap.Remove( coord );
			}
		}

		if ( sources == null )
			return;

		foreach ( var source in sources )
		{
			if ( source == null || !source.Enabled )
				continue;

			Vector3 worldPos = source.WorldPosition;
			float radius = Math.Max( 1.0f, source.Radius );

			// Only stamp interaction when the source is close enough to the terrain
			// (within the source's radius). This prevents airborne sources from
			// affecting the grass.
			float groundHeight = SampleTerrainHeight( worldPos.x, worldPos.y ) + TerrainHeightOffset;
			float verticalDist = MathF.Abs( worldPos.z - groundHeight );
			if ( verticalDist > radius )
				continue;

			// Instant response: stamp strength is no longer frame-time dependent.
			float baseStrength = Math.Clamp( source.Strength, 0.0f, 1.0f ) * 4.0f;
			float rateBoost = Math.Max( 1.0f, InteractionStampRate * 0.25f );
			float strength = baseStrength * rateBoost;

			int minChunkX = MathX.FloorToInt( ( worldPos.x - radius ) / ChunkSize );
			int maxChunkX = MathX.FloorToInt( ( worldPos.x + radius ) / ChunkSize );
			int minChunkY = MathX.FloorToInt( ( worldPos.y - radius ) / ChunkSize );
			int maxChunkY = MathX.FloorToInt( ( worldPos.y + radius ) / ChunkSize );

			for ( int chunkX = minChunkX; chunkX <= maxChunkX; chunkX++ )
			{
				for ( int chunkY = minChunkY; chunkY <= maxChunkY; chunkY++ )
				{
					ChunkCoord chunkCoord = new( chunkX, chunkY );
					GrassChunkInteractionMap interactionMap = GetOrCreateChunkInteractionMap( chunkCoord, ChunkResolution );

					Vector2 chunkOrigin = new( chunkX * ChunkSize, chunkY * ChunkSize );
					Vector2 localCenter = new(
						( worldPos.x - chunkOrigin.x ) / ChunkSize,
						( worldPos.y - chunkOrigin.y ) / ChunkSize
					);

					if ( source.IsCutter )
					{
						interactionMap.StampCircleCut( localCenter, radius / ChunkSize );
					}
					else
					{
						interactionMap.StampCircle( localCenter, radius / ChunkSize, strength );
					}
				}
			}
		}
	}

	public void SetDensityResource( GrassDensityMapResource resource )
	{
		if ( DensityResource == resource )
			return;

		DensityResource = resource;
		ReplaceDensityMap( resource?.ToRuntimeMap() ?? new Dictionary<Vector2Int, GrassChunkDensityMap>() );
	}

	public void SyncDensityResource()
	{
		if ( DensityResource == null )
			return;

		Dictionary<Vector2Int, GrassChunkDensityMap> map = DensityMap.ToDictionary(
			pair => new Vector2Int( pair.Key.X, pair.Key.Y ),
			pair => pair.Value
		);

		DensityResource.FromRuntimeMap( map );
	}

	public void RefreshChunk( ChunkCoord coord )
	{
		if ( !ActiveChunks.TryGetValue( coord, out GrassChunk chunk ) )
			return;

		InitializeChunkHeightData( chunk );
	}

	public void ProcessPendingDestroy()
	{
		while ( PendingDestroy.Count > 0 )
		{
			var chunk = PendingDestroy.Dequeue();

			chunk.PositionsBuffer?.Dispose();
			chunk.ArgsBuffer?.Dispose();
			chunk.ArgsBufferLOD?.Dispose();
		}
	}

	// =========================================================

	private GrassChunk CreateChunk( ChunkCoord coord )
	{
		GrassChunk chunk = new();

		chunk.Coord = coord;

		int instanceCount = ChunkResolution * ChunkResolution;

		chunk.ArgsBuffer = new GpuBuffer(
			1,
			5 * sizeof( uint ),
			GpuBuffer.UsageFlags.ByteAddress
		);

		chunk.ArgsBufferLOD = new GpuBuffer(
			1,
			5 * sizeof( uint ),
			GpuBuffer.UsageFlags.ByteAddress
		);

		chunk.ArgsBuffer.SetData(
			BuildArgs( GrassModel, (uint)instanceCount )
		);

		chunk.ArgsBufferLOD.SetData(
			BuildArgs( GrassModelLOD, (uint)instanceCount )
		);

		chunk.PositionsBuffer =
			new GpuBuffer<GrassData>( instanceCount );

		Vector3 mins = new(
			coord.X * ChunkSize,
			coord.Y * ChunkSize,
			TerrainProbeBottom
		);

		Vector3 maxs = new(
			mins.x + ChunkSize,
			mins.y + ChunkSize,
			TerrainProbeTop
		);

		chunk.Bounds = new BBox( mins, maxs );

		chunk.Material = GrassMaterial.CreateCopy();

		return chunk;
	}

	// =========================================================

	private void InitializeChunkHeightData( GrassChunk chunk )
	{
		int instanceCount = ChunkResolution * ChunkResolution;
		GrassData[] data = new GrassData[instanceCount];

		DensityMap.TryGetValue( chunk.Coord, out GrassChunkDensityMap densityMap );
		densityMap?.EnsureResolution( ChunkResolution );

		int P = Math.Max( 1, HeightProbeResolution );

		// Sample a low-resolution probe grid across the chunk
		float[,] probes = new float[P, P];

		for ( int py = 0; py < P; py++ )
		{
			for ( int px = 0; px < P; px++ )
			{
				float probeU = (P == 1) ? 0.5f : (float)px / (float)(P - 1);
				float probeV = (P == 1) ? 0.5f : (float)py / (float)(P - 1);

				float probeWorldX = (chunk.Coord.X * ChunkSize) + probeU * ChunkSize;
				float probeWorldY = (chunk.Coord.Y * ChunkSize) + probeV * ChunkSize;

				probes[px, py] = SampleTerrainHeight( probeWorldX, probeWorldY ) + TerrainHeightOffset;
			}
		}

		float minHeight = float.MaxValue;
		float maxHeight = float.MinValue;

		// Bilinearly upsample probe grid to fill full chunk resolution
		for ( int y = 0; y < ChunkResolution; y++ )
		{
			for ( int x = 0; x < ChunkResolution; x++ )
			{
				int index = x + y * ChunkResolution;

				float u = ( x + 0.5f ) / ChunkResolution;
				float v = ( y + 0.5f ) / ChunkResolution;

				float fx = (P == 1) ? 0.0f : u * (P - 1);
				float fy = (P == 1) ? 0.0f : v * (P - 1);

				int ix = (P == 1) ? 0 : Math.Clamp( (int)MathF.Floor( fx ), 0, P - 2 );
				int iy = (P == 1) ? 0 : Math.Clamp( (int)MathF.Floor( fy ), 0, P - 2 );

				float tx = (P == 1) ? 0.0f : fx - ix;
				float ty = (P == 1) ? 0.0f : fy - iy;

				float h00 = probes[ix, iy];
				float h10 = probes[ix + 1, iy];
				float h01 = probes[ix, iy + 1];
				float h11 = probes[ix + 1, iy + 1];

				float hx0 = h00 + ( h10 - h00 ) * tx;
				float hx1 = h01 + ( h11 - h01 ) * tx;
				float height = hx0 + ( hx1 - hx0 ) * ty;

				float worldX = (chunk.Coord.X * ChunkSize) + u * ChunkSize;
				float worldY = (chunk.Coord.Y * ChunkSize) + v * ChunkSize;

				data[index].Position = new Vector4( worldX, worldY, height, 1.0f );
				data[index].Uv = new Vector2( u, v );

				minHeight = MathF.Min( minHeight, height );
				maxHeight = MathF.Max( maxHeight, height );
			}
		}

		chunk.PositionsBuffer.SetData( data );

		if ( minHeight == float.MaxValue )
		{
			minHeight = FallbackHeight;
			maxHeight = FallbackHeight;
		}

		Vector3 mins = new(
			chunk.Coord.X * ChunkSize,
			chunk.Coord.Y * ChunkSize,
			minHeight - GrassHeightPadding
		);

		Vector3 maxs = new(
			mins.x + ChunkSize,
			mins.y + ChunkSize,
			maxHeight + GrassHeightPadding
		);

		chunk.Bounds = new BBox( mins, maxs );
	}

	private float SampleTerrainHeight( float worldX, float worldY )
	{
		if ( Game.ActiveScene == null )
			return FallbackHeight;

		Vector3 start = new( worldX, worldY, TerrainProbeTop );
		Vector3 end = new( worldX, worldY, TerrainProbeBottom );

		var trace = Game.ActiveScene.Trace.Ray( start, end ).WithTag("terrain").Run();

		if ( !trace.Hit )
			return FallbackHeight;

		return trace.EndPosition.z;
	}

	// =========================================================
	// RENDER THREAD
	// =========================================================

	public override void RenderSceneObject()
	{
		base.RenderSceneObject();

		GenerateWind();

		foreach ( var pair in ActiveChunks )
		{
			RenderChunk( pair.Value );
		}
	}

	private void RenderChunk( GrassChunk chunk )
	{
		if ( !DensityMap.TryGetValue( chunk.Coord, out GrassChunkDensityMap densityMap ) || densityMap == null || !densityMap.HasAnyDensity() )
			return;

		InteractionMap.TryGetValue( chunk.Coord, out GrassChunkInteractionMap interactionMap );

		Vector3 cameraPosition;
		Frustum cameraFrustum;

		if ( CullingCamera != null )
		{
			cameraPosition = CullingCamera.WorldPosition;
			cameraFrustum = CullingCamera.GetFrustum();
		}
		else
		{
			return;
		}

		if ( !cameraFrustum.IsInside( chunk.Bounds, true ) )
			return;

		float distSq = chunk.Bounds.Center.DistanceSquared(
			cameraPosition
		);

		float maxCull = DistanceCutoff + DistanceTransitionRange;
		if ( distSq > maxCull * maxCull )
			return;

		densityMap.EnsureTexture();

		chunk.Material.Attributes.Set(
			"_GrassDataBuffer",
			chunk.PositionsBuffer
		);

		chunk.Material.Attributes.Set(
			"_WindMap",
			Wind
		);

		chunk.Material.Attributes.Set(
			"_DensityMap",
			densityMap.GetTexture()
		);

		chunk.Material.Attributes.Set(
			"_InteractionMap",
			interactionMap?.GetTexture() ?? GrassChunkInteractionMap.EmptyTexture
		);

		chunk.Material.Attributes.Set(
			"_InteractionStrength",
			InteractionStrength
		);

		// Compute distance-based blend factors for smooth LOD transitions
		float dist = MathF.Sqrt( distSq );

		// Helper smoothstep implementation
		float SmoothStep( float a, float b, float x )
		{
			if ( x <= a ) return 0.0f;
			if ( x >= b ) return 1.0f;
			float t = ( x - a ) / ( b - a );
			return t * t * ( 3.0f - 2.0f * t );
		}

		float lod0Factor = 1.0f - SmoothStep( LodCutoff - LodTransitionRange, LodCutoff + LodTransitionRange, dist );

		float lod1In = SmoothStep( LodCutoff - LodTransitionRange, LodCutoff + LodTransitionRange, dist );
		float lod1Out = 1.0f - SmoothStep( DistanceCutoff - DistanceTransitionRange, DistanceCutoff + DistanceTransitionRange, dist );
		float lod1Factor = lod1In * lod1Out;

		// Draw both LODs where appropriate, scaling their effective density via _LodScale
		if ( lod0Factor > 0.001f )
		{
			chunk.Material.Attributes.Set( "_LodScale", lod0Factor );

			Graphics.DrawModelInstancedIndirect(
				GrassModel,
				chunk.ArgsBuffer,
				0,
				chunk.Material.Attributes
			);
		}

		if ( lod1Factor > 0.001f )
		{
			chunk.Material.Attributes.Set( "_LodScale", lod1Factor );

			Graphics.DrawModelInstancedIndirect(
				GrassModelLOD,
				chunk.ArgsBufferLOD,
				0,
				chunk.Material.Attributes
			);
		}
	}

	// =========================================================
	// WIND
	// =========================================================

	private void GenerateWind()
	{


		GenerateWindShader.Attributes.Set(
			"_WindMap",
			Wind
		);

		GenerateWindShader.Attributes.Set(
			"_Time",
			Time.Now * WindSpeed
		);

		GenerateWindShader.Attributes.Set(
			"_Frequency",
			WindFrequency
		);

		GenerateWindShader.Attributes.Set(
			"_Amplitude",
			WindAmplitude
		);

		GenerateWindShader.Attributes.Set(
			"_Direction",
			WindDirection.Normal
		);

		GenerateWindShader.Dispatch(
			Wind.Width,
			Wind.Height,
			1
		);
	}

	// =========================================================

	private ChunkCoord WorldToChunk( Vector3 worldPos )
	{
		return new ChunkCoord(
			MathX.FloorToInt( worldPos.x / ChunkSize ),
			MathX.FloorToInt( worldPos.y / ChunkSize )
		);
	}

	// =========================================================

	private uint[] BuildArgs( Model model, uint instanceCount )
	{
		return new uint[]
		{
			(uint)model.GetIndexCount( 0 ),
			instanceCount,
			(uint)model.GetIndexStart( 0 ),
			(uint)model.GetBaseVertex( 0 ),
			0
		};
	}

	// =========================================================

	public void Disable()
	{
		foreach ( var pair in ActiveChunks )
		{
			PendingDestroy.Enqueue( pair.Value );
		}

		ActiveChunks.Clear();

		ProcessPendingDestroy();
	}
}