// =====================================================================================================================
//  fennec, a free and open source game engine
//  Copyright (C) 2025  Medusa Slockbower
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <https://www.gnu.org/licenses/>.
// =====================================================================================================================

///
/// \file allocator.h
/// \brief This header contains structures and classes related to allocating blocks of memory
///
///
/// \details
/// \author Medusa Slockbower
///
/// \copyright Copyright © 2025 Medusa Slockbower ([GPLv3](https://www.gnu.org/licenses/gpl-3.0.en.html))
///
///


#ifndef FENNEC_MEMORY_ALLOCATOR_H
#define FENNEC_MEMORY_ALLOCATOR_H

#include <fennec/memory/ptr_traits.h>
#include <fennec/memory/new.h>

#include <fennec/lang/conditional_types.h>
#include <fennec/lang/numeric_transforms.h>
#include <fennec/lang/types.h>
#include <fennec/lang/type_traits.h>

#include <fennec/math/ext/constants.h>
#include <fennec/math/common.h>


#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wchanges-meaning"
#endif


namespace fennec
{

///
/// \brief Helper structure for obtaining traits of an allocator class
/// \tparam Alloc The allocator class to analyze
template<class Alloc>
struct allocator_traits
{
private:
	// These help with using concepts in `detect_t`
	template<typename ClassT> using __pointer            = typename ClassT::pointer_t;
	template<typename ClassT> using __const_pointer      = typename ClassT::const_pointer_t;
	template<typename ClassT> using __void_pointer       = typename ClassT::void_pointer_t;
	template<typename ClassT> using __void_const_pointer = typename ClassT::void_const_pointer_t;

	// Propagation Patterns
	template<typename ClassT> using __propagate_on_containter_copy_assignment = typename ClassT::propagate_on_containter_copy_assignment;
	template<typename ClassT> using __propagate_on_containter_move_assignment = typename ClassT::propagate_on_containter_move_assignment;
	template<typename ClassT> using __propagate_on_containter_swap            = typename ClassT::propagate_on_containter_swap;

	template<typename ClassT> using __is_always_equal = typename ClassT::is_always_equal;

	template<typename AllocT, typename TypeT>
    struct __rebind : replace_first_element<AllocT, TypeT> {};

	template<typename AllocT, typename TypeT>
	requires requires { typename AllocT::template rebind<TypeT>::other; }
	struct __rebind<AllocT, TypeT> { using type = typename AllocT::template rebind<TypeT>::other; };

	// This detects AllocT::diff_t if present, otherwise uses the diff_t associated with PtrT
	// It works using SFINAE, 'typename = void' forces the second __diff to be evaluated first when
	// __diff is substituted using only two template arguments. That one uses 'type = typename AllocT::diff_t'
	// however, if it fails, the compiler moves on to the original definition. __size works in the same manner.
	template<typename AllocT, typename PtrT, typename = void>
	struct __diff { using type = typename ptr_traits<PtrT>::diff_t; };

	template<typename AllocT, typename PtrT>
	struct __diff<AllocT, PtrT, void_t<typename AllocT::diff_t>> { using type = typename AllocT::diff_t; };

	template<typename AllocT, typename DiffT, typename = void>
	struct __size : make_unsigned<DiffT> {};

	template<typename AllocT, typename DiffT>
	struct __size<AllocT, DiffT, void_t<typename AllocT::size_t>> { using type = typename AllocT::size_t; };

public:

	/// \brief Alias for the allocator type
	using alloc_t = Alloc;

	/// \brief Alias for the value type of the allocator
	using value_t = typename Alloc::value_t;

	/// \brief Alias for a pointer to the value type. Will use `Alloc::pointer_t` if present
	using pointer_t = detect_t<value_t*, __pointer, Alloc>;

	/// \brief Alias for a const pointer to the value type. Will use `Alloc::const_pointer_t` if present
	using const_pointer_t = detect_t<const value_t*, __const_pointer, Alloc>;

	/// \brief Alias for a pointer to void. Will use `Alloc::void_pointer_t` if present
	using void_pointer_t = detect_t<void*, __void_pointer, Alloc>;

	/// \brief Alias for a const pointer to void. Will use `Alloc::const_void_pointer_t` if present
	using const_void_pointer_t = detect_t<const void*, __void_const_pointer, Alloc>;

	/// \brief Alias for differences between pointers. Will use `Alloc::diff_t` if present
	using diff_t = typename __diff<Alloc, pointer_t>::type;

	/// \brief Alias for the size of allocations. Will use `Alloc::size_t` if present
	using size_t = typename __size<Alloc, pointer_t>::type;

	// TODO: Document propagation
	using propagate_on_container_copy_assignment = detect_t<false_type, __propagate_on_containter_copy_assignment, Alloc>;
	using propagate_on_container_move_assignment = detect_t<false_type, __propagate_on_containter_move_assignment, Alloc>;
	using propagate_on_container_swap = detect_t<false_type, __propagate_on_containter_swap, Alloc>;

	/// \brief Checks if this allocator type is always equal to another allocator of similar type
	using is_always_equal = detect_t<false_type, __is_always_equal, Alloc>;

	/// \brief Rebinds the allocator type to produce an element type of type `TypeT`
	template<typename TypeT> using rebind = typename __rebind<Alloc, TypeT>::type;

	// TODO: allocator_traits static functions
};


///
/// \brief Allocator implementation, uses `new` and `delete` operators.
/// \tparam T The data type to allocate
template<typename T>
class allocator
{
public:
	/// \brief Alias for the data type used for metaprogramming
	using value_t = T;

	/// \brief Metaprogramming utility to rebind an allocator to a different data type
	template<typename R> using rebind = allocator<R>;

	/// \brief Default Constructor
	constexpr allocator() = default;

	/// \brief Default Destructor
	constexpr ~allocator() = default;

	/// \brief Copy Constructor
	constexpr allocator(const allocator&) = default;

	/// \brief Copy Assignment
	constexpr allocator& operator=(const allocator&) = default;


	/// \brief Equality operator
	constexpr bool_t operator==(const allocator&) {
		return true;
	}

	/// \brief Inequality operator
	constexpr bool_t operator!=(const allocator&) {
		return false;
	}

	/// \brief Equality operator for allocators of same type but with different data type
	template<typename U> constexpr bool_t operator==(const allocator<U>&) {
		return true;
	}

	/// \brief Inequality operator for allocators of same type but with different data type
	template<typename U> constexpr bool_t operator!=(const allocator<U>&) {
		return true;
	}

	/// \brief Allocate a block of memory large enough to hold `n` elements of type `T`
	constexpr T* allocate(size_t n) {
		return static_cast<T*>(::operator new(n * sizeof(T)));
	}

	/// \brief Allocate a block of memory large enough to hold `n` elements of type `T`
	constexpr T* allocate(size_t n, align_t align) {
		return static_cast<T*>(::operator new(n * sizeof(T), align));
	}

	/// \brief Deallocate a block of memory with type `T`
	constexpr void deallocate(T* ptr) {
		return ::operator delete(ptr);
	}

	/// \brief Deallocate a block of memory with type `T`
	constexpr void deallocate(T* ptr, align_t align) {
		return ::operator delete(ptr, align);
	}
};


///
/// \brief Allocator implementation, uses `new` and `delete` operators.
/// \tparam T The data type to allocate
template<typename T>
class allocator<T[]>
{
public:
	/// \brief Alias for the data type used for metaprogramming
	using value_t = T;

	/// \brief Metaprogramming utility to rebind an allocator to a different data type
	template<typename R> using rebind = allocator<R>;

	/// \brief Default Constructor
	constexpr allocator() = default;

	/// \brief Default Destructor
	constexpr ~allocator() = default;

	/// \brief Copy Constructor
	constexpr allocator(const allocator&) = default;

	/// \brief Copy Assignment
	constexpr allocator& operator=(const allocator&) = default;


	/// \brief Equality operator
	constexpr bool_t operator==(const allocator&) {
		return true;
	}

	/// \brief Inequality operator
	constexpr bool_t operator!=(const allocator&) {
		return false;
	}

	/// \brief Equality operator for allocators of same type but with different data type
	template<typename U> constexpr bool_t operator==(const allocator<U>&) {
		return true;
	}

	/// \brief Inequality operator for allocators of same type but with different data type
	template<typename U> constexpr bool_t operator!=(const allocator<U>&) {
		return true;
	}

	/// \brief Allocate a block of memory large enough to hold `n` elements of type `T`
	constexpr T* allocate(size_t n) {
		return static_cast<T*>(::operator new[](n * sizeof(T)));
	}

	/// \brief Allocate a block of memory large enough to hold `n` elements of type `T`
	constexpr T* allocate(size_t n, align_t align) {
		return static_cast<T*>(::operator new[](n * sizeof(T), align));
	}

	/// \brief Deallocate a block of memory with type `T`
	constexpr void deallocate(T* ptr) {
		return ::operator delete[](ptr);
	}

	/// \brief Deallocate a block of memory with type `T`
	constexpr void deallocate(T* ptr, align_t align) {
		return ::operator delete[](ptr, align);
	}
};


///
/// \brief Container to hold a memory allocation
/// \tparam T The data type of the allocation
///
/// \details This simply acts as a proxy for allocating memory. It does not call any constructors or
///          initialize any values as if they were the provided data type. Any operations present work
///          only on individual bytes.
template<typename T, class AllocT>
struct allocation
{
public:
	/// \brief alias for the allocator type
	using alloc_t = typename allocator_traits<AllocT>::template rebind<T>;

	/// \brief alias for the data type
	using value_t = T;

	/// \brief size type definition for ptr_traits
	using size_t = size_t;

	/// \brief diff type definition for ptr_traits
	using diff_t = ptrdiff_t;


	// Cosntructors ========================================================================================================

	///
	/// \brief Default Constructor, initializes internal data to `null` and the capacity to `0`
	constexpr allocation() noexcept
		: _data(nullptr), _capacity(0) {
	}

	///
	/// \brief Sized Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes
	/// \param n The number of elements of type `T` to allocate for
	constexpr allocation(size_t n) noexcept
		: _data(nullptr), _capacity(0) {
		allocate(n);
	}

	///
	/// \brief Buffer Copy Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes.
	///        Then, the contents of data are copied into the allocation.
	/// \param data the buffer to copy
	/// \param n the number of elements
	constexpr allocation(const T* data, size_t n)
		: allocation(n) {
		fennec::memcpy(_data, data, n);
	}

	///
	/// \brief Sized Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes
	/// \param n The number of elements of type `T` to allocate for
	constexpr allocation(size_t n, align_t align) noexcept
		: _data(nullptr)
		, _capacity(0)
		, _alignment(zero<align_t>()) {
		allocate(n, align);
	}

	///
	/// \brief Buffer Copy Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes.
	///        Then, the contents of data are copied into the allocation.
	/// \param data the buffer to copy
	/// \param n the number of elements
	constexpr allocation(const T* data, size_t n, align_t align)
		: allocation(n, align) {
		fennec::memcpy(_data, data, n);
	}

	///
	/// \brief Allocator Constructor
	/// \param alloc The allocation object to copy.
	///
	/// \details This constructor should be used when the type `AllocT` needs internal data.
	constexpr allocation(const alloc_t& alloc) noexcept
		: _alloc(alloc)
		, _data(nullptr)
		, _capacity(0) {
	}

	///
	/// \brief Sized Allocator Constructor
	/// \param n The number of elements of type `T` to allocate for
	/// \param alloc The allocation object to copy.
	///
	/// \details This constructor should be used when the type `AllocT` needs internal data.
	constexpr allocation(size_t n, const alloc_t& alloc) noexcept
		: _alloc(alloc)
		, _data(nullptr)
		, _capacity(0) {
		allocate(n);
	}

	///
	/// \brief Buffer Copy Allocator Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes.
	///        Then, the contents of data are copied into the allocation.
	/// \param data the buffer to copy
	/// \param n the number of elements
	/// \param alloc The allocation object to copy.
	///
	/// \details This constructor should be used when the type `AllocT` needs internal data.
	constexpr allocation(const T* data, size_t n, const alloc_t& alloc)
		: allocation(n, alloc) {
		fennec::memcpy(_data, data, n);
	}

	///
	/// \brief Sized Allocator Constructor
	/// \param n The number of elements of type `T` to allocate for
	/// \param alloc The allocation object to copy.
	///
	/// \details This constructor should be used when the type `AllocT` needs internal data.
	constexpr allocation(size_t n, align_t align, const alloc_t& alloc) noexcept
		: _alloc(alloc)
		, _data(nullptr)
		, _capacity(0)
		, _alignment(zero<align_t>()) {
		allocate(n, align);
	}

	///
	/// \brief Buffer Copy Allocator Constructor, initializes the allocation with a block of size `n * sizeof(T)` bytes.
	///        Then, the contents of data are copied into the allocation.
	/// \param data the buffer to copy
	/// \param n the number of elements
	/// \param alloc The allocation object to copy.
	///
	/// \details This constructor should be used when the type `AllocT` needs internal data.
	constexpr allocation(const T* data, size_t n, align_t align, const alloc_t& alloc)
		: allocation(n, align, alloc) {
		fennec::memcpy(_data, data, n);
	}

	///
	/// \brief Copy Constructor, creates an allocation of equal size and performs a byte-wise copy
	/// \param alloc The allocation to copy
	constexpr allocation(const allocation& alloc) noexcept
		: _alloc(alloc._alloc)
		, _data(_alloc.allocate(alloc._capacity))
		, _capacity(alloc._capacity) {
		fennec::memcpy(_data, alloc._data, alloc._capacity * sizeof(T));
	}

	///
	/// \brief Move Constructor, moves the data in `alloc` to the new object and cleans `alloc` so that it
	///        can safely destruct
	/// \param alloc The allocation to move
	constexpr allocation(allocation&& alloc) noexcept
		: _alloc(alloc._alloc)
		, _data(alloc._data)
		, _capacity(alloc._capacity) {
		alloc._data = nullptr; alloc._capacity = 0;
	}


	///
	/// \brief Default Destructor, releases the memory block if still present
	constexpr ~allocation() noexcept {
		if (_data) _alloc.deallocate(_data);
	}


	// Assignment ==========================================================================================================

	///
	/// \brief Copy Assignment Operator
	/// \param alloc the allocation to copy
	/// \returns a reference to `this`
	constexpr allocation& operator=(const allocation& alloc) {
		allocation::allocate(alloc.capacity());
		fennec::memcpy(_data, alloc, size());
		return *this;
	}

	///
	/// \brief Move Assignment Operator
	/// \param alloc the allocation to copy
	/// \returns a reference to `this`
	constexpr allocation& operator=(allocation&& alloc) noexcept {
		// Copy contents
		_alloc = alloc._alloc;
		_data = alloc._data;
		_capacity = alloc._capacity;

		// Cleanup alloc
		alloc._data = nullptr;
		alloc._capacity = 0;

		return *this;
	}


	// Allocation and Deallocation =========================================================================================

	///
	/// \brief Allocate a block of memory for the allocation.
	///        If there is already an allocated block of memory, the previous allocation is released.
	/// \param n The number of elements of type `T` to allocate for
	constexpr void allocate(size_t n, align_t align = zero<align_t>()) noexcept {
		deallocate();

		if (_alignment != zero<align_t>()) {
			_data = _alloc.allocate(_capacity = n, _alignment = align);
		} else {
			_data = _alloc.allocate(_capacity = n);
		}
	}

	///
	/// \brief Release the block of memory.
	constexpr void deallocate() noexcept {
		if (_data) {
			if (_alignment != zero<align_t>()) {
				_alloc.deallocate(_data, _alignment);
			} else {
				_alloc.deallocate(_data);
			}
		}

		_data = nullptr;
		_capacity = 0;
		_alignment = zero<align_t>();
	}

	///
	/// \brief Reallocate the block with a new size.
	///        Contents are copied to the new allocation.
	constexpr void reallocate(size_t n, align_t align = zero<align_t>()) noexcept {
		if (_data == nullptr) {
			allocate(n, align);
		}

		value_t* old = _data;
		_data = _alloc.allocate(n);
		fennec::memcpy(_data, old, min(_capacity, n) * sizeof(T));
		_alloc.deallocate(old);
		_capacity = n;
	}


	// Access ==============================================================================================================

	constexpr value_t& operator[](size_t i) {
		assertd(i < size(), "Array Out of Bounds");
		return _data[i];
	}

	constexpr const value_t operator[](size_t i) const {
		assertd(i < size(), "Array Out of Bounds");
		return _data[i];
	}


// Modification ========================================================================================================

	///
	/// \brief Clear the block of memory, setting all bytes to 0.
	constexpr void clear() noexcept {
		fennec::memset(_data, 0, _capacity * sizeof(T));
	}

	///
	/// \brief Getter for the byte size of the allocation.
	/// \returns the size of the allocation in bytes
	constexpr size_t size() const {
		return _capacity * sizeof(T);
	}

	///
	/// \brief Getter for the number of elements `n` of type `T` that the allocation can hold.
	/// \returns the size of the allocation in elements
	constexpr size_t capacity() const {
		return _capacity;
	}

	///
	/// \brief Getter for the real pointer to the allocated block of memory
	/// \returns Pointer to the allocated memory.
	constexpr value_t* data() {
		return _data;
	}

	///
	/// \brief Getter for the real pointer to the allocated block of memory
	/// \returns Pointer to the allocated memory.
	constexpr const value_t* data() const {
		return _data;
	}

private:
	alloc_t     _alloc; // Allocator object
	value_t*     _data; // Handle for the memory block
	size_t   _capacity; // Capacity of the memory block in elements.
	align_t _alignment; // Alignment information
};

#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif

}

#endif // FENNEC_MEMORY_ALLOCATOR_H