#include <iostream>
#include <bitset>
#include "bitstream.h"

#define min(x,y) (((x) > (y)) ? (y) : (x))

// i.e when stream is 0b12345678 0bABCDEFGH and command is 
// to read 12 bits:
// out = 00000000 00000000 0000EFGH 12345678 -- wrong
// to read 2 bits:
// out = 00000000 00000000 00000000 00000078
// TODO:
// to read 12 bits:
// out = 00000000 00000000 00008765 4321HGFE
// to read 2 bits:
// out = 00000000 00000000 00000000 00000087
int ibitstream::getbits(size_t n) {
    // std::cerr << "To read " << n << " count: " << this->count << " " << this->cache << std::endl;
    int out = 0, read = 0, to_read;

    if (n > 32 || n < 0) {
        throw std::runtime_error("When reading bits from bitstream n must be <= 32");
    }
    
    // read to cache
    if (this->count == -1 || this->count >= 8) {
        if (!(this->is >> cache)) throw std::runtime_error("Stream EOF");
        this->count = 0;
    }
    
    while (n > 0) {
        to_read = min(n, 8 - this->count);
        // std::cerr << "Iter n: " << n << " count: " << this->count << " "
        //     << this->cache << " to_read: " << to_read << " already read: " << read << std::endl;

        // cache & 0b11111000 if count = 3;
        // cache & 0b10000000 if count = 7;
        // cache & 0b11111111 if count = 0, etc;
        uint8_t mask = (((1 << to_read) - 1) << this->count);
        uint8_t chunk = ((cache & mask) >> this->count);

        // todo inverse chunk
        uint8_t inv = 0;
        for (size_t i = 0; i < to_read; i++)
        {
            inv |= ((chunk >> i) & 1) << (to_read - i - 1);
        }

        out <<= to_read; // shift by length of chunk
        out |= inv; // concat with chunk

        // out |= inv << read;

        // std::cerr << "Mask " << std::bitset<8>(mask) << " chunk " << std::bitset<8>(chunk) <<
        //             " inv " << std::bitset<8>(inv) << " out " << std::bitset<32>(out) << std::endl;
        
        this->count += to_read;
        read += to_read;
        n -= to_read;

        if (this->count == 8){
            // read another byte
            if (!(this->is >> cache)) throw std::runtime_error("Stream EOF");
            this->count = 0;
        }
    }
    
    return out;
}

void obitstream::writebits(short bits, size_t n) {
    // std::cerr << "Write: " << std::bitset<16>(bits) << " n: " << n << " count: " << this->count << " cache: " << std::bitset<8>(this->cache) << std::endl;
    int written = 0, to_write;
    while (n > 0) {
        to_write = min(n, 8 - this->count);
        
        // mask:
        // ------- n -----------------------------
        // 00000..00000 1...1111111111 000000..000 == bits
        // - written - | - to_write - | - offset -
        // so offset = n - to_write
        // because n = n - to_write at every step, so written is already present

        uint16_t mask = (((1 << to_write) - 1) << (n - to_write)); 
        uint8_t chunk = (bits & mask) >> (n - to_write);

        // todo inverse bits in chunk
        uint8_t inv = 0;
        for (size_t i = 0; i < to_write; i++)
        {
            inv |= ((chunk >> i) & 1) << (to_write - i - 1);
        }

        // std::cerr << "Chunk " << std::bitset<8>(chunk) << " inv " << std::bitset<8>(inv) << " to_write " 
        //         << to_write << " written " << written << " n " << n << " mask " << std::bitset<16>(mask) 
        //         << " offset " << (n - to_write) << std::endl;

        this->cache |= (inv << this->count);

        this->count += to_write;
        written += to_write;
        n -= to_write;
        
        if (this->count == 8){
            // flush chunk
            // std::cerr << "Flush: " << std::bitset<8>(this->cache) << std::endl;

            os << this->cache;
            os.flush();

            this->count = 0;
            this->cache = 0;
        }
    }
    
}

void obitstream::flush() {
    // std::cerr << "Flush: " << std::bitset<8>(this->cache) << std::endl;
    os << this->cache;
    os.flush();

    this->cache = 0;
    this->count = 0;
}