using System;
using Sandbox;

internal sealed class PbrGeneratorPixelResult
{
	public int Width { get; set; }
	public int Height { get; set; }
	public bool AlbedoIsFloatingPoint { get; set; }
	public Color[] AlbedoPixels { get; set; }
	public Color[] HeightPixels { get; set; }
	public Color[] NormalPixels { get; set; }
	public Color[] RoughnessPixels { get; set; }
	public Color[] AmbientOcclusionPixels { get; set; }
	public Color[] MetallicPixels { get; set; }
	public Color[] OrmPixels { get; set; }
}

public static class PbrGenerator
{
	public static PbrGeneratorResult Generate( Bitmap source, PbrGeneratorSettings settings )
	{
		if ( source == null || !source.IsValid )
			return null;

		settings ??= new PbrGeneratorSettings();

		var width = source.Width;
		var height = source.Height;

		if ( width <= 0 || height <= 0 )
			return null;

		var sourcePixels = source.GetPixels();

		var pixels = GeneratePixels( sourcePixels, width, height, source.IsFloatingPoint, settings );
		return CreateBitmapResult( pixels );
	}

	internal static PbrGeneratorPixelResult GeneratePixels( Color[] sourcePixels, int width, int height, bool albedoIsFloatingPoint, PbrGeneratorSettings settings )
	{
		if ( sourcePixels == null || width <= 0 || height <= 0 || sourcePixels.Length < width * height )
			return null;

		settings ??= new PbrGeneratorSettings();

		var gray = CreateLuminanceValues( sourcePixels );
		var heightValues = CreateHeightValues( gray, width, height, settings );
		var normalValues = CreateNormalSourceValues( heightValues, width, height, settings );
		var roughnessValues = CreateRoughnessValues( gray, heightValues, width, height, settings );
		var aoValues = CreateAoValues( heightValues, width, height, settings );
		var metallicValues = CreateFlatValues( width * height, settings.MetallicValue );

		return new PbrGeneratorPixelResult
		{
			Width = width,
			Height = height,
			AlbedoIsFloatingPoint = albedoIsFloatingPoint,
			AlbedoPixels = (Color[])sourcePixels.Clone(),
			HeightPixels = CreateGrayscalePixels( heightValues ),
			NormalPixels = CreateNormalPixels( normalValues, width, height, settings ),
			RoughnessPixels = CreateGrayscalePixels( roughnessValues ),
			AmbientOcclusionPixels = CreateGrayscalePixels( aoValues ),
			MetallicPixels = CreateGrayscalePixels( metallicValues ),
			OrmPixels = CreateOrmPixels( aoValues, roughnessValues, metallicValues )
		};
	}

	internal static PbrGeneratorResult CreateBitmapResult( PbrGeneratorPixelResult pixels )
	{
		if ( pixels == null || pixels.Width <= 0 || pixels.Height <= 0 )
			return null;

		return new PbrGeneratorResult
		{
			Albedo = CreateBitmap( pixels.AlbedoPixels, pixels.Width, pixels.Height, pixels.AlbedoIsFloatingPoint ),
			Height = CreateBitmap( pixels.HeightPixels, pixels.Width, pixels.Height, false ),
			Normal = CreateBitmap( pixels.NormalPixels, pixels.Width, pixels.Height, false ),
			Roughness = CreateBitmap( pixels.RoughnessPixels, pixels.Width, pixels.Height, false ),
			AmbientOcclusion = CreateBitmap( pixels.AmbientOcclusionPixels, pixels.Width, pixels.Height, false ),
			Metallic = CreateBitmap( pixels.MetallicPixels, pixels.Width, pixels.Height, false ),
			Orm = CreateBitmap( pixels.OrmPixels, pixels.Width, pixels.Height, false )
		};
	}

	public static string GetMapLabel( PbrMapType mapType )
	{
		return mapType switch
		{
			PbrMapType.Albedo => "Albedo",
			PbrMapType.Height => "Height",
			PbrMapType.Normal => "Normal",
			PbrMapType.Roughness => "Roughness",
			PbrMapType.AmbientOcclusion => "AO",
			PbrMapType.Metallic => "Metallic",
			PbrMapType.Orm => "ORM",
			_ => "Map"
		};
	}

	public static string GetMapSuffix( PbrMapType mapType )
	{
		return mapType switch
		{
			PbrMapType.Albedo => "_albedo",
			PbrMapType.Height => "_height",
			PbrMapType.Normal => "_normal",
			PbrMapType.Roughness => "_roughness",
			PbrMapType.AmbientOcclusion => "_ao",
			PbrMapType.Metallic => "_metallic",
			PbrMapType.Orm => "_orm",
			_ => "_map"
		};
	}

	private static float[] CreateLuminanceValues( Color[] sourcePixels )
	{
		var values = new float[sourcePixels.Length];

		for ( var i = 0; i < sourcePixels.Length; i++ )
		{
			values[i] = SeamlessSuiteImageUtility.GetLuminance( sourcePixels[i] );
		}

		return values;
	}

	private static float[] CreateHeightValues( float[] gray, int width, int height, PbrGeneratorSettings settings )
	{
		if ( settings.AdvancedHeightEnabled )
			return CreateAdvancedHeightValues( gray, width, height, settings );

		return CreateSimpleHeightValues( gray, width, height, settings );
	}

	private static float[] CreateSimpleHeightValues( float[] gray, int width, int height, PbrGeneratorSettings settings )
	{
		var values = new float[gray.Length];
		var contrast = MathF.Max( 0f, settings.HeightContrast );
		var strength = MathF.Max( 0f, settings.HeightStrength );
		var blackPoint = SeamlessSuiteImageUtility.Clamp01( settings.HeightBlackPoint );
		var whitePoint = SeamlessSuiteImageUtility.Clamp01( settings.HeightWhitePoint );
		var midpointOffset = 0.5f - SeamlessSuiteImageUtility.Clamp01( settings.HeightMidpoint );
		var cavityEmphasis = SeamlessSuiteImageUtility.Clamp01( settings.HeightCavityEmphasis );

		if ( whitePoint <= blackPoint + 0.001f )
		{
			whitePoint = Math.Min( 1f, blackPoint + 0.001f );

			if ( whitePoint <= blackPoint )
				blackPoint = Math.Max( 0f, whitePoint - 0.001f );
		}

		for ( var i = 0; i < gray.Length; i++ )
		{
			var value = (gray[i] - blackPoint) / (whitePoint - blackPoint);
			value = SeamlessSuiteImageUtility.Clamp01( value );
			value = 0.5f + (value - 0.5f) * contrast;
			value = 0.5f + (value - 0.5f) * strength;
			value += midpointOffset;
			value = ApplyHeightCavityEmphasis( value, cavityEmphasis );
			values[i] = SeamlessSuiteImageUtility.Clamp01( value );
		}

		if ( settings.HeightBlur > 0.01f )
			values = FastBlurValues( values, width, height, settings.HeightBlur );

		if ( settings.InvertHeight )
		{
			for ( var i = 0; i < values.Length; i++ )
			{
				values[i] = 1f - values[i];
			}
		}

		return values;
	}

	private static float[] CreateAdvancedHeightValues( float[] gray, int width, int height, PbrGeneratorSettings settings )
	{
		var scale = MathF.Max( 1f, settings.AdvancedHeightScaleRadius );
		var fineRadius = Math.Max( 1, (int)MathF.Round( scale ) );
		var mediumRadius = Math.Max( fineRadius + 1, (int)MathF.Round( scale * 3f ) );
		var largeRadius = Math.Max( mediumRadius + 1, (int)MathF.Round( scale * 8f ) );
		var broadRadius = Math.Max( largeRadius + 1, (int)MathF.Round( scale * 16f ) );

		var fineBlur = FastBlurValues( gray, width, height, fineRadius );
		var mediumBlur = FastBlurValues( gray, width, height, mediumRadius );
		var largeBlur = FastBlurValues( gray, width, height, largeRadius );
		var broadBlur = FastBlurValues( gray, width, height, broadRadius );

		var fineStrength = MathF.Max( 0f, settings.AdvancedHeightFineStrength );
		var mediumStrength = MathF.Max( 0f, settings.AdvancedHeightMediumStrength );
		var largeStrength = MathF.Max( 0f, settings.AdvancedHeightLargeStrength );
		var broadStrength = MathF.Max( 0f, settings.AdvancedHeightBroadStrength );
		var finalContrast = MathF.Max( 0f, settings.AdvancedHeightFinalContrast );
		var gain = MathF.Max( 0f, settings.AdvancedHeightGain );
		var bias = settings.AdvancedHeightBias;
		var values = new float[gray.Length];

		for ( var i = 0; i < gray.Length; i++ )
		{
			var baseHeight = gray[i] - 0.5f;
			var fineDetail = gray[i] - fineBlur[i];
			var mediumDetail = fineBlur[i] - mediumBlur[i];
			var largeDetail = mediumBlur[i] - largeBlur[i];
			var broadDetail = largeBlur[i] - broadBlur[i];

			var value = 0.5f + baseHeight
				+ fineDetail * fineStrength
				+ mediumDetail * mediumStrength
				+ largeDetail * largeStrength
				+ broadDetail * broadStrength;

			value = 0.5f + (value - 0.5f) * gain;
			value = 0.5f + (value - 0.5f) * finalContrast;
			value += bias;
			values[i] = SeamlessSuiteImageUtility.Clamp01( value );
		}

		if ( settings.AdvancedHeightSmoothing > 0.01f )
			values = FastBlurValues( values, width, height, settings.AdvancedHeightSmoothing );

		if ( settings.InvertHeight )
		{
			for ( var i = 0; i < values.Length; i++ )
			{
				values[i] = 1f - values[i];
			}
		}

		return values;
	}

	private static float ApplyHeightCavityEmphasis( float value, float amount )
	{
		if ( amount <= 0f || value >= 0.5f )
			return value;

		value = SeamlessSuiteImageUtility.Clamp01( value );
		var recessed = value / 0.5f;
		var emphasized = recessed * recessed * 0.5f;
		return value + (emphasized - value) * amount;
	}

	private static float[] CreateNormalSourceValues( float[] heightValues, int width, int height, PbrGeneratorSettings settings )
	{
		if ( settings.NormalSmoothing <= 0.01f )
			return heightValues;

		return FastBlurValues( heightValues, width, height, settings.NormalSmoothing );
	}

	private static float[] CreateRoughnessValues( float[] gray, float[] heightValues, int width, int height, PbrGeneratorSettings settings )
	{
		var values = new float[gray.Length];
		var baseValue = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessBase );
		var contrast = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessContrast );
		var blackPoint = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessBlackPoint );
		var whitePoint = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessWhitePoint );
		var cavity = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessCavity );
		var variation = SeamlessSuiteImageUtility.Clamp01( settings.RoughnessVariation );
		var blurred = BoxBlurValues( gray, width, height, 2 );

		if ( whitePoint <= blackPoint + 0.001f )
		{
			whitePoint = Math.Min( 1f, blackPoint + 0.001f );

			if ( whitePoint <= blackPoint )
				blackPoint = Math.Max( 0f, whitePoint - 0.001f );
		}

		for ( var y = 0; y < height; y++ )
		{
			for ( var x = 0; x < width; x++ )
			{
				var index = y * width + x;
				var detail = SeamlessSuiteImageUtility.Clamp01( MathF.Abs( gray[index] - blurred[index] ) * 2f );
				var value = baseValue + (detail - 0.5f) * contrast;

				if ( cavity > 0f )
				{
					var recessed = SeamlessSuiteImageUtility.Clamp01( (0.5f - heightValues[index]) * 2f );
					value = SeamlessSuiteImageUtility.Lerp( value, 1f, recessed * recessed * cavity );
				}

				if ( variation > 0f )
				{
					var noise = GetWrappedValueNoise( x, y, width, height );
					value += (noise - 0.5f) * 0.25f * variation;
				}

				value = SeamlessSuiteImageUtility.Clamp01( value );
				value = SeamlessSuiteImageUtility.Clamp01( (value - blackPoint) / (whitePoint - blackPoint) );

				if ( settings.InvertRoughness )
					value = 1f - value;

				values[index] = value;
			}
		}

		return values;
	}

	private static float[] CreateAoValues( float[] heightValues, int width, int height, PbrGeneratorSettings settings )
	{
		var values = new float[heightValues.Length];
		var strength = Math.Clamp( settings.AoStrength, 0f, 2f );
		var radius = Math.Clamp( (int)MathF.Round( settings.AoRadius ), 1, Math.Max( width, height ) );
		var minimum = SeamlessSuiteImageUtility.Clamp01( settings.AoMinimum );
		var sampleOffsets = new[]
		{
			( radius, 0 ),
			( -radius, 0 ),
			( 0, radius ),
			( 0, -radius ),
			( radius, radius ),
			( -radius, radius ),
			( radius, -radius ),
			( -radius, -radius )
		};

		for ( var y = 0; y < height; y++ )
		{
			for ( var x = 0; x < width; x++ )
			{
				var index = y * width + x;
				var center = heightValues[index];
				var cavity = 0f;

				foreach ( var offset in sampleOffsets )
				{
					var sample = SampleWrapped( heightValues, width, height, x + offset.Item1, y + offset.Item2 );
					cavity += MathF.Max( 0f, sample - center );
				}

				cavity /= sampleOffsets.Length;
				var ao = 1f - cavity * strength * 2f;
				values[index] = Math.Clamp( ao, minimum, 1f );
			}
		}

		return values;
	}

	private static float[] CreateFlatValues( int count, float value )
	{
		var values = new float[count];
		var clamped = SeamlessSuiteImageUtility.Clamp01( value );

		for ( var i = 0; i < values.Length; i++ )
		{
			values[i] = clamped;
		}

		return values;
	}

	private static Color[] CreateGrayscalePixels( float[] values )
	{
		var pixels = new Color[values.Length];

		for ( var i = 0; i < values.Length; i++ )
		{
			var value = SeamlessSuiteImageUtility.Clamp01( values[i] );
			pixels[i] = new Color( value, value, value, 1f );
		}

		return pixels;
	}

	private static Color[] CreateNormalPixels( float[] heightValues, int width, int height, PbrGeneratorSettings settings )
	{
		var pixels = new Color[heightValues.Length];
		var strength = MathF.Max( 0f, settings.NormalStrength );

		for ( var y = 0; y < height; y++ )
		{
			for ( var x = 0; x < width; x++ )
			{
				var left = SampleWrapped( heightValues, width, height, x - 1, y );
				var right = SampleWrapped( heightValues, width, height, x + 1, y );
				var up = SampleWrapped( heightValues, width, height, x, y - 1 );
				var down = SampleWrapped( heightValues, width, height, x, y + 1 );

				var slopeX = (left - right) * strength;
				var slopeY = (up - down) * strength;

				if ( settings.FlipNormalY )
					slopeY = -slopeY;

				var normalLength = MathF.Sqrt( slopeX * slopeX + slopeY * slopeY + 1f );
				var nx = slopeX / normalLength;
				var ny = slopeY / normalLength;
				var nz = 1f / normalLength;

				pixels[y * width + x] = new Color(
					nx * 0.5f + 0.5f,
					ny * 0.5f + 0.5f,
					nz * 0.5f + 0.5f,
					1f
				);
			}
		}

		return pixels;
	}

	private static Color[] CreateOrmPixels( float[] aoValues, float[] roughnessValues, float[] metallicValues )
	{
		var pixels = new Color[aoValues.Length];

		for ( var i = 0; i < pixels.Length; i++ )
		{
			pixels[i] = new Color(
				SeamlessSuiteImageUtility.Clamp01( aoValues[i] ),
				SeamlessSuiteImageUtility.Clamp01( roughnessValues[i] ),
				SeamlessSuiteImageUtility.Clamp01( metallicValues[i] ),
				1f
			);
		}

		return pixels;
	}

	private static Bitmap CreateBitmap( Color[] pixels, int width, int height, bool floatingPoint )
	{
		if ( pixels == null || pixels.Length < width * height )
			return null;

		var bitmap = new Bitmap( width, height, floatingPoint );
		bitmap.SetPixels( pixels );
		return bitmap;
	}

	private static float[] FastBlurValues( float[] values, int width, int height, float radius )
	{
		if ( radius <= 0.01f )
			return (float[])values.Clone();

		var blur = values;
		var boxRadii = GetGaussianBoxRadii( radius, 3 );

		foreach ( var boxRadius in boxRadii )
		{
			if ( boxRadius <= 0 )
				continue;

			blur = FastBoxBlurValues( blur, width, height, boxRadius );
		}

		return ReferenceEquals( blur, values ) ? (float[])values.Clone() : blur;
	}

	private static int[] GetGaussianBoxRadii( float sigma, int passCount )
	{
		sigma = MathF.Max( 0.01f, sigma );
		var idealWidth = MathF.Sqrt( 12f * sigma * sigma / passCount + 1f );
		var lowerWidth = (int)MathF.Floor( idealWidth );

		if ( lowerWidth % 2 == 0 )
			lowerWidth--;

		lowerWidth = Math.Max( 1, lowerWidth );
		var upperWidth = lowerWidth + 2;
		var passSplit = (int)MathF.Round(
			(12f * sigma * sigma - passCount * lowerWidth * lowerWidth - 4f * passCount * lowerWidth - 3f * passCount)
			/ (-4f * lowerWidth - 4f)
		);
		passSplit = Math.Clamp( passSplit, 0, passCount );

		var radii = new int[passCount];

		for ( var i = 0; i < passCount; i++ )
		{
			var width = i < passSplit ? lowerWidth : upperWidth;
			radii[i] = Math.Max( 0, (width - 1) / 2 );
		}

		return radii;
	}

	private static float[] BoxBlurValues( float[] values, int width, int height, int radius )
	{
		if ( radius <= 0 )
			return (float[])values.Clone();

		var blurred = new float[values.Length];
		var diameter = radius * 2 + 1;
		var sampleCount = diameter * diameter;

		for ( var y = 0; y < height; y++ )
		{
			for ( var x = 0; x < width; x++ )
			{
				var sum = 0f;

				for ( var sampleY = -radius; sampleY <= radius; sampleY++ )
				{
					for ( var sampleX = -radius; sampleX <= radius; sampleX++ )
					{
						sum += SampleWrapped( values, width, height, x + sampleX, y + sampleY );
					}
				}

				blurred[y * width + x] = sum / sampleCount;
			}
		}

		return blurred;
	}

	private static float[] FastBoxBlurValues( float[] values, int width, int height, int radius )
	{
		if ( radius <= 0 )
			return (float[])values.Clone();

		var diameter = radius * 2 + 1;
		var horizontal = new float[values.Length];
		var blurred = new float[values.Length];

		for ( var y = 0; y < height; y++ )
		{
			var row = y * width;
			var sum = 0f;

			for ( var sampleX = -radius; sampleX <= radius; sampleX++ )
			{
				var sourceX = Math.Clamp( sampleX, 0, width - 1 );
				sum += values[row + sourceX];
			}

			for ( var x = 0; x < width; x++ )
			{
				horizontal[row + x] = sum / diameter;

				var removeX = Math.Clamp( x - radius, 0, width - 1 );
				var addX = Math.Clamp( x + radius + 1, 0, width - 1 );
				sum += values[row + addX] - values[row + removeX];
			}
		}

		for ( var x = 0; x < width; x++ )
		{
			var sum = 0f;

			for ( var sampleY = -radius; sampleY <= radius; sampleY++ )
			{
				var sourceY = Math.Clamp( sampleY, 0, height - 1 );
				sum += horizontal[sourceY * width + x];
			}

			for ( var y = 0; y < height; y++ )
			{
				blurred[y * width + x] = sum / diameter;

				var removeY = Math.Clamp( y - radius, 0, height - 1 );
				var addY = Math.Clamp( y + radius + 1, 0, height - 1 );
				sum += horizontal[addY * width + x] - horizontal[removeY * width + x];
			}
		}

		return blurred;
	}

	private static float SampleWrapped( float[] values, int width, int height, int x, int y )
	{
		var wrappedX = SeamlessSuiteImageUtility.WrapIndex( x, width );
		var wrappedY = SeamlessSuiteImageUtility.WrapIndex( y, height );
		return values[wrappedY * width + wrappedX];
	}

	private static float GetWrappedValueNoise( int x, int y, int width, int height )
	{
		var gridX = Math.Clamp( width / 48, 4, 32 );
		var gridY = Math.Clamp( height / 48, 4, 32 );
		var u = x / Math.Max( 1f, width ) * gridX;
		var v = y / Math.Max( 1f, height ) * gridY;
		var x0 = (int)MathF.Floor( u );
		var y0 = (int)MathF.Floor( v );
		var xBlend = SmoothStep( u - x0 );
		var yBlend = SmoothStep( v - y0 );

		var a = Hash01Wrapped( x0, y0, gridX, gridY );
		var b = Hash01Wrapped( x0 + 1, y0, gridX, gridY );
		var c = Hash01Wrapped( x0, y0 + 1, gridX, gridY );
		var d = Hash01Wrapped( x0 + 1, y0 + 1, gridX, gridY );
		var top = SeamlessSuiteImageUtility.Lerp( a, b, xBlend );
		var bottom = SeamlessSuiteImageUtility.Lerp( c, d, xBlend );

		return SeamlessSuiteImageUtility.Lerp( top, bottom, yBlend );
	}

	private static float Hash01Wrapped( int x, int y, int gridX, int gridY )
	{
		x = SeamlessSuiteImageUtility.WrapIndex( x, gridX );
		y = SeamlessSuiteImageUtility.WrapIndex( y, gridY );

		unchecked
		{
			var hash = x * 374761393 + y * 668265263;// random large numbers for noise generation
			hash = (hash ^ (hash >> 13)) * 1274126177;
			return ((hash ^ (hash >> 16)) & 0x7fffffff) / (float)int.MaxValue;
		}
	}

	private static float SmoothStep( float value )
	{
		value = SeamlessSuiteImageUtility.Clamp01( value );
		return value * value * (3f - 2f * value);
	}
}