/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * 
 * -----------------------------------------------------------------------------
 * File: Allocator.h
 * Description:
 *  This file implements the functions for allocating, freeing and managing memory.
 *  It includes functions to allocate, reallocate, free, copy, move and set memory.
 * 
 * Author: Mineplay
 * -----------------------------------------------------------------------------
 */
#include "../../Include/Hallocy/Core/Allocator.h"
#include "../../Include/Hallocy/Core/Memory.h"

#if defined(_WIN32)
#include <windows.h>

static INIT_ONCE HALLOCY_INIT_ONCE = INIT_ONCE_STATIC_INIT;

static HANDLE hallocy_heap = NULL;
static CRITICAL_SECTION hallocy_critical_section;
#elif defined(__linux__)
#include <unistd.h>
#include <sys/mman.h>
#include <errno.h>
#include <syscall.h>

static const int HALLOCY_FUTEX_WAKE = 0;
static const int HALLOCY_FUTEX_WAIT = 1;

static int hallocy_futex_address = 0;
#endif

typedef struct HallocyMemoryHeader {
    size_t size;
    struct HallocyMemoryHeader *next;
} HallocyMemoryHeader;

static size_t hallocy_medium_memory_freed = 0;
static size_t hallocy_medium_memory_allocated = 0;
static HallocyMemoryHeader *hallocy_medium_memory_bin = NULL;

static _Thread_local size_t hallocy_small_memory_freed = 0;
static _Thread_local size_t hallocy_small_memory_allocated = 0;
static _Thread_local HallocyMemoryHeader *hallocy_small_memory_bin = NULL;

static size_t hallocy_small_allocation_size = 0;
static size_t hallocy_medium_allocation_size = 0;

#if defined(_WIN32)
static BOOL CALLBACK hallocy_initialize_mutex(PINIT_ONCE init_once, PVOID parameter, PVOID *context) {
    (void)init_once;
    (void)parameter;
    (void)context;

    return InitializeCriticalSectionEx(&hallocy_critical_section, 0x00000400, 0);
}
#endif

void *hallocy_allocate(const size_t size, const bool zero_memory) {
    static size_t page_size = 0;
    if (page_size == 0) {
        #if defined(_WIN32)
        SYSTEM_INFO system_info;
        GetSystemInfo(&system_info);

        page_size = system_info.dwPageSize;
        #elif defined(__linux__)
        page_size = sysconf(_SC_PAGE_SIZE);
        #endif

        hallocy_small_allocation_size = page_size;
        hallocy_medium_allocation_size = page_size * 10;
    }

    size_t aligned_size = page_size * (((size + sizeof(HallocyMemoryHeader)) / page_size) + 1);
    HallocyMemoryHeader *memory_pointer = NULL;
    if (aligned_size <= hallocy_small_allocation_size) {
        HallocyMemoryHeader *previous_header = NULL;
        memory_pointer = hallocy_small_memory_bin;
        while (memory_pointer != NULL) {
            if (memory_pointer->size >= aligned_size) {
                hallocy_small_memory_allocated += memory_pointer->size;
                hallocy_small_memory_freed -= memory_pointer->size;

                if (previous_header == NULL) {
                    hallocy_small_memory_bin = hallocy_small_memory_bin->next; 
                } else {
                    previous_header->next = memory_pointer->next;
                }

                memory_pointer->next = NULL;
                if (zero_memory) {
                    hallocy_set_memory(memory_pointer + 1, 0, aligned_size - sizeof(HallocyMemoryHeader));
                }

                return (void*)(memory_pointer + 1);
            }

            previous_header = memory_pointer;
            memory_pointer = memory_pointer->next;
        }

        #if defined(_WIN32)
        if (hallocy_heap == NULL) {
            hallocy_heap = GetProcessHeap();
        }

        if (zero_memory) {
            memory_pointer = HeapAlloc(hallocy_heap, HEAP_ZERO_MEMORY, aligned_size);
        } else {
            memory_pointer = HeapAlloc(hallocy_heap, 0, aligned_size);
        }

        if (memory_pointer == NULL) {
            return NULL;
        }
        #elif defined(__linux__)
        memory_pointer = mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (memory_pointer == MAP_FAILED) {
            return NULL;
        }
        #endif

        hallocy_small_memory_allocated += aligned_size;
    } else if (aligned_size <= hallocy_medium_allocation_size) { 
        #if defined(_WIN32)
        InitOnceExecuteOnce(&HALLOCY_INIT_ONCE, hallocy_initialize_mutex, NULL, NULL);
        EnterCriticalSection(&hallocy_critical_section);
        #elif defined(__linux__)
        bool locked = false;
        while (!locked) {
            if (__sync_bool_compare_and_swap(&hallocy_futex_address, 0, 1)) {
                locked = true;
            } else {
                syscall(SYS_futex, &hallocy_futex_address, HALLOCY_FUTEX_WAIT, 1, NULL, NULL, 0);
            }
        }
        #endif

        HallocyMemoryHeader *previous_header = NULL;
        memory_pointer = hallocy_medium_memory_bin;
        while (memory_pointer != NULL) {
            if (memory_pointer->size >= aligned_size) {
                hallocy_medium_memory_allocated += memory_pointer->size;
                hallocy_medium_memory_freed -= memory_pointer->size;

                if (previous_header == NULL) {
                    hallocy_medium_memory_bin = hallocy_medium_memory_bin->next; 
                } else {
                    previous_header->next = memory_pointer->next;
                }

                #if defined(_WIN32)
                LeaveCriticalSection(&hallocy_critical_section);
                #elif defined(__linux__)
                hallocy_futex_address = 0;
                syscall(SYS_futex, &hallocy_futex_address, HALLOCY_FUTEX_WAKE, 1, NULL, NULL, 0);
                #endif

                memory_pointer->next = NULL;
                if (zero_memory) {
                    hallocy_set_memory(memory_pointer + 1, 0, aligned_size - sizeof(HallocyMemoryHeader));
                }

                return (void*)(memory_pointer + 1);
            }

            previous_header = memory_pointer;
            memory_pointer = memory_pointer->next;
        }

        #if defined(_WIN32)
        if (hallocy_heap == NULL) {
            hallocy_heap = GetProcessHeap();
        }

        if (zero_memory) {
            memory_pointer = HeapAlloc(hallocy_heap, HEAP_ZERO_MEMORY, aligned_size);
        } else {
            memory_pointer = HeapAlloc(hallocy_heap, 0, aligned_size);
        }

        if (memory_pointer == NULL) {
            return NULL;
        }

        hallocy_medium_allocation_size += aligned_size;
        LeaveCriticalSection(&hallocy_critical_section);
        #elif defined(__linux__)
        memory_pointer = mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (memory_pointer == MAP_FAILED) {
            return NULL;
        }

        hallocy_medium_allocation_size += aligned_size;

        hallocy_futex_address = 0;
        syscall(SYS_futex, &hallocy_futex_address, HALLOCY_FUTEX_WAKE, 1, NULL, NULL, 0);
        #endif
    } else {
        #if defined(_WIN32)
        memory_pointer = (HallocyMemoryHeader*)VirtualAlloc(NULL, aligned_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
        if (memory_pointer == NULL) {
            return NULL;
        }
        #elif defined(__linux__)
        memory_pointer = (HallocyMemoryHeader*)mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (memory_pointer == MAP_FAILED) {
            return NULL;
        }
        #endif
    }

    memory_pointer->size = aligned_size;
    memory_pointer->next = NULL;

    return (void*)(memory_pointer + 1);
}

void *hallocy_realloc(void *memory_pointer, const size_t size) {
    if (memory_pointer == NULL) {
        return hallocy_allocate(size, false);
    }

    if (size == 0) {
        hallocy_free(memory_pointer);
        return NULL;
    }

    HallocyMemoryHeader *memory_header = ((HallocyMemoryHeader*)memory_pointer) - 1;
    if (memory_header->size - sizeof(HallocyMemoryHeader) >= size) {
        return memory_pointer;
    }

    void *new_memory = hallocy_allocate(size, false);
    hallocy_copy_memory(new_memory, memory_pointer, memory_header->size);
    hallocy_free(memory_pointer);
    
    return new_memory;
}

HallocyError hallocy_free(void *pointer) {
    if (pointer == NULL) {
        return HALLOCY_ERROR_INVALID_POINTER;
    }

    HallocyMemoryHeader *memory_header = ((HallocyMemoryHeader*)pointer) - 1;
    if (memory_header->size <= hallocy_small_allocation_size) {
        memory_header->next = hallocy_small_memory_bin;
        hallocy_small_memory_bin = memory_header;
        
        hallocy_small_memory_allocated -= memory_header->size;
        hallocy_small_memory_freed += memory_header->size;
        if (hallocy_small_memory_allocated == 0 && hallocy_small_memory_freed > hallocy_medium_allocation_size) {
            HallocyMemoryHeader *previous_header = NULL;
            HallocyMemoryHeader *current_header = hallocy_small_memory_bin;
            while (current_header != NULL) {
                previous_header = current_header;
                current_header = current_header->next;

                #if defined(_WIN32)
                bool result = HeapFree(hallocy_heap, 0, previous_header);
                if (result == false) {
                    DWORD error = GetLastError();
                    switch (error) {
                        case ERROR_INVALID_PARAMETER:
                            return HALLOCY_ERROR_INVALID_PARAM;
        
                        case ERROR_NOT_ENOUGH_MEMORY:
                            return HALLOCY_ERROR_OUT_OF_MEMORY;
        
                        case ERROR_INVALID_HANDLE:
                            return HALLOCY_ERROR_INVALID_POINTER;
        
                        default:
                            return HALLOCY_ERROR_UNKNOWN;
                    }
                }
                #elif defined(__linux__)
                int result = munmap(previous_header, previous_header->size);
                if (result == -1) {
                    switch (errno) {
                        case EINVAL:
                            return HALLOCY_ERROR_INVALID_PARAM;
        
                        case ENOMEM:
                            return HALLOCY_ERROR_OUT_OF_MEMORY;
        
                        case EFAULT:
                            return HALLOCY_ERROR_INVALID_POINTER;
        
                        default:
                            return HALLOCY_ERROR_UNKNOWN;
                    }
                }
                #endif
            }

            hallocy_small_memory_bin = NULL;
            hallocy_small_memory_allocated = 0;
            hallocy_small_memory_freed = 0;
        }
    } else if (memory_header->size <= hallocy_medium_allocation_size) {
        #if defined(_WIN32)
        InitOnceExecuteOnce(&HALLOCY_INIT_ONCE, hallocy_initialize_mutex, NULL, NULL);
        EnterCriticalSection(&hallocy_critical_section);
        #elif defined(__linux__)
        bool locked = false;
        while (!locked) {
            if (__sync_bool_compare_and_swap(&hallocy_futex_address, 0, 1)) {
                locked = true;
            } else {
                syscall(SYS_futex, &hallocy_futex_address, HALLOCY_FUTEX_WAIT, 1, NULL, NULL, 0);
            }
        }
        #endif

        memory_header->next = hallocy_medium_memory_bin;
        hallocy_medium_memory_bin = memory_header;
        
        hallocy_medium_memory_allocated -= memory_header->size;
        hallocy_medium_memory_freed += memory_header->size;
        if (hallocy_medium_memory_allocated == 0 && hallocy_medium_memory_freed > hallocy_medium_allocation_size) {
            HallocyMemoryHeader *previous_header = NULL;
            HallocyMemoryHeader *current_header = hallocy_medium_memory_bin;
            while (current_header != NULL) {
                previous_header = current_header;
                current_header = current_header->next;

                #if defined(_WIN32)
                bool result = HeapFree(hallocy_heap, 0, previous_header);
                if (result == false) {
                    DWORD error = GetLastError();
                    switch (error) {
                        case ERROR_INVALID_PARAMETER:
                            return HALLOCY_ERROR_INVALID_PARAM;
        
                        case ERROR_NOT_ENOUGH_MEMORY:
                            return HALLOCY_ERROR_OUT_OF_MEMORY;
        
                        case ERROR_INVALID_HANDLE:
                            return HALLOCY_ERROR_INVALID_POINTER;
        
                        default:
                            return HALLOCY_ERROR_UNKNOWN;
                    }
                }
                #elif defined(__linux__)
                int result = munmap(previous_header, previous_header->size);
                if (result == -1) {
                    switch (errno) {
                        case EINVAL:
                            return HALLOCY_ERROR_INVALID_PARAM;
        
                        case ENOMEM:
                            return HALLOCY_ERROR_OUT_OF_MEMORY;
        
                        case EFAULT:
                            return HALLOCY_ERROR_INVALID_POINTER;
        
                        default:
                            return HALLOCY_ERROR_UNKNOWN;
                    }
                }
                #endif
            }

            hallocy_medium_memory_bin = NULL;
            hallocy_medium_memory_allocated = 0;
            hallocy_medium_memory_freed = 0;
        }

        #if defined(_WIN32)
        LeaveCriticalSection(&hallocy_critical_section);
        #elif defined(__linux__)
        hallocy_futex_address = 0;
        syscall(SYS_futex, &hallocy_futex_address, HALLOCY_FUTEX_WAKE, 1, NULL, NULL, 0);
        #endif
    } else {
        #if defined(_WIN32)
        bool result = VirtualFree(memory_header, 0, MEM_RELEASE);
        if (result == false) {
            DWORD error = GetLastError();
            switch (error) {
                case ERROR_INVALID_PARAMETER:
                    return HALLOCY_ERROR_INVALID_PARAM;

                case ERROR_NOT_ENOUGH_MEMORY:
                    return HALLOCY_ERROR_OUT_OF_MEMORY;

                case ERROR_INVALID_HANDLE:
                    return HALLOCY_ERROR_INVALID_POINTER;

                case ERROR_INVALID_ADDRESS:
                    return HALLOCY_ERROR_INVALID_POINTER;

                default:
                    return HALLOCY_ERROR_UNKNOWN;
            }
        }
        #elif defined(__linux__)
        int result = munmap(memory_header, memory_header->size);
        if (result == -1) {
            switch (errno) {
                case EINVAL:
                    return HALLOCY_ERROR_INVALID_PARAM;

                case ENOMEM:
                    return HALLOCY_ERROR_OUT_OF_MEMORY;

                case EFAULT:
                    return HALLOCY_ERROR_INVALID_POINTER;

                default:
                    return HALLOCY_ERROR_UNKNOWN;
            }
        }
        #endif
    }

    return HALLOCY_ERROR_NONE;
}