#include "precomp.h"

#include <stdlib.h>
#include <string.h>
#include "bitstream3.h"

const unsigned char BitStream3::hiMask[] = { 0, 0x80, 0xC0, 0xE0, 0xF0, 0xF8, 0xFC, 0xFE, 0xFF };

BitStream3::BitStream3() :
ownMemory(false),
currentBit(0),
currentByte(0),
buffer(NULL)
{

}

BitStream3::BitStream3(int bitCount)
{
	currentBit = 0;
	currentByte = 0;
	Create(bitCount);
}

BitStream3::~BitStream3()
{
	if (ownMemory)
		free(buffer);
}

void BitStream3::Create(int bitCount)
{
	bufSize = (bitCount + 7) >> 3;
	bitSize = bitCount;
	buffer = (unsigned char*)malloc((bitCount + 7) >> 3);
	ownMemory = true;
	memset(buffer, 0, bufSize);
}

void BitStream3::Attach(void * buf, int bitCount)
{
	buffer = (unsigned char *)buf;
	bitSize = bitCount;
	currentBit = 0;
	currentByte = 0;
}

int BitStream3::Write(const unsigned char * buf, int bitCount)
{
	int i = 0, gap = 8 - (currentBit & 7);

	if (bitCount >= gap)
	{
		if ((gap & 7) != 0)
		{
			buffer[currentByte] = (unsigned char)((buffer[currentByte] << gap) | (buf[0] >> (bitCount < 8 ? bitCount - gap : 8 - gap)));
			currentBit += gap; bitCount -= gap; currentByte++;

			for (; bitCount >= 8; i++, currentByte++, currentBit += 8, bitCount -= 8)
				buffer[currentByte] = (unsigned char)((buf[i] << gap) | (buf[i + 1] >> (8 - gap)));
		}
		else
			for (; bitCount >= 8; currentByte++, i++, currentBit += 8, bitCount -= 8)
				buffer[currentByte] = buf[i];
	}

	if (bitCount > 0)
	{
		buffer[currentByte] = (unsigned char)((buffer[currentByte] << bitCount) | (buf[i] & ((1 << bitCount) - 1)));
		currentBit += bitCount;
	}

	return bitCount;
}

int BitStream3::WriteUInt16(unsigned short val)
{
    unsigned char split[2]; split[0] = (unsigned char)(val >> 8); split[1] = (unsigned char)(val & 0xFF);
    return Write(split, 16);
}

int BitStream3::Read(unsigned char * readBuf, int bitCount)
{
	int i = 0, gap = currentBit & 7;

	if ((gap & 7) != 0)
	{
		for (; bitCount >= 8 + 8 - gap; bitCount -= 8, currentByte++, currentBit += 8, i++)
			readBuf[i] = (unsigned char)((buffer[currentByte] << gap) | (buffer[currentByte + 1] >> (8 - gap)));

		// byte align 
		if (bitCount > 0)
		{
			if (bitCount < 8 - gap)
			{
				readBuf[i] = (unsigned char)((buffer[currentByte] >> (8 - gap - bitCount)) & ((1 << bitCount) - 1));
				currentBit += bitCount;
				bitCount = 0;
			}
			else
			{
				readBuf[i] = (unsigned char)(buffer[currentByte] & ((1 << (8 - gap)) - 1));

				bitCount -= (8 - gap);
				currentBit += (8 - gap);
				currentByte++;

				readBuf[i] <<= (bitCount < gap ? bitCount : gap);
			}
		}
	}
	else
		for (; bitCount >= 8; bitCount -= 8, currentByte++, currentBit += 8, i++)
			readBuf[i] = buffer[currentByte];

	// already aligned
	if (bitCount > 0)
	{
		if (currentBit + 8 <= bitSize) // not last byte
			readBuf[i] |= (unsigned char)((buffer[currentByte] >> (8 - bitCount)) & ((1 << bitCount) - 1));
		else // last byte 
			readBuf[i] = (unsigned char)((buffer[currentByte] >> ((bitSize & 7) - bitCount)) & ((1 << bitCount) - 1));

		currentBit += bitCount;
	}

	return bitCount;
}

int BitStream3::ReadUInt16(unsigned short * val)
{
    unsigned char split[2];

    int result = Read(split, 16);

    *val = (unsigned short)(((unsigned short)split[0] << 8) | (unsigned short)split[1]);

    return result;
}

unsigned char * BitStream3::GetBuffer()
{
	return buffer;
}