#include <Utilities/DenseArray.hpp>
#include <Utilities/Utilities.hpp>

uint32_t DenseArray::extractByte(size_t bitpos) const
{
    size_t index = bitpos >> 3;
    size_t shift = bitpos & 7;

    uint32_t a = (uint32_t) bytes[index] << 8;
    uint32_t b = (index + 1) < numBytes ? bytes[index + 1] : 0;

    uint32_t combined = (a | b) << shift;
    return (combined >> 8) & 0xFF;
}

DenseArray::DenseArray(uint8_t bitsPerEntry, uint32_t numEntries) :
    bitsPerEntry(bitsPerEntry),
    numEntries(numEntries)
{
    clear(bitsPerEntry, numEntries);
}

DenseArray::DenseArray(const DenseArray &other) :
    bitsPerEntry(other.bitsPerEntry),
    numEntries(other.numEntries),
    numBytes(other.numBytes)
{
    bytes = new uint8_t[numBytes];
    std::memcpy(bytes, other.bytes, numBytes);
}

DenseArray::DenseArray(DenseArray &&other) :
    bitsPerEntry(other.bitsPerEntry),
    numEntries(other.numEntries),
    bytes(other.bytes),
    numBytes(other.numBytes)
{
    other.bytes = nullptr;
    other.numBytes = 0;
    other.numEntries = 0;
    other.bitsPerEntry = 0;
}

DenseArray::~DenseArray()
{
    delete[] bytes;
}

DenseArray& DenseArray::operator=(const DenseArray& other)
{
    if (this != &other)
    {
        delete[] bytes;

        bitsPerEntry = other.bitsPerEntry;
        numEntries = other.numEntries;
        numBytes = other.numBytes;

        bytes = new uint8_t[numBytes];
        std::memcpy(bytes, other.bytes, numBytes);
    }

    return *this;
}

DenseArray& DenseArray::operator=(DenseArray&& other)
{
    if (this != &other)
    {
        delete[] bytes;

        bitsPerEntry = other.bitsPerEntry;
        numEntries = other.numEntries;
        numBytes = other.numBytes;
        bytes = other.bytes;

        other.bytes = nullptr;
        other.numBytes = 0;
        other.numEntries = 0;
        other.bitsPerEntry = 0;
    }

    return *this;
}

void DenseArray::clear(uint8_t bitsPerEntry, uint32_t numEntries)
{
#ifdef DEBUG_BUILD
    if (bitsPerEntry > 32)
    {
        crash(formatString("DenseArray: invalid bitsPerEntry=$+", (int) bitsPerEntry));
    }
#endif

    this->bitsPerEntry = bitsPerEntry;
    this->numEntries = numEntries;

    delete[] bytes;
    numBytes = (bitsPerEntry * numEntries + 7) / 8;
    bytes = new uint8_t[numBytes];
    std::memset(bytes, 0, numBytes);
}

uint32_t DenseArray::get(size_t index) const
{
#ifdef DEBUG_BUILD
    if (index >= numEntries)
    {
        crash(formatString("DenseArray::get: index ($+) out of bounds (numEntries=$+)", index, numEntries));
    }
#endif

    size_t bitpos = index * bitsPerEntry;
    
    uint32_t result = 0;
    uint32_t bitsLeft = bitsPerEntry;

    while (bitsLeft >= 8)
    {
        result = (result << 8) | extractByte(bitpos);
        bitsLeft -= 8;
        bitpos += 8;
    }

    if (bitsLeft != 0)
    {
        uint32_t piece = extractByte(bitpos);
        piece >>= (8 - bitsLeft);

        result <<= bitsLeft;
        result |= piece;
    }

    return result;
}

void DenseArray::set(size_t index, uint32_t value)
{
    if (value >= (1ULL << bitsPerEntry))
    {
        size_t newBitsPerEntry = bitsPerEntry;
        while (value >= (1ULL << newBitsPerEntry))
        {
            newBitsPerEntry++;
        }

        transform(newBitsPerEntry, [](uint32_t a) { return a; });
    }

    size_t bitpos = index * bitsPerEntry;
    size_t bitsLeft = bitsPerEntry;

    while (bitsLeft > 0)
    {
        size_t byteIndex = bitpos / 8;
        size_t bitOffset = bitpos % 8;
        size_t bitsInCurrentByte = std::min(bitsLeft, 8 - bitOffset);

        uint8_t mask = ((1 << bitsInCurrentByte) - 1) << (8 - bitOffset - bitsInCurrentByte);

        uint8_t bitsToWriteValue = (value >> (bitsLeft - bitsInCurrentByte)) << (8 - bitOffset - bitsInCurrentByte);

        bytes[byteIndex] = (bytes[byteIndex] & ~mask) | bitsToWriteValue;

        bitpos += bitsInCurrentByte;
        bitsLeft -= bitsInCurrentByte;
    }
}

uint8_t DenseArray::getBitsPerEntry() const
{
    return bitsPerEntry;
}

uint32_t DenseArray::getNumEntries() const
{
    return numEntries;
}

void DenseArray::transform(size_t newBitsPerEntry, std::function<uint32_t(uint32_t)> transformFn)
{
    size_t newNumBytes = (newBitsPerEntry * numEntries + 7) / 8;
    uint8_t* newBytes = new uint8_t[newNumBytes];
    std::memset(newBytes, 0, newNumBytes);

    size_t newBitPos = 0;

    for (size_t i = 0; i < numEntries; i++)
    {
        uint32_t currentValue = transformFn(get(i));

        size_t bitsToWrite = newBitsPerEntry;

        while (bitsToWrite > 0)
        {
            size_t byteIndex = newBitPos / 8;
            size_t bitOffset = newBitPos % 8;
            size_t bitsInCurrentByte = std::min(bitsToWrite, 8 - bitOffset);

            uint8_t mask = ((1 << bitsInCurrentByte) - 1) << (8 - bitOffset - bitsInCurrentByte);
            uint8_t bitsToWriteValue = (currentValue >> (bitsToWrite - bitsInCurrentByte)) << (8 - bitOffset - bitsInCurrentByte);

            newBytes[byteIndex] = (newBytes[byteIndex] & ~mask) | bitsToWriteValue;

            newBitPos += bitsInCurrentByte;
            bitsToWrite -= bitsInCurrentByte;
        }
    }

    delete[] bytes;
    bytes = newBytes;
    numBytes = newNumBytes;
    bitsPerEntry = newBitsPerEntry;
}

void DenseArray::push(uint32_t value)
{
    uint32_t newIndex = numEntries++;

    uint32_t newNumBytes = (bitsPerEntry * numEntries + 7) / 8;
    if (numBytes != newNumBytes)
    {
        uint8_t *newBytes = new uint8_t[newNumBytes];
        std::memcpy(newBytes, bytes, numBytes);
        std::memset(newBytes + numBytes, 0, newNumBytes - numBytes);
        delete[] bytes;
        bytes = newBytes;
        numBytes = newNumBytes;
    }

    set(newIndex, value);
}

std::optional<size_t> DenseArray::find(uint32_t value) const
{
    for (size_t i = 0; i < numEntries; i++)
    {
        if (get(i) == value)
        {
            return i;
        }
    }

    return std::nullopt;
}

bool DenseArray::operator==(const DenseArray &other) const
{
    if (bitsPerEntry != other.bitsPerEntry || numEntries != other.numEntries || numBytes != other.numBytes)
    {
        return false;
    }

    return std::memcmp(bytes, other.bytes, numBytes) == 0;
}