using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using static flecs_hub.flecs.Runtime;

namespace gaemstone.Utility;

public interface IAllocator
{
	nint Allocate(int byteCount);
	void Free(nint pointer);
}

public unsafe static class AllocatorExtensions
{
	public static Span<T> Allocate<T>(this IAllocator allocator, int count) where T : unmanaged
		=> new((void*)allocator.Allocate(sizeof(T) * count), count);
	public static void Free<T>(this IAllocator allocator, Span<T> span) where T : unmanaged
		=> allocator.Free((nint)Unsafe.AsPointer(ref span[0]));

	public static Span<T> AllocateCopy<T>(this IAllocator allocator, ReadOnlySpan<T> orig) where T : unmanaged
		{ var copy = allocator.Allocate<T>(orig.Length); orig.CopyTo(copy); return copy; }

	public static ref T Allocate<T>(this IAllocator allocator) where T : unmanaged
		=> ref Unsafe.AsRef<T>((void*)allocator.Allocate(sizeof(T)));
	public static void Free<T>(this IAllocator allocator, ref T value) where T : unmanaged
		=> allocator.Free((nint)Unsafe.AsPointer(ref value));

	public static CString AllocateCString(this IAllocator allocator, string? value)
	{
		if (value == null) return default;
		var bytes = Encoding.UTF8.GetByteCount(value);
		var span  = allocator.Allocate<byte>(bytes + 1);
		Encoding.UTF8.GetBytes(value, span);
		span[^1] = 0; // In case the allocated span is not cleared.
		return new((nint)Unsafe.AsPointer(ref span[0]));
	}

	public static CString AllocateCString(this IAllocator allocator, ReadOnlySpan<byte> utf8)
	{
		var copy = allocator.Allocate<byte>(utf8.Length + 1);
		utf8.CopyTo(copy);
		copy[^1] = 0; // In case the allocated span is not cleared.
		return new((nint)Unsafe.AsPointer(ref copy[0]));
	}
}

public static class TempAllocator
{
	public const int Capacity = 1024 * 1024; // 1 MB
	private static readonly ThreadLocal<ArenaAllocator> _allocator
		= new(() => new(Capacity));

	public static ResetOnDispose Use()
	{
		var allocator = _allocator.Value!;
		return new(allocator, allocator.Used);
	}

	public sealed class ResetOnDispose
		: IAllocator
		, IDisposable
	{
		private readonly ArenaAllocator _allocator;
		private readonly int _start;

		public ResetOnDispose(ArenaAllocator allocator, int start)
			{ _allocator = allocator; _start = start; }

		// IAllocator implementation
		public nint Allocate(int byteCount) => _allocator.Allocate(byteCount);
		public void Free(nint pointer) { /* Do nothing. */ }

		// IDisposable implementation
		public void Dispose() => _allocator.Reset(_start);
	}
}

public class GlobalHeapAllocator
	: IAllocator
{
	public nint Allocate(int byteCount)
		=> Marshal.AllocHGlobal(byteCount);
	public void Free(nint pointer)
		=> Marshal.FreeHGlobal(pointer);
}

public sealed class ArenaAllocator
	: IAllocator
	, IDisposable
{
	private nint _buffer;
	public int Capacity { get; private set; }
	public int Used { get; private set; } = 0;

	public ArenaAllocator(int capacity)
	{
		if (capacity <= 0) throw new ArgumentOutOfRangeException(nameof(capacity));
		_buffer  = Marshal.AllocHGlobal(capacity);
		Capacity = capacity;
	}

	public void Dispose()
	{
		Marshal.FreeHGlobal(_buffer);
		_buffer  = default;
		Capacity = 0;
	}

	public nint Allocate(int byteCount)
	{
		if (_buffer == default) throw new ObjectDisposedException(nameof(ArenaAllocator));
		if (Used + byteCount > Capacity) throw new InvalidOperationException(
			$"Cannot allocate more than {Capacity} bytes with this {nameof(ArenaAllocator)}");
		var ptr = _buffer + Used;
		Used += byteCount;
		return ptr;
	}

	public void Free(nint pointer)
		{ /* Do nothing. */ }

	public unsafe void Reset(int start = 0)
	{
		if (start > Used) throw new ArgumentOutOfRangeException(nameof(start));
		new Span<byte>((void*)(_buffer + start), Used - start).Clear();
		Used = start;
	}
}

public sealed class RingAllocator
	: IAllocator
	, IDisposable
{
	private nint _buffer;
	private int _current = 0;
	public int Capacity { get; private set; }

	public RingAllocator(int capacity)
	{
		if (capacity <= 0) throw new ArgumentOutOfRangeException(nameof(capacity));
		_buffer  = Marshal.AllocHGlobal(capacity);
		Capacity = capacity;
	}

	public void Dispose()
	{
		Marshal.FreeHGlobal(_buffer);
		_buffer  = default;
		Capacity = 0;
	}

	public nint Allocate(int byteCount)
	{
		if (_buffer == default) throw new ObjectDisposedException(nameof(RingAllocator));
		if (byteCount > Capacity) throw new ArgumentOutOfRangeException(nameof(byteCount));
		if (_current + byteCount > Capacity) _current = 0;
		var ptr = _buffer + _current;
		_current += byteCount;
		return ptr;
	}

	public void Free(nint pointer)
		{ /* IGNORE */ }
}