#include <luna/Bitmap.h>
#include <luna/CString.h>
#include <luna/Check.h>
#include <luna/Heap.h>

Bitmap::Bitmap()
{
}

Bitmap::Bitmap(void* location, usize size_in_bytes) : m_location((u8*)location), m_size_in_bytes(size_in_bytes)
{
}

void Bitmap::initialize(void* location, usize size_in_bytes)
{
    m_location = (u8*)location;
    m_size_in_bytes = size_in_bytes;
}

Result<void> Bitmap::allocate(usize size_in_bytes)
{
    initialize(TRY(malloc_impl(size_in_bytes)), size_in_bytes);
    return {};
}

Result<void*> Bitmap::resize(usize new_size_in_bytes)
{
    m_location = (u8*)TRY(realloc_impl(m_location, new_size_in_bytes));
    m_size_in_bytes = new_size_in_bytes;
    return (void*)m_location;
}

Result<void> Bitmap::deallocate()
{
    return free_impl(m_location);
}

void* Bitmap::move(void* new_location, usize new_location_size_in_bytes)
{
    if (new_location_size_in_bytes > m_size_in_bytes) memcpy(new_location, m_location, m_size_in_bytes);
    else
        memcpy(new_location, m_location, new_location_size_in_bytes);

    void* const old_location = (void*)m_location;

    m_location = (u8*)new_location;
    m_size_in_bytes = new_location_size_in_bytes;

    return old_location;
}

void Bitmap::set(usize index, bool value)
{
    expect(index < size(), "Bitmap access out of range");
    const u64 byte_index = index / 8;
    const u8 bit_mask = (u8)(0b10000000 >> (index % 8));
    m_location[byte_index] &= (u8)(~bit_mask);
    if (value) { m_location[byte_index] |= bit_mask; }
}

bool Bitmap::get(usize index) const
{
    expect(index < size(), "Bitmap access out of range");
    const usize byte_index = index / 8;
    const u8 bit_mask = (u8)(0b10000000 >> (index % 8));
    return (m_location[byte_index] & bit_mask) > 0;
}

void Bitmap::clear(bool value)
{
    memset(m_location, value_byte(value), m_size_in_bytes);
}

void Bitmap::clear_region(usize start, usize bits, bool value)
{
    expect((start + bits) <= size(), "Bitmap clear out of range");

    if (!bits) return;

    // Set individual bits while not on a byte boundary.
    while ((start % 8) && bits)
    {
        set(start, value);
        start++;
        bits--;
    }

    // Clear out the rest in bytes.
    const usize bytes = bits / 8;

    memset(&m_location[start / 8], value_byte(value), bytes);
    start += bytes * 8;
    bits -= bytes * 8;

    // Set the remaining individual bits.
    while (bits--)
    {
        set(start, value);
        start++;
    }
}

Option<usize> Bitmap::find(bool value, usize begin) const
{
    const usize size = this->size();

    expect(begin < size, "Start index out of range");

    while (begin % 8)
    {
        if (get(begin) == value) return begin;
        begin++;
    }

    if (begin == size) return {};

    usize i = begin / 8;
    const u8 byte_that_does_not_contain_value = value_byte(!value);
    while (i < m_size_in_bytes)
    {
        if (m_location[i] == byte_that_does_not_contain_value)
        {
            i++;
            continue;
        }

        usize index = i * 8;
        for (usize j = 0; j < 8; j++, index++)
        {
            if (get(index) == value) return index;
        }

        // Once we've located a byte that contains the value, we should succeed in finding it.
        unreachable();
    }

    return {};
}

Option<usize> Bitmap::find_and_toggle(bool value, usize begin)
{
    const usize index = TRY(find(value, begin));
    set(index, !value);
    return index;
}

Option<usize> Bitmap::find_region(bool value, usize count, usize begin) const
{
    // FIXME: Optimize this using bit and byte manipulation.
    u64 region_bits_found = 0;
    u64 region_start = 0;

    for (u64 index = begin; index < m_size_in_bytes * 8; index++)
    {
        if (get(index) != value)
        {
            region_bits_found = 0;
            continue;
        }

        if (region_bits_found == 0)
        {
            region_start = index;
            region_bits_found++;
        }
        else
            region_bits_found++;

        if (region_bits_found == count) return region_start;
    }

    return {};
}

Option<usize> Bitmap::find_and_toggle_region(bool value, usize count, usize begin)
{
    const usize index = TRY(find_region(value, count, begin));
    clear_region(index, count, !value);
    return index;
}

bool Bitmap::match_region_impl(usize start, usize bits, bool value)
{
    if (!bits) return true;

    // Match individual bits while not on a byte boundary.
    while ((start % 8) && bits)
    {
        if (get(start) != value) return false;
        start++;
        bits--;
    }

    // Match the rest in bytes.
    const usize bytes = bits / 8;
    const u8 byte_that_contains_only_value = value_byte(value);

    for (usize i = start; i < start + bytes; i += 8)
    {
        if (m_location[i / 8] != byte_that_contains_only_value) return false;
    }

    start += bytes * 8;
    bits -= bytes * 8;

    // Match the remaining individual bits.
    while (bits--)
    {
        if (get(start) != value) return false;
        start++;
    }

    return true;
}

bool Bitmap::match_region(usize start, usize bits, bool value)
{
    expect((start + bits) <= size(), "Bitmap match out of range");

    return match_region_impl(start, bits, value);
}

Result<bool> Bitmap::try_match_region(usize start, usize bits, bool value)
{
    if ((start + bits) > size()) return err(EINVAL);

    return match_region_impl(start, bits, value);
}