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

namespace TerrainLines;

// Rasterization only reads storage. Preview and commit use the same tile results.
internal static class LineRaster
{
	internal const int TileSize = 32;

	internal static float Distance2D( Vector3 a, Vector3 b )
	{
		float x = a.x - b.x, y = a.y - b.y;
		return MathF.Sqrt( x * x + y * y );
	}

	internal static float SampleHeight( TerrainStorage storage, float x, float y, Func<int, int, ushort> read = null )
	{
		int n = storage.Resolution;
		float px = Math.Clamp( x / storage.TerrainSize * n, 0, n - 1 );
		float py = Math.Clamp( y / storage.TerrainSize * n, 0, n - 1 );
		int x0 = (int)px, y0 = (int)py;
		int x1 = Math.Min( x0 + 1, n - 1 ), y1 = Math.Min( y0 + 1, n - 1 );
		float top = Lerp( read is null ? storage.HeightMap[y0 * n + x0] : read( x0, y0 ), read is null ? storage.HeightMap[y0 * n + x1] : read( x1, y0 ), px - x0 );
		float bottom = Lerp( read is null ? storage.HeightMap[y1 * n + x0] : read( x0, y1 ), read is null ? storage.HeightMap[y1 * n + x1] : read( x1, y1 ), px - x0 );
		return Lerp( top, bottom, py - y0 ) * storage.TerrainHeight / ushort.MaxValue;
	}

	internal static List<LineTile> Build( TerrainStorage storage, IReadOnlyList<Vector3> points, LineSettings settings, float worldScale )
	{
		if ( !settings.IsFinite || !float.IsFinite( worldScale ) || worldScale <= 0 )
			throw new ArgumentException( "Line settings must be finite and terrain scale must be positive." );
		if ( storage.Resolution < 2 || storage.TerrainSize <= 0 || storage.TerrainHeight <= 0 )
			throw new ArgumentException( "Terrain has invalid dimensions." );
		if ( points.Count < 2 ) return [];
		int resolution = storage.Resolution;
		float texel = storage.TerrainSize / resolution;
		float halfWidth = Math.Max( 1, settings.PathWidth ) * 0.5f / worldScale;
		float shoulder = Math.Max( 0, settings.ShoulderWidth ) / worldScale;
		float radius = halfWidth + shoulder + (settings.PaintMaterials ? texel : 0);
		float totalLength = 0;
		for ( int i = 1; i < points.Count; i++ ) totalLength += Distance2D( points[i - 1], points[i] );
		if ( totalLength < 0.001f ) return [];

		bool paint = settings.PaintMaterials && settings.PaintStrength > 0;
		if ( paint && (settings.SurfaceMaterial < 0 || settings.SurfaceMaterial >= Math.Min( 32, storage.Materials.Count )) )
			throw new ArgumentException( "Select a surface material already assigned to this terrain." );
		if ( paint && (settings.ShoulderMaterial < -1 || settings.ShoulderMaterial >= Math.Min( 32, storage.Materials.Count )) )
			throw new ArgumentException( "Select a shoulder material already assigned to this terrain." );
		if ( !settings.ShapeTerrain && !paint ) return [];

		var tiles = new Dictionary<(int, int), WorkingTile>();
		float travelled = 0;
		var profile = settings.ShapeTerrain ? new LineProfile( storage, points, totalLength, settings.EvenRamp, settings.SlopeLimit ) : null;
		for ( int segment = 1; segment < points.Count; segment++ )
		{
			var a = points[segment - 1];
			var b = points[segment];
			float length = Distance2D( a, b );
			if ( length < 0.001f ) continue;
			// Bound work along diagonal straight lines instead of visiting their whole bounding rectangle.
			int steps = Math.Max( 1, (int)MathF.Ceiling( length / Math.Max( texel * 4, radius ) ) );
			for ( int step = 0; step < steps; step++ )
			{
				float t0 = (float)step / steps, t1 = (float)(step + 1) / steps;
				var from = a + (b - a) * t0;
				var to = a + (b - a) * t1;
				int left = Math.Max( 0, (int)MathF.Floor( (Math.Min( from.x, to.x ) - radius) / texel ) );
				int right = Math.Min( resolution - 1, (int)MathF.Ceiling( (Math.Max( from.x, to.x ) + radius) / texel ) );
				int top = Math.Max( 0, (int)MathF.Floor( (Math.Min( from.y, to.y ) - radius) / texel ) );
				int bottom = Math.Min( resolution - 1, (int)MathF.Ceiling( (Math.Max( from.y, to.y ) + radius) / texel ) );
				float dx = to.x - from.x, dy = to.y - from.y;
				float squaredLength = dx * dx + dy * dy;
				for ( int y = top; y <= bottom; y++ )
				for ( int x = left; x <= right; x++ )
				{
					int index = y * resolution + x;
					if ( new CompactTerrainMaterial( storage.ControlMap[index] ).IsHole ) continue;
					float t = Math.Clamp( ((x * texel - from.x) * dx + (y * texel - from.y) * dy) / squaredLength, 0, 1 );
					float cx = from.x + dx * t, cy = from.y + dy * t;
					float distanceSquared = MathF.Pow( x * texel - cx, 2 ) + MathF.Pow( y * texel - cy, 2 );
					if ( distanceSquared > radius * radius ) continue;
					var key = (x / TileSize, y / TileSize);
					if ( !tiles.TryGetValue( key, out var tile ) )
					{
						tile = new WorkingTile( storage, key.Item1 * TileSize, key.Item2 * TileSize, settings.ShapeTerrain, paint );
						tiles.Add( key, tile );
					}
					int local = (y - tile.Result.Y) * tile.Result.Width + x - tile.Result.X;
					// Nearest segment wins at crossings; later segments win exact ties. Never accumulate stamps.
					if ( distanceSquared > tile.Distance[local] ) continue;
					tile.Distance[local] = distanceSquared;
					float segmentT = Lerp( t0, t1, t );
					if ( profile is not null )
						tile.Height[local] = profile.Sample( travelled + length * segmentT ) + settings.Elevation / worldScale;
				}
			}
			travelled += length;
		}

		var result = new List<LineTile>();
		foreach ( var tile in tiles.Values )
		{
			var output = tile.Result;
			for ( int i = 0; i < tile.Distance.Length; i++ )
			{
				if ( !float.IsFinite( tile.Distance[i] ) ) continue;
				float distance = MathF.Sqrt( tile.Distance[i] );
				float edge = distance <= halfWidth ? 1 : shoulder > 0 ? Math.Clamp( 1 - (distance - halfWidth) / shoulder, 0, 1 ) : 0;
				float weight = ShoulderWeight( edge, settings.ShoulderCurve );
				if ( output.AfterHeight is not null )
				{
					float target = Math.Clamp( tile.Height[i] / storage.TerrainHeight, 0, 1 ) * ushort.MaxValue;
					output.AfterHeight[i] = (ushort)Math.Clamp( (int)MathF.Round( Lerp( output.BeforeHeight[i], target, weight ) ), 0, ushort.MaxValue );
				}
				if ( output.AfterControl is not null )
				{
					// Paint uses a filtered edge at the path boundary, independently of the sculpt shoulder.
					float core = Coverage( distance, halfWidth, texel );
					float strength = Math.Clamp( settings.PaintStrength, 0, 1 );
					float outer = Coverage( distance, halfWidth + shoulder, texel );
					output.AfterControl[i] = shoulder > 0 && settings.ShoulderMaterial >= 0
						? PaintBands( output.BeforeControl[i], settings.SurfaceMaterial, settings.ShoulderMaterial, core, outer, strength )
						: Paint( output.BeforeControl[i], settings.SurfaceMaterial, core * strength );
				}
			}
			if ( output.BeforeHeight is not null && output.BeforeHeight.SequenceEqual( output.AfterHeight ) )
				output.BeforeHeight = output.AfterHeight = null;
			if ( output.BeforeControl is not null && output.BeforeControl.SequenceEqual( output.AfterControl ) )
				output.BeforeControl = output.AfterControl = null;
			if ( output.Flags != 0 ) result.Add( output );
		}
		return result;
	}

	internal static float ShoulderWeight( float edge, ShoulderCurve curve ) => curve switch
	{
		ShoulderCurve.Linear => edge,
		ShoulderCurve.RoundIn => edge * edge,
		ShoulderCurve.RoundOut => 1 - (1 - edge) * (1 - edge),
		_ => edge * edge * (3 - 2 * edge)
	};

	internal static float Coverage( float distance, float radius, float texel )
	{
		float t = Math.Clamp( (radius - distance) / texel + 0.5f, 0, 1 );
		return t * t * (3 - 2 * t);
	}

	internal static uint PaintBands( uint packed, int surface, int shoulder, float core, float outer, float strength )
	{
		if ( surface == shoulder ) return Paint( packed, surface, outer * strength );
		if ( core >= 1 ) return Paint( packed, surface, strength );
		if ( core <= 0 ) return Paint( packed, shoulder, outer * strength );
		var material = new CompactTerrainMaterial( packed );
		if ( material.IsHole || strength <= 0 ) return packed;
		// Blend both bands once, retaining the two strongest layers supported by native terrain.
		Span<float> weights = stackalloc float[32];
		weights.Clear();
		float retained = 1 - outer * strength;
		float overlay = material.BlendFactor / 255f;
		weights[material.BaseTextureId] += (1 - overlay) * retained;
		weights[material.OverlayTextureId] += overlay * retained;
		weights[surface] += core * strength;
		weights[shoulder] += (outer - core) * strength;
		int first = 0;
		for ( int i = 1; i < 32; i++ ) if ( weights[i] > weights[first] ) first = i;
		int second = first == 0 ? 1 : 0;
		for ( int i = 0; i < 32; i++ ) if ( i != first && weights[i] > weights[second] ) second = i;
		material.BaseTextureId = (byte)first;
		material.OverlayTextureId = (byte)second;
		material.BlendFactor = (byte)Math.Clamp( (int)MathF.Round( weights[second] / (weights[first] + weights[second]) * 255 ), 0, 255 );
		return material.Packed;
	}

	internal static uint Paint( uint packed, int id, float strength )
	{
		if ( strength <= 0 ) return packed;
		var material = new CompactTerrainMaterial( packed );
		if ( material.IsHole ) return packed;
		float overlay = material.BlendFactor / 255f;
		if ( material.BaseTextureId == id ) overlay *= 1 - strength;
		else if ( material.OverlayTextureId == id ) overlay = Lerp( overlay, 1, strength );
		else
		{
			if ( overlay > 0.5f ) material.BaseTextureId = material.OverlayTextureId;
			material.OverlayTextureId = (byte)id;
			overlay = strength;
		}
		if ( overlay > 0.5f )
		{
			(material.BaseTextureId, material.OverlayTextureId) = (material.OverlayTextureId, material.BaseTextureId);
			overlay = 1 - overlay;
		}
		material.BlendFactor = (byte)Math.Clamp( (int)MathF.Round( overlay * 255 ), 0, 255 );
		return material.Packed;
	}

	private static float Lerp( float a, float b, float t ) => a + (b - a) * t;

	private sealed class WorkingTile
	{
		internal readonly LineTile Result;
		internal readonly float[] Distance;
		internal readonly float[] Height;
		internal WorkingTile( TerrainStorage storage, int x, int y, bool shape, bool paint )
		{
			Result = new LineTile( storage, x, y, shape, paint );
			Distance = new float[Result.Width * Result.Height];
			Array.Fill( Distance, float.PositiveInfinity );
			Height = new float[Distance.Length];
		}
	}
}

internal sealed class LineTile
{
	internal readonly int X, Y, Width, Height;
	internal ushort[] BeforeHeight, AfterHeight;
	internal uint[] BeforeControl, AfterControl;
	internal Terrain.SyncFlags Flags => (BeforeHeight is null ? 0 : Terrain.SyncFlags.Height) | (BeforeControl is null ? 0 : Terrain.SyncFlags.Control);

	internal LineTile( TerrainStorage storage, int x, int y, bool shape, bool paint )
	{
		X = x; Y = y;
		Width = Math.Min( LineRaster.TileSize, storage.Resolution - x );
		Height = Math.Min( LineRaster.TileSize, storage.Resolution - y );
		if ( shape )
		{
			BeforeHeight = Copy( storage.HeightMap, storage.Resolution );
			AfterHeight = (ushort[])BeforeHeight.Clone();
		}
		if ( paint )
		{
			BeforeControl = Copy( storage.ControlMap, storage.Resolution );
			AfterControl = (uint[])BeforeControl.Clone();
		}
	}

	private T[] Copy<T>( T[] source, int stride )
	{
		var result = new T[Width * Height];
		for ( int row = 0; row < Height; row++ ) Array.Copy( source, (Y + row) * stride + X, result, row * Width, Width );
		return result;
	}

	internal bool Matches( TerrainStorage storage )
	{
		for ( int row = 0; row < Height; row++ )
		for ( int column = 0; column < Width; column++ )
		{
			int local = row * Width + column, index = (Y + row) * storage.Resolution + X + column;
			if ( BeforeHeight is not null && BeforeHeight[local] != AfterHeight[local] && storage.HeightMap[index] != BeforeHeight[local] ) return false;
			if ( BeforeControl is not null && BeforeControl[local] != AfterControl[local] && storage.ControlMap[index] != BeforeControl[local] ) return false;
		}
		return true;
	}
	internal void UploadPreview( Terrain terrain )
	{
		if ( BeforeHeight is not null ) terrain.HeightMap.Update<ushort>( AfterHeight, X, Y, Width, Height );
		if ( BeforeControl is not null ) terrain.ControlMap.Update<uint>( AfterControl, X, Y, Width, Height );
	}

	internal void WriteStorage( TerrainStorage storage, bool after )
	{
		// Only restore changed samples, preserving unrelated edits inside the same snapshot tile.
		for ( int row = 0; row < Height; row++ )
		for ( int column = 0; column < Width; column++ )
		{
			int local = row * Width + column, index = (Y + row) * storage.Resolution + X + column;
			if ( BeforeHeight is not null && BeforeHeight[local] != AfterHeight[local] ) storage.HeightMap[index] = after ? AfterHeight[local] : BeforeHeight[local];
			if ( BeforeControl is not null && BeforeControl[local] != AfterControl[local] ) storage.ControlMap[index] = after ? AfterControl[local] : BeforeControl[local];
		}
	}
	internal void RestorePreview( Terrain terrain )
	{
		if ( BeforeHeight is not null ) terrain.HeightMap.Update<ushort>( Copy( terrain.Storage.HeightMap, terrain.Storage.Resolution ), X, Y, Width, Height );
		if ( BeforeControl is not null ) terrain.ControlMap.Update<uint>( Copy( terrain.Storage.ControlMap, terrain.Storage.Resolution ), X, Y, Width, Height );
	}
}