#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include "librng.h"

#ifdef DEBUG
 #include <stdio.h>
 #include <stdlib.h>
 long* results = 0;
 long* results2 = 0;
 unsigned long dMax = 0;
 int pushed[64];
 long readData = 0;
 long readBuff = 0;
 long readBits = 0;
 long validBits = 0;
 long bitsPushed = 0;
 long readExtra = 0;
 int bits = 0;
 
 unsigned long totalBits = 0;
 unsigned long used = 0;
 unsigned long wasted = 0;
 
 unsigned long power(int exp) {
   unsigned long x = 1;
   
   for (int i = 0; i < exp; i++)
     x *= 2;
   
   return x;
 }
 
 void dump_results() {
   fprintf(stderr, "Rounds for each number:\n");
   for (unsigned long i = 0; i < dMax; i++)
     fprintf(stderr, "%li = %li\t", i, results[i]);
   fprintf(stderr, "\n");
   
   fprintf(stderr, "Rounds for each initial number:\n");
   for (unsigned long i = 0; i < power(bits); i++)
     fprintf(stderr, "%li = %li\t", i, results2[i]);
   fprintf(stderr, "\n");
   
   fprintf(stderr, "Rounds for extra bits: total pushed: \t%li\n", bitsPushed);
   for (int i = 0; i < 64; i++)
     fprintf(stderr, "%i = %i\t", i, pushed[i]);
   fprintf(stderr, "\n");
   
   fprintf(stderr, "Device reads = %li\n", readData);
   fprintf(stderr, "Buffer reads = %li\n", readBuff);
   fprintf(stderr, "64bit reads = %li\n", readBits);
   fprintf(stderr, "Valid results = %li\n", validBits);
   fprintf(stderr, "Extra bit reads = %li\n", readExtra);
   fprintf(stderr, "Num bits per sample = %i\n", bits);   
   fprintf(stderr, "Pagesize = %li\n", sysconf(_SC_PAGE_SIZE));
   fprintf(stderr, "Maximum value = %lu\n", dMax);
   fprintf(stderr, "Total bits read = %lu\n", totalBits);
   fprintf(stderr, "Bits output = %li\n", used);
   fprintf(stderr, "Bits wasted = %li\n", wasted);
   //fprintf(stderr, " = %li\n");
 }
#endif

Rng::Rng(char* device, uint64_t max) {
  iRndFD = -1;
  lCursor = 0;
  lBuffer = 0;
  pStart = 0;
  pNext = 0;
  pEnd = 0;
    
  long lPagesize = sysconf(_SC_PAGE_SIZE);
  if (lPagesize % sizeof(uint64_t) != 0)
    error("Pagesize not multiple of uint64_t");
  
  pStart = (uint64_t*)mmap(NULL, lPagesize, PROT_READ | PROT_WRITE, 
			   MAP_SHARED | MAP_ANONYMOUS | MAP_NORESERVE, -1 ,0);
  if (pStart == MAP_FAILED)
    error("mmap");
  
  lBuffer = lPagesize;
  lCursor = 0;
  pNext = pEnd = pStart;
  
  setMax(max);
  setDevice(device);
}

Rng::~Rng() {
  if (iRndFD != -1)
    close(iRndFD);
  
  iRndFD = -1;
  
  munmap(pStart, pEnd - pStart);
  pStart = pEnd = pNext = 0;

#ifdef DEBUG
  dump_results();
  
  if (results)
    free(results);
  if (results2)
    free(results2);
#endif
}

uint64_t Rng::getMax() {
  return lMax;
}

const uint64_t Rng::setMax(const uint64_t max) {
  lMax = max;
  lOutMask = 0;
  cOutBits = 0;
  
  while (lMax > lOutMask) {
    lOutMask = (lOutMask << 1) + 1;
    cOutBits++;
  }

#ifdef DEBUG
  if (results)
    free(results);
  if (results2)
    free(results2);
  results = (long *)malloc((lMax + 1) * sizeof(long));
  results2 = (long *)malloc((power(cOutBits)) * sizeof(long));
  dMax = lMax + 1;
  bits = cOutBits;
#endif
  
  return lMax;
}

int Rng::setDevice(const char* device) {
  if (iRndFD != -1)
    close(iRndFD);
  
  iRndFD = open(device, O_RDONLY);
  if (iRndFD == -1)
    error("Open failed");
  
  return iRndFD;
}
  
uint64_t Rng::getRnd() {
  uint64_t lOut = 0;
  bool valid = 0;
  
  if (iRndFD == -1 || pStart == 0) {
    error("MMAP or Open failed");
    return lOut;
  }
  
  if (lMax == 0) {
    error("Maximun set to zero");
    return lOut;
  }
  
  do {
    if (sRnd.has(cOutBits)) {
      lOut = sRnd.get(cOutBits);
      
#ifdef DEBUG
      wasted += cOutBits;
      readBits++;
#endif
  
      if (lOut <= lMax)
	valid = 1;
      else
	saveBits(lOut);
    }
    else
      readBuffer();
    
  } while (!valid);
  
#ifdef DEBUG
  results[lOut]++;
  validBits++;
  used += cOutBits;
  wasted -= cOutBits;
#endif
  return lOut;
}

void Rng::saveBits(uint64_t save) {
  uint64_t level = 1;
  
  uint64_t max = lOutMask;
  max -= lMax;
  
  save -= lMax;
  save --;
  
  do {
    level <<= 1;
  } while (level <= max);
  
  level >>= 1;
  
  processBits(max, level, save);
}

inline void Rng::processBits(uint64_t max, uint64_t level, uint64_t data) {
  if (data < level) {
    sRnd.set(data, calculateBits(level));

#ifdef DEBUG
    pushed[calculateBits(level)]++;
    bitsPushed++;
    wasted -= calculateBits(level);
#endif
  }
  else {
    max -= level;
    data -= level;
    do {
      level >>= 1;
    } while (level > max);
    
    if (level > 1)
      processBits(max, level, data);
#ifdef DEBUG
    else {
      pushed[0]++;
    }
#endif
  }
}

void Rng::readBuffer() {
  if (pNext == pEnd)
    fillBuffer();
  
  sRnd.fill(pNext++);

#ifdef DEBUG
  readBuff++;
#endif
}

void Rng::fillBuffer() {
  size_t count = 0;
  ssize_t bytes = 0;
  uint8_t* start = (uint8_t *)pStart;

  if (lCursor == lBuffer) {
    pNext = pEnd = pStart;
    lCursor = 0;
  }
  
  do {
    bytes = read(iRndFD, start + lCursor, lBuffer - lCursor);
    if (bytes == -1)
      error("Read interrupted");
    else {
      lCursor += bytes;
      count += bytes;
    }
  } while (count < sizeof(uint64_t) && lBuffer != lCursor);
  
  pEnd = pStart + (lCursor / sizeof(uint64_t));
  
#ifdef DEBUG
  totalBits += count * 8;
  readData++;
#endif
}

inline void Rng::error (const char* str) {
#ifdef DEBUG
  fprintf(stderr, "RNG error: %s\n", str);
#endif
}

BitContainer::BitContainer() {
  pri = (struct BinaryData) {0, 0};
  sec = pri;
}

BitContainer::~BitContainer() {
}

bool BitContainer::has(int8_t bits) {
  if (pri.bits >= bits)
    return true;
  
  if (sec.bits)
    xfer();
  
  if (pri.bits >= bits)
    return true;

  return false;
}

uint64_t Rng::getMask(int8_t bits) {
  uint64_t mask = 0;

  mask = ~mask;

  //64bit shift at once is undefined:
  mask <<= (bits - 1);
  mask <<= 1;
  
  mask = ~mask;
  
  return mask;
}

int8_t Rng::calculateBits(uint64_t level) {
  int8_t bits = 0;
  
  while (level > 1) {
    level >>= 1;
    bits++;
  }
  
  return bits;
}

uint64_t BitContainer::get(int8_t bits) {
  uint64_t data = 0;
  uint64_t mask = Rng::getMask(bits);
  
  data = pri.data & mask;
  pri.data >>= bits;
  pri.bits -= bits;
  
  return data;
}

void BitContainer::xfer() {
  int8_t bits = 0;
  int8_t free = (sizeof(uint64_t) * 8) - pri.bits;
  
  bits = free < sec.bits ? free : sec.bits;
  
  cpy(&pri, &sec, bits);
}

void BitContainer::cpy(struct BinaryData *dest, struct BinaryData *src, int8_t bits) {
  dest->data <<= bits;
  dest->data |= (src->data & Rng::getMask(bits));
  src->data >>= bits;
  
  src->bits -= bits;
  added(dest->bits, bits);
}

void BitContainer::added(int8_t &stored, int8_t bits) {
  stored += bits;
  
  if (stored > ((int8_t)sizeof(uint64_t) * 8))
    stored = sizeof(uint64_t) * 8;
}

int8_t BitContainer::set(uint64_t data, int8_t bits) {
  struct BinaryData src = (struct BinaryData) {data, bits};
  
  int8_t free = (sizeof(uint64_t) * 8) - sec.bits;
  free = free < bits ? free : bits;
  
  cpy(&sec, &src, free);
  
  return free;
}

void BitContainer::fill(uint64_t *data) {
  sec = pri;
  
  pri.data = *data;
  pri.bits = sizeof(uint64_t) * 8;
}