#include "BitStream.h"

#include "Bits.h"

BitStream::~BitStream()
{

}

Byte BitStream::getCurrentByte()
{
    if (mByteOffset < 0)
    {
        Byte buffer;
        readNextByte(buffer);
    }
    return mCurrentByte;
}

void BitStream::write(DWord data)
{
    for(size_t idx=0; idx<sizeof(DWord); idx++)
    {
        writeByte(Bits::getByteN(data, idx));
    }
}

void BitStream::writeWord(Word data)
{
    const auto byte0 = static_cast<Byte>(data >> 8);
    const auto byte1 = static_cast<Byte>((data << 8) >> 8);
    writeByte(byte0);
    writeByte(byte1);
}

int BitStream::getCurrentByteOffset() const
{
    return mByteOffset;
}

size_t BitStream::getCurrentBitOffset() const
{
    return mBitOffset;
}

String BitStream::logLocation()
{
    String ret;
    ret << _s("Byte offset ") << mByteOffset<< _s(" | Bit offset ") << mBitOffset;
    ret << _s(" | Working byte ") << Bits::toString(getCurrentByte()) << _s('\n');
    return ret;
}

String BitStream::logNextNBytes(size_t n) const
{
    String sstr;
    size_t count{0};
    VecBytes bytes;
    peekNextNBytes(n, bytes);
    for(auto byte : bytes)
    {
        sstr << mByteOffset + count << _s(" | ") << Bits::toString(byte) + _s('\n');
        count++;
    }
    return sstr;
}

void BitStream::writeNBits(DWord data, size_t length)
{
    const auto num_left = 8 - mBitOffset;
    const int overshoot = length - num_left;

    if (overshoot > 0)
    {
        Byte lower_bits = Bits::getLowerNBits(data, num_left);
        mCurrentByte |= lower_bits << mBitOffset;

        writeByte(mCurrentByte, false);

        size_t num_bytes = overshoot / 8;
        for (size_t idx=0; idx< num_bytes; idx++)
        {
            mCurrentByte = Bits::getMBitsAtN(static_cast<Byte>(data), overshoot, idx*8 + num_left);
            writeByte(mCurrentByte, false);
        }

        if (const auto remainder = overshoot % 8; remainder > 0)
        {
            mCurrentByte = Bits::getMBitsAtN(static_cast<Byte>(data), remainder, num_bytes*8 + num_left);
            mBitOffset = remainder;
        }
        else
        {
            mCurrentByte = 0;
            mBitOffset = 0;
        }
    }
    else
    {
        mCurrentByte |= (static_cast<Byte>(data) << mBitOffset);
        mBitOffset += length;
        if (mBitOffset == 8)
        {
            writeByte(mCurrentByte, false);
            mCurrentByte = 0;
            mBitOffset = 0;
        }
    }
}

bool BitStream::readNextNBits(size_t n, Byte& buffer)
{
    Byte internal_buffer;
    if (mByteOffset < 0)
    {
        if (!readNextByte(internal_buffer))
        {
            return false;
        }
    }

    int overshoot = n + mBitOffset - 8;
    if (overshoot > 0)
    {
        const auto last_byte = mCurrentByte;
        if (!readNextByte(internal_buffer))
        {
            return false;
        }

        auto num_lower = 8 - mBitOffset;
        const auto lower_bits = Bits::getHigherNBits(last_byte, num_lower);
        const auto higher_bits = Bits::getLowerNBits(mCurrentByte, overshoot);

        buffer = (higher_bits << num_lower) | lower_bits;

        mBitOffset = overshoot;
        return true;
    }
    else
    {
        buffer = Bits::getMBitsAtN(mCurrentByte, n, mBitOffset);
        mBitOffset += n;
        return true;
    }
}

void BitStream::resetOffsets()
{
    mEndByteOffset = mByteOffset;
    mEndBitOffset = mBitOffset;
    mEndByte = mCurrentByte;

    mCurrentByte = 0;
    mByteOffset = -1;
    mBitOffset = 0;
}

void BitStream::reset()
{
    resetOffsets();
    mCurrentByte = 0;
}

void BitStream::flushRemainingBits()
{
    if (mBitOffset > 0)
    {
        writeByte(mCurrentByte, false);
        mBitOffset = 0;
    }
}

Pair<Byte, size_t> BitStream::getRemainingBits() const
{
    return {mEndByte, mEndBitOffset};
}

void BitStream::setChecksumCalculator(AbstractChecksumCalculator* calc)
{
    mChecksumCalculator = calc;
}

void BitStream::clearChecksumCalculator()
{
    mChecksumCalculator = nullptr;
}