using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;

namespace Nerfed.Runtime.Graphics;

// FIXME: can we map the transfer buffer instead of maintaining a local pointer

/// <summary>
/// A convenience structure for creating resources and uploading data to them.
///
/// Note that Upload or UploadAndWait must be called after the Create methods for the data to actually be uploaded.
///
/// Note that this structure does not magically keep memory usage down -
/// you may want to stagger uploads over multiple submissions to minimize memory usage.
/// </summary>
public unsafe class ResourceUploader : GraphicsResource
{
	TransferBuffer BufferTransferBuffer;
	TransferBuffer TextureTransferBuffer;

	byte* bufferData;
	uint bufferDataOffset = 0;
	uint bufferDataSize = 1024;

	byte* textureData;
	uint textureDataOffset = 0;
	uint textureDataSize = 1024;

	List<(uint, BufferRegion, bool)> BufferUploads = new List<(uint, BufferRegion, bool)>();
	List<(uint, TextureRegion, bool)> TextureUploads = new List<(uint, TextureRegion, bool)>();

	public ResourceUploader(GraphicsDevice device) : base(device)
	{
		bufferData = (byte*) NativeMemory.Alloc(bufferDataSize);
		textureData = (byte*) NativeMemory.Alloc(textureDataSize);
	}

	// Buffers

	/// <summary>
	/// Creates a Buffer with data to be uploaded.
	/// </summary>
	public Buffer CreateBuffer<T>(Span<T> data, BufferUsageFlags usageFlags) where T : unmanaged
	{
		uint lengthInBytes = (uint) (Marshal.SizeOf<T>() * data.Length);
		Buffer buffer = new Buffer(Device, usageFlags, lengthInBytes);

		SetBufferData(buffer, 0, data, false);

		return buffer;
	}

	/// <summary>
	/// Prepares upload of data into a Buffer.
	/// </summary>
	public void SetBufferData<T>(Buffer buffer, uint bufferOffsetInElements, Span<T> data, bool cycle) where T : unmanaged
	{
		uint elementSize = (uint) Marshal.SizeOf<T>();
		uint offsetInBytes = elementSize * bufferOffsetInElements;
		uint lengthInBytes = (uint) (elementSize * data.Length);

		uint resourceOffset;
		fixed (void* spanPtr = data)
		{
			resourceOffset = CopyBufferData(spanPtr, lengthInBytes);
		}

		BufferRegion bufferRegion = new BufferRegion(buffer, offsetInBytes, lengthInBytes);
		BufferUploads.Add((resourceOffset, bufferRegion, cycle));
	}

	// Textures

	public Texture CreateTexture2D<T>(Span<T> pixelData, uint width, uint height) where T : unmanaged
	{
		Texture texture = Texture.CreateTexture2D(Device, width, height, TextureFormat.R8G8B8A8, TextureUsageFlags.Sampler);
		SetTextureData(texture, pixelData, false);
		return texture;
	}

	/// <summary>
	/// Creates a 2D Texture from compressed image data to be uploaded.
	/// </summary>
	public Texture CreateTexture2DFromCompressed(Span<byte> compressedImageData)
	{
		ImageUtils.ImageInfoFromBytes(compressedImageData, out uint width, out uint height, out uint _);
		Texture texture = Texture.CreateTexture2D(Device, width, height, TextureFormat.R8G8B8A8, TextureUsageFlags.Sampler);
		SetTextureDataFromCompressed(texture, compressedImageData);
		return texture;
	}

	/// <summary>
	/// Creates a 2D Texture from a compressed image stream to be uploaded.
	/// </summary>
	public Texture CreateTexture2DFromCompressed(Stream compressedImageStream)
	{
		long length = compressedImageStream.Length;
		void* buffer = NativeMemory.Alloc((nuint) length);
		Span<byte> span = new Span<byte>(buffer, (int) length);
		compressedImageStream.ReadExactly(span);

		Texture texture = CreateTexture2DFromCompressed(span);

		NativeMemory.Free(buffer);

		return texture;
	}

	/// <summary>
	/// Creates a 2D Texture from a compressed image file to be uploaded.
	/// </summary>
	public Texture CreateTexture2DFromCompressed(string compressedImageFilePath)
	{
		FileStream fileStream = new FileStream(compressedImageFilePath, FileMode.Open, FileAccess.Read);
		return CreateTexture2DFromCompressed(fileStream);
	}

	/// <summary>
	/// Creates a texture from a DDS stream.
	/// </summary>
	public Texture CreateTextureFromDDS(Stream stream)
	{
		using BinaryReader reader = new BinaryReader(stream);
		Texture texture;
		int faces;
		ImageUtils.ParseDDS(reader, out TextureFormat format, out int width, out int height, out int levels, out bool isCube);

		if (isCube)
		{
			texture = Texture.CreateTextureCube(Device, (uint) width, format, TextureUsageFlags.Sampler, (uint) levels);
			faces = 6;
		}
		else
		{
			texture = Texture.CreateTexture2D(Device, (uint) width, (uint) height, format, TextureUsageFlags.Sampler, (uint) levels);
			faces = 1;
		}

		for (int face = 0; face < faces; face += 1)
		{
			for (int level = 0; level < levels; level += 1)
			{
				int levelWidth = width >> level;
				int levelHeight = height >> level;

				int levelSize = ImageUtils.CalculateDDSLevelSize(levelWidth, levelHeight, format);
				void* byteBuffer = NativeMemory.Alloc((nuint) levelSize);
				Span<byte> byteSpan = new Span<byte>(byteBuffer, levelSize);
				stream.ReadExactly(byteSpan);

				TextureRegion textureRegion = new TextureRegion
				{
					TextureSlice = new TextureSlice
					{
						Texture = texture,
						Layer = (uint) face,
						MipLevel = (uint) level
					},
					X = 0,
					Y = 0,
					Z = 0,
					Width = (uint) levelWidth,
					Height = (uint) levelHeight,
					Depth = 1
				};

				SetTextureData(textureRegion, byteSpan, false);

				NativeMemory.Free(byteBuffer);
			}
		}

		return texture;
	}

	/// <summary>
	/// Creates a texture from a DDS file.
	/// </summary>
	public Texture CreateTextureFromDDS(string path)
	{
		FileStream stream = new FileStream(path, FileMode.Open, FileAccess.Read);
		return CreateTextureFromDDS(stream);
	}

	public void SetTextureDataFromCompressed(TextureRegion textureRegion, Span<byte> compressedImageData)
	{
		byte* pixelData = ImageUtils.GetPixelDataFromBytes(compressedImageData, out uint _, out uint _, out uint sizeInBytes);
		Span<byte> pixelSpan = new Span<byte>((void*) pixelData, (int) sizeInBytes);

		SetTextureData(textureRegion, pixelSpan, false);

		ImageUtils.FreePixelData(pixelData);
	}

	public void SetTextureDataFromCompressed(TextureRegion textureRegion, Stream compressedImageStream)
	{
		long length = compressedImageStream.Length;
		void* buffer = NativeMemory.Alloc((nuint) length);
		Span<byte> span = new Span<byte>(buffer, (int) length);
		compressedImageStream.ReadExactly(span);
		SetTextureDataFromCompressed(textureRegion, span);
		NativeMemory.Free(buffer);
	}

	public void SetTextureDataFromCompressed(TextureRegion textureRegion, string compressedImageFilePath)
	{
		FileStream fileStream = new FileStream(compressedImageFilePath, FileMode.Open, FileAccess.Read);
		SetTextureDataFromCompressed(textureRegion, fileStream);
	}

	/// <summary>
	/// Prepares upload of pixel data into a TextureSlice.
	/// </summary>
	public void SetTextureData<T>(TextureRegion textureRegion, Span<T> data, bool cycle) where T : unmanaged
	{
		int elementSize = Marshal.SizeOf<T>();
		uint dataLengthInBytes = (uint) (elementSize * data.Length);

		uint resourceOffset;
		fixed (T* dataPtr = data)
		{
			resourceOffset = CopyTextureData(dataPtr, dataLengthInBytes, Texture.BytesPerPixel(textureRegion.TextureSlice.Texture.Format));
		}

		TextureUploads.Add((resourceOffset, textureRegion, cycle));
	}

	// Upload

	/// <summary>
	/// Uploads all the data corresponding to the created resources.
	/// </summary>
	public void Upload()
	{
		CopyToTransferBuffer();

		CommandBuffer commandBuffer = Device.AcquireCommandBuffer();
		RecordUploadCommands(commandBuffer);
		Device.Submit(commandBuffer);
	}

	/// <summary>
	/// Uploads and then blocks until the upload is finished.
	/// This is useful for keeping memory usage down during threaded upload.
	/// </summary>
	public void UploadAndWait()
	{
		CopyToTransferBuffer();

		CommandBuffer commandBuffer = Device.AcquireCommandBuffer();
		RecordUploadCommands(commandBuffer);
		Fence fence = Device.SubmitAndAcquireFence(commandBuffer);
		Device.WaitForFence(fence);
		Device.ReleaseFence(fence);
	}

	// Helper methods

	private void CopyToTransferBuffer()
	{
		if (BufferUploads.Count > 0)
		{
			if (BufferTransferBuffer == null || BufferTransferBuffer.Size < bufferDataSize)
			{
				BufferTransferBuffer?.Dispose();
				BufferTransferBuffer = new TransferBuffer(Device, TransferBufferUsage.Upload, bufferDataSize);
			}

			Span<byte> dataSpan = new Span<byte>(bufferData, (int) bufferDataSize);
			BufferTransferBuffer.SetData(dataSpan, true);
		}


		if (TextureUploads.Count > 0)
		{
			if (TextureTransferBuffer == null || TextureTransferBuffer.Size < textureDataSize)
			{
				TextureTransferBuffer?.Dispose();
				TextureTransferBuffer = new TransferBuffer(Device, TransferBufferUsage.Upload, textureDataSize);
			}

			Span<byte> dataSpan = new Span<byte>(textureData, (int) textureDataSize);
			TextureTransferBuffer.SetData(dataSpan, true);
		}
	}

	private void RecordUploadCommands(CommandBuffer commandBuffer)
	{
		CopyPass copyPass = commandBuffer.BeginCopyPass();

		foreach ((uint transferOffset, BufferRegion bufferRegion, bool option) in BufferUploads)
		{
			copyPass.UploadToBuffer(
				new TransferBufferLocation(BufferTransferBuffer, transferOffset),
				bufferRegion,
				option
			);
		}

		foreach ((uint transferOffset, TextureRegion textureRegion, bool option) in TextureUploads)
		{
			copyPass.UploadToTexture(
				new TextureTransferInfo(TextureTransferBuffer, transferOffset),
				textureRegion,
				option
			);
		}

		commandBuffer.EndCopyPass(copyPass);

		BufferUploads.Clear();
		TextureUploads.Clear();
		bufferDataOffset = 0;
	}

	private uint CopyBufferData(void* ptr, uint lengthInBytes)
	{
		if (bufferDataOffset + lengthInBytes >= bufferDataSize)
		{
			bufferDataSize = bufferDataOffset + lengthInBytes;
			bufferData = (byte*) NativeMemory.Realloc(bufferData, bufferDataSize);
		}

		uint resourceOffset = bufferDataOffset;

		NativeMemory.Copy(ptr, bufferData + bufferDataOffset, lengthInBytes);
		bufferDataOffset += lengthInBytes;

		return resourceOffset;
	}

	private uint CopyTextureData(void* ptr, uint lengthInBytes, uint alignment)
	{
		textureDataOffset = RoundToAlignment(textureDataOffset, alignment);

		if (textureDataOffset + lengthInBytes >= textureDataSize)
		{
			textureDataSize = textureDataOffset + lengthInBytes;
			textureData = (byte*) NativeMemory.Realloc(textureData, textureDataSize);
		}

		uint resourceOffset = textureDataOffset;

		NativeMemory.Copy(ptr, textureData + textureDataOffset, lengthInBytes);
		textureDataOffset += lengthInBytes;

		return resourceOffset;
	}

	private uint RoundToAlignment(uint value, uint alignment)
	{
		return alignment * ((value + alignment - 1) / alignment);
	}

	// Dispose

	/// <summary>
	/// It is valid to immediately call Dispose after calling Upload.
	/// </summary>
	protected override void Dispose(bool disposing)
	{
		if (!IsDisposed)
		{
			if (disposing)
			{
				BufferTransferBuffer?.Dispose();
				TextureTransferBuffer?.Dispose();
			}

			NativeMemory.Free(bufferData);
		}
		base.Dispose(disposing);
	}
}
