#include "memory_pool_impl.hpp"

#include <ruby.h>

namespace cumo {
namespace internal {

void CheckStatus(cudaError_t status) {
    if (status != 0) {
        throw CUDARuntimeError(status);
    }
}

Memory::Memory(size_t size) : size_(size) {
    if (size_ > 0) {
        CheckStatus(cudaGetDevice(&device_id_));
        CheckStatus(cudaMallocManaged(&ptr_, size_, cudaMemAttachGlobal));
        // std::cout << "cudaMalloc " << ptr_ << std::endl;
    }
}

Memory::~Memory() {
    if (size_ > 0) {
        // std::cout << "cudaFree  " << ptr_ << std::endl;
        cudaError_t status = cudaFree(ptr_);
        // CUDA driver may shut down before freeing memory inside memory pool.
        // It is okay to simply ignore because CUDA driver automatically frees memory.
        if (status != cudaErrorCudartUnloading) {
            CheckStatus(status);
        }
    }
}

std::shared_ptr<Chunk> Split(std::shared_ptr<Chunk>& self, size_t size) {
    assert(self->size_ >= size);
    if (self->size_ == size) {
        return nullptr;
    }

    auto remaining = std::make_shared<Chunk>(self->mem_, self->offset_ + size, self->size_ - size, self->stream_ptr_);
    self->size_ = size;

    if (self->next_) {
        remaining->set_next(std::move(self->next_));
        remaining->next()->set_prev(remaining);
    }
    self->next_ = remaining;
    remaining->set_prev(self);

    return remaining;
}


void Merge(std::shared_ptr<Chunk>& self, std::shared_ptr<Chunk> remaining) {
    assert(remaining != nullptr);
    assert(self->stream_ptr_ == remaining->stream_ptr());
    self->size_ += remaining->size();
    self->next_ = remaining->next();
    if (remaining->next() != nullptr) {
        self->next_->set_prev(self);
    }
}

void SingleDeviceMemoryPool::AppendToFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr) {
    assert(chunk != nullptr && !chunk->in_use());
    int bin_index = GetBinIndex(size);

    std::lock_guard<std::recursive_mutex> lock{mutex_};

    Arena& arena = GetArena(stream_ptr);
    ArenaIndexMap& arena_index_map = GetArenaIndexMap(stream_ptr);
    int arena_index = std::lower_bound(arena_index_map.begin(), arena_index_map.end(), bin_index) - arena_index_map.begin();
    int length = static_cast<int>(arena_index_map.size());
    if (arena_index >= length || arena_index_map.at(arena_index) != bin_index) {
        arena_index_map.insert(arena_index_map.begin() + arena_index, bin_index);
        arena.insert(arena.begin() + arena_index, FreeList{});
    }
    FreeList& free_list = arena[arena_index];
    free_list.emplace_back(chunk);
}

bool SingleDeviceMemoryPool::RemoveFromFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr) {
    assert(chunk != nullptr && !chunk->in_use());
    int bin_index = GetBinIndex(size);

    std::lock_guard<std::recursive_mutex> lock{mutex_};

    Arena& arena = GetArena(stream_ptr);
    ArenaIndexMap& arena_index_map = GetArenaIndexMap(stream_ptr);
    if (arena_index_map.size() == 0) {
        return false;
    }
    int arena_index = std::lower_bound(arena_index_map.begin(), arena_index_map.end(), bin_index) - arena_index_map.begin();
    if (static_cast<size_t>(arena_index) == arena_index_map.size()) {
        // Bin does not exist for the given chunk size.
        return false;
    }
    if (arena_index_map.at(arena_index) != bin_index) {
        return false;
    }
    assert(arena.size() > static_cast<size_t>(arena_index));
    FreeList& free_list = arena[arena_index];
    return EraseFromFreeList(free_list, chunk);
}

intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
    size = GetRoundedSize(size);
    std::shared_ptr<Chunk> chunk = nullptr;

    {
        std::lock_guard<std::recursive_mutex> lock{mutex_};

        // find best-fit, or a smallest larger allocation
        Arena& arena = GetArena(stream_ptr);
        int arena_index = GetArenaIndex(size);
        int arena_length = static_cast<int>(arena.size());
        for (int i = arena_index; i < arena_length; ++i) {
            FreeList& free_list = arena[i];
            if (free_list.empty()) {
                continue;
            }
            chunk = PopFromFreeList(free_list);
            // TODO(sonots): compact_index
            break;
        }
    }

    if (chunk != nullptr) {
        std::shared_ptr<Chunk> remaining = Split(chunk, size);
        if (remaining != nullptr) {
            AppendToFreeList(remaining->size(), remaining, stream_ptr);
        }
    } else {
        // cudaMalloc if a cache is not found
        std::shared_ptr<Memory> mem = nullptr;
        try {
            mem = std::make_shared<Memory>(size);
        } catch (const CUDARuntimeError& e) {
            if (e.status() != cudaErrorMemoryAllocation) {
                throw;
            }
            // Retry after free all free blocks.
            // NOTE: Anotehr retry after GC is done at cumo_cuda_runtime_malloc.
            FreeAllBlocks();
            try {
                mem = std::make_shared<Memory>(size);
            } catch (const CUDARuntimeError& e) {
                if (e.status() != cudaErrorMemoryAllocation) {
                    throw;
                }
                size_t total = size + GetTotalBytes();
                throw OutOfMemoryError(size, total);
            }
        }
        chunk = std::make_shared<Chunk>(mem, 0, size, stream_ptr);
    }

    assert(chunk != nullptr);
    assert(chunk->stream_ptr() == stream_ptr);
    {
        std::lock_guard<std::recursive_mutex> lock{mutex_};

        chunk->set_in_use(true);
        in_use_.emplace(chunk->ptr(), chunk);
    }
    return chunk->ptr();
}

void SingleDeviceMemoryPool::Free(intptr_t ptr, cudaStream_t stream_ptr) {
    std::shared_ptr<Chunk> chunk = nullptr;

    {
        std::lock_guard<std::recursive_mutex> lock{mutex_};

        chunk = in_use_[ptr];
        // assert(chunk != nullptr);
        if (!chunk) return;
        chunk->set_in_use(false);
        in_use_.erase(ptr);
    }

    if (chunk->next() != nullptr && !chunk->next()->in_use()) {
        if (RemoveFromFreeList(chunk->next()->size(), chunk->next(), stream_ptr)) {
            Merge(chunk, chunk->next());
        }
    }
    if (chunk->prev() != nullptr && !chunk->prev()->in_use()) {
        if (RemoveFromFreeList(chunk->prev()->size(), chunk->prev(), stream_ptr)) {
            chunk = chunk->prev();
            Merge(chunk, chunk->next());
        }
    }
    AppendToFreeList(chunk->size(), chunk, stream_ptr);
}

void SingleDeviceMemoryPool::CompactIndex(cudaStream_t stream_ptr, bool free) {
    // need lock ouside this function
    if (!HasArena(stream_ptr)) return;

    Arena new_arena;
    ArenaIndexMap new_arena_index_map;
    Arena& arena = GetArena(stream_ptr);
    ArenaIndexMap& arena_index_map = GetArenaIndexMap(stream_ptr);
    size_t arena_length = arena.size();
    for (size_t arena_index = 0; arena_index < arena_length; ++arena_index) {
        FreeList& free_list = arena[arena_index];
        if (free_list.empty()) {
            continue;
        }
        if (free) {
            FreeList keep_list;
            for (auto chunk : free_list) {
                if (chunk->prev() != nullptr || chunk->next() != nullptr) {
                    keep_list.emplace_back(chunk);
                }
            }
            if (keep_list.size() == 0) {
                continue;
            }
            new_arena_index_map.emplace_back(arena_index_map[arena_index]);
            new_arena.emplace_back(keep_list);
        } else {
            new_arena_index_map.emplace_back(arena_index_map[arena_index]);
            new_arena.emplace_back(free_list);
        }
    }
    if (new_arena.empty()) {
        index_.erase(stream_ptr);
        free_.erase(stream_ptr);
    } else {
        arena_index_map.swap(new_arena_index_map);
        arena.swap(new_arena);
    }
}

// Free all **non-split** chunks in all arenas
void SingleDeviceMemoryPool::FreeAllBlocks() {
    std::lock_guard<std::recursive_mutex> lock{mutex_};

    std::vector<cudaStream_t> keys(free_.size());
    transform(free_.begin(), free_.end(), keys.begin(), [](auto pair) { return pair.first; });
    for (cudaStream_t stream_ptr : keys) {
        CompactIndex(stream_ptr, true);
    }
}

// Free all **non-split** chunks in specified arena
void SingleDeviceMemoryPool::FreeAllBlocks(cudaStream_t stream_ptr) {
    std::lock_guard<std::recursive_mutex> lock{mutex_};

    CompactIndex(stream_ptr, true);
}

size_t SingleDeviceMemoryPool::GetNumFreeBlocks() {
    size_t n = 0;

    std::lock_guard<std::recursive_mutex> lock{mutex_};

    for (auto kv : free_) {
        Arena& arena = kv.second;
        for (auto free_list : arena) {
            n += free_list.size();
        }
    }
    return n;
}

size_t SingleDeviceMemoryPool::GetUsedBytes() {
    size_t size = 0;

    std::lock_guard<std::recursive_mutex> lock{mutex_};

    for (auto kv : in_use_) {
        std::shared_ptr<Chunk>& chunk = kv.second;
        if (chunk) size += chunk->size();
    }
    return size;
}

size_t SingleDeviceMemoryPool::GetFreeBytes() {
    size_t size = 0;

    std::lock_guard<std::recursive_mutex> lock{mutex_};

    for (auto kv : free_) {
        Arena& arena = kv.second;
        for (auto free_list : arena) {
            for (auto chunk : free_list) {
                if (chunk) size += chunk->size();
            }
        }
    }
    return size;
}

} // namespace internal
} // namespace cumo