#include "memory/AddressSpace.h"
#include "Log.h"
#include "arch/MMU.h"
#include "memory/Heap.h"
#include <luna/CString.h>
#include <luna/ScopeGuard.h>

static constexpr u64 VM_START = ARCH_PAGE_SIZE;
static constexpr u64 VM_END = 0x0000800000000000;

Result<OwnedPtr<AddressSpace>> AddressSpace::try_create()
{
    OwnedPtr<AddressSpace> ptr = TRY(make_owned<AddressSpace>());

    TRY(ptr->create_null_region());
    TRY(ptr->create_default_region());

    ptr->m_directory = TRY(MMU::create_page_directory_for_userspace());

    return move(ptr);
}

Result<void> AddressSpace::create_null_region()
{
    // Create a small region at the start of the address space to prevent anyone from mapping page 0.
    auto* region = TRY(make<VMRegion>());
    region->start = 0;
    region->end = VM_START;
    region->count = 1;
    region->used = true;
    region->persistent = true;
    m_regions.append(region);
    return {};
}

Result<void> AddressSpace::create_default_region()
{
    // Create a free region covering the rest of the address space.
    auto* region = TRY(make<VMRegion>());
    region->start = VM_START;
    region->end = VM_END;
    region->count = (VM_END / ARCH_PAGE_SIZE) - 1;
    region->used = false;
    m_regions.append(region);
    return {};
}

Result<OwnedPtr<AddressSpace>> AddressSpace::clone()
{
    OwnedPtr<AddressSpace> ptr = TRY(make_owned<AddressSpace>());

    for (const auto* region : m_regions)
    {
        auto* new_region = TRY(make<VMRegion>());
        memcpy(new_region, region, sizeof(*region));
        ptr->m_regions.append(new_region);
    }

    ptr->m_directory = TRY(MMU::clone_userspace_page_directory(m_directory));

    return move(ptr);
}

AddressSpace::AddressSpace()
{
}

AddressSpace& AddressSpace::operator=(AddressSpace&& other)
{
    if (&other == this) return *this;

    m_regions.consume([](VMRegion* region) { delete region; });
    if (m_directory) MMU::delete_userspace_page_directory(m_directory);

    m_regions = other.m_regions;
    m_directory = other.m_directory;

    other.m_regions.reset();
    other.m_directory = nullptr;

    return *this;
}

Result<u64> AddressSpace::alloc_region(usize count, bool persistent)
{
    for (auto* region = m_regions.expect_last(); region; region = m_regions.previous(region).value_or(nullptr))
    {
        if (!region->used)
        {
            if (region->count < count) continue;
            if (region->count == count)
            {
                region->used = true;
                region->persistent = persistent;
                u64 address = region->start;
                try_merge_region_with_neighbors(region);
                return address;
            }

            u64 boundary = region->end - (count * ARCH_PAGE_SIZE);

            auto* new_region = TRY(split_region(region, boundary));
            new_region->used = true;
            new_region->persistent = persistent;
            try_merge_region_with_neighbors(new_region);

            return boundary;
        }
    }

    return err(ENOMEM);
}

Result<bool> AddressSpace::set_region(u64 address, usize count, bool used, bool persistent)
{
    if (address >= VM_END) return err(EINVAL);

    u64 end = address + (count * ARCH_PAGE_SIZE);

    for (auto* region : m_regions)
    {
        if (region->end < address) continue;
        if (region->start > end) return false;

        if (region->persistent) return false;
        if (region->used == used)
        {
            if (used) return false;
            continue;
        }

        if (region->start >= address && region->end <= end)
        {
            region->used = used;
            region->persistent = persistent;
            if (region->start == address && region->end == end)
            {
                try_merge_region_with_neighbors(region);
                return true;
            }
            continue;
        }

        if (region->end > end && region->start < address)
        {
            auto* middle_region = TRY(split_region(region, address));
            TRY(split_region(middle_region, end));
            middle_region->used = used;
            middle_region->persistent = persistent;
            return true;
        }

        if (region->start < address)
        {
            bool finished = region->end == end;
            auto* split = TRY(split_region(region, address));
            split->used = used;
            split->persistent = persistent;
            try_merge_region_with_neighbors(split);
            if (!finished) continue;
            return true;
        }

        if (region->end > end)
        {
            TRY(split_region(region, end));
            region->used = used;
            region->persistent = persistent;
            try_merge_region_with_neighbors(region);
            return true;
        }
    }

    return true;
}

void AddressSpace::merge_contiguous_regions(VMRegion* a, VMRegion* b)
{
    a->end = b->end;
    a->count += b->count;
    m_regions.remove(b);
    delete b;
}

void AddressSpace::try_merge_region_with_neighbors(VMRegion* region)
{
    auto prev = m_regions.previous(region);
    if (prev.has_value() && (*prev)->used == region->used && (*prev)->persistent == region->persistent)
    {
        merge_contiguous_regions(*prev, region);
        region = *prev;
    }

    auto next = m_regions.next(region);
    if (next.has_value() && (*next)->used == region->used && (*next)->persistent == region->persistent)
    {
        merge_contiguous_regions(region, *next);
    }
}

Result<VMRegion*> AddressSpace::split_region(VMRegion* parent, u64 boundary)
{
    auto* region = TRY(make<VMRegion>());

    region->start = boundary;
    region->end = parent->end;
    region->count = (region->end - region->start) / ARCH_PAGE_SIZE;
    region->used = parent->used;
    region->persistent = parent->persistent;
    m_regions.add_after(parent, region);

    parent->end = boundary;
    parent->count -= region->count;

    return region;
}

AddressSpace::~AddressSpace()
{
    m_regions.consume([](VMRegion* region) { delete region; });
    if (m_directory) MMU::delete_userspace_page_directory(m_directory);
}