using AssetRipper.IO.Endian;
using AssetRipper.SourceGenerated.Classes.ClassID_28;
using AssetRipper.SourceGenerated.Enums;
using System.Diagnostics;
using System.Drawing;

namespace AssetRipper.SourceGenerated.Extensions;

public static class SwitchSwizzle
{
	private const int GobXTexelCount = 4;
	private const int GobYTexelCount = 8;
	private const int TexelByteSize = 16;

	/// <summary>
	/// In this case, "block" refers to a compressed texture block, not a gob one.
	/// </summary>
	private const int BlocksInGob = GobXTexelCount * GobYTexelCount;

	private static readonly int[] GobXPoses =
	[
		0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3
	];

	private static readonly int[] GobYPoses =
	[
		0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7, 6, 7, 0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7, 6, 7
	];

	private static int CeilDivide(int a, int b)
	{
		return (a + b - 1) / b;
	}

	public static byte[] Unswizzle(ITexture2D texture, byte[] data)
	{
		// Just in case
		if (texture.Mips || texture.ImageCount_C28 != 1)
		{
			return data;
		}

		TextureFormat realFormat = GetCorrectedSwitchTextureFormat(texture.Format_C28E);

		// Format is unsupported, we back out
		if (GetTextureFormatBlockSize(realFormat) is not { } blockSize)
		{
			return data;
		}

		byte[] newData = new byte[data.Length];

		Debug.Assert(texture.PlatformBlob_C28 is not null && texture.PlatformBlob_C28.Length >= 12);
		int blockHeight = GetBlockHeightByPlatformBlob(texture.PlatformBlob_C28, texture.Collection.EndianType);

		Size paddedSize = GetPaddedTextureSize(texture.Width_C28, texture.Height_C28, blockSize.Width, blockSize.Height, blockHeight);

		int blockCountX = CeilDivide(paddedSize.Width, blockSize.Width);
		int blockCountY = CeilDivide(paddedSize.Height, blockSize.Height);

		int gobCountX = blockCountX / GobXTexelCount;
		int gobCountY = blockCountY / GobYTexelCount;

		int srcPos = 0;
		for (int i = 0; i < gobCountY / blockHeight; i++)
		{
			for (int j = 0; j < gobCountX; j++)
			{
				for (int k = 0; k < blockHeight; k++)
				{
					for (int l = 0; l < BlocksInGob; l++)
					{
						int gobX = GobXPoses[l];
						int gobY = GobYPoses[l];
						int gobDstX = j * GobXTexelCount + gobX;
						int gobDstY = (i * blockHeight + k) * GobYTexelCount + gobY;
						int gobDstLinPos = gobDstY * blockCountX * TexelByteSize + gobDstX * TexelByteSize;

						data.AsSpan(srcPos, TexelByteSize).CopyTo(newData.AsSpan(gobDstLinPos, TexelByteSize));

						srcPos += TexelByteSize;
					}
				}
			}
		}

		// We need to remove padding from the data now, so that it matches the expected layout for later decoding.
		// The resulting arrays will have extra unused data at the end, but we need to leave it there
		// because TextureConverter expects the data to be at least texture.CompleteImageSize bytes long.

		if (blockSize.Height > 1)
		{
			int dstBlockCountX = CeilDivide(texture.Width_C28, blockSize.Width);
			if (dstBlockCountX != blockCountX)
			{
				byte[] croppedImageData = new byte[newData.Length];
				int dstBlockCountY = CeilDivide(texture.Height_C28, blockSize.Height);
				for (int y = 0; y < dstBlockCountY; y++)
				{
					newData.AsSpan(y * blockCountX * TexelByteSize, dstBlockCountX * TexelByteSize).CopyTo(croppedImageData.AsSpan(y * dstBlockCountX * TexelByteSize));
				}
				return croppedImageData;
			}

			return newData;
		}

		int pixelSize = TexelByteSize / (blockSize.Width * blockSize.Height);
		if (newData.Length <= texture.Width_C28 * texture.Height_C28 * pixelSize)
		{
			return newData;
		}

		byte[] croppedData = new byte[newData.Length];
		int rowBytesSrc = paddedSize.Width * pixelSize;
		int rowBytesDst = texture.Width_C28 * pixelSize;
		for (int y = 0; y < texture.Height_C28; y++)
		{
			newData.AsSpan(y * rowBytesSrc, rowBytesDst).CopyTo(croppedData.AsSpan(y * rowBytesDst, rowBytesDst));
		}

		if (realFormat == texture.Format_C28E)
		{
			return croppedData;
		}

		// Convert RGBA to RGB
		Debug.Assert(pixelSize % 4 == 0);
		int actualPixelSize = pixelSize / 4 * 3;
		byte[] finalData = new byte[croppedData.Length];
		for (int y = 0; y < texture.Height_C28; y++)
		{
			int srcOffset = y * texture.Width_C28 * pixelSize;
			int dstOffset = y * texture.Width_C28 * actualPixelSize;
			for (int x = 0; x < texture.Width_C28; x++)
			{
				int srcIndex = srcOffset + x * pixelSize;
				int dstIndex = dstOffset + x * actualPixelSize;
				croppedData.AsSpan(srcIndex, actualPixelSize).CopyTo(finalData.AsSpan(dstIndex, actualPixelSize));
			}
		}
		return finalData;
	}

	/// <remarks>
	/// Some of these are guesses based on others and that texels are always 16 bytes in size.
	/// </remarks>
	private static Size? GetTextureFormatBlockSize(TextureFormat textureFormat) => textureFormat switch
	{
		TextureFormat.Alpha8 => new Size(16, 1),
		TextureFormat.ARGB4444 => new Size(8, 1),
		TextureFormat.RGBA4444 => new Size(8, 1),
		TextureFormat.ARGB32 => new Size(4, 1),
		TextureFormat.BGRA32_14 or TextureFormat.BGRA32_37 => new Size(4, 1),
		TextureFormat.RGB565 => new Size(8, 1),
		TextureFormat.DXT1 => new Size(8, 4),
		TextureFormat.DXT5 => new Size(4, 4),
		TextureFormat.BC4 => new Size(8, 4),
		TextureFormat.BC5 => new Size(4, 4),
		TextureFormat.BC6H => new Size(4, 4),
		TextureFormat.BC7 => new Size(4, 4),
		TextureFormat.ASTC_RGB_4x4 or TextureFormat.ASTC_RGBA_4x4 or TextureFormat.ASTC_HDR_4x4 => new Size(4, 4),
		TextureFormat.ASTC_RGB_5x5 or TextureFormat.ASTC_RGBA_5x5 or TextureFormat.ASTC_HDR_5x5 => new Size(5, 5),
		TextureFormat.ASTC_RGB_6x6 or TextureFormat.ASTC_RGBA_6x6 or TextureFormat.ASTC_HDR_6x6 => new Size(6, 6),
		TextureFormat.ASTC_RGB_8x8 or TextureFormat.ASTC_RGBA_8x8 or TextureFormat.ASTC_HDR_8x8 => new Size(8, 8),
		TextureFormat.ASTC_RGB_10x10 or TextureFormat.ASTC_RGBA_10x10 or TextureFormat.ASTC_HDR_10x10 => new Size(10, 10),
		TextureFormat.ASTC_RGB_12x12 or TextureFormat.ASTC_RGBA_12x12 or TextureFormat.ASTC_HDR_12x12 => new Size(12, 12),
		TextureFormat.R8 or TextureFormat.R8_SIGNED => new Size(16, 1),
		TextureFormat.RG16 or TextureFormat.RG16_SIGNED => new Size(8, 1),
		TextureFormat.RGBA32 or TextureFormat.RGBA32_SIGNED => new Size(4, 1),
		TextureFormat.R16 or TextureFormat.R16_SIGNED or TextureFormat.RHalf => new Size(8, 1),
		TextureFormat.RG32 or TextureFormat.RG32_SIGNED or TextureFormat.RGHalf => new Size(4, 1),
		TextureFormat.RGBA64 or TextureFormat.RGBA64_SIGNED or TextureFormat.RGBAHalf => new Size(2, 1),
		TextureFormat.RFloat => new Size(4, 1),
		TextureFormat.RGFloat => new Size(2, 1),
		TextureFormat.RGBAFloat => new Size(1, 1),
		_ => null,
	};

	private static Size GetPaddedTextureSize(int width, int height, int blockWidth, int blockHeight, int gobsPerBlock)
	{
		width = CeilDivide(width, blockWidth * GobXTexelCount) * blockWidth * GobXTexelCount;
		height = CeilDivide(height, blockHeight * GobYTexelCount * gobsPerBlock) * blockHeight * GobYTexelCount * gobsPerBlock;
		return new Size(width, height);
	}

	private static int GetBlockHeightByPlatformBlob(byte[] platformBlob, EndianType endianType)
	{
		return 1 << new EndianSpanReader(platformBlob.AsSpan(8), endianType).ReadInt32();
	}

	private static TextureFormat GetCorrectedSwitchTextureFormat(TextureFormat format) => format switch
	{
		TextureFormat.RGB24 => TextureFormat.RGBA32,
		TextureFormat.RGB48 => TextureFormat.RGBA64,
		TextureFormat.RGB24_SIGNED => TextureFormat.RGBA32_SIGNED,
		TextureFormat.RGB48_SIGNED => TextureFormat.RGBA64_SIGNED,
		_ => format
	};
}