#pragma once

/*
  #define VECTOR_UTILS
  #define BUFFER
*/

#include <vector>
#include <cstdint>
#include <string>
#include <iostream>
#include <fstream>
#include <cstring>

#ifdef BUFFER
class Buffer
{
public:
  uint8_t *buffer = nullptr;
  size_t taken = 0;
  size_t size = 0;
  std::string file_path;

  Buffer(size_t size);
  Buffer();
  ~Buffer();
  bool resize(size_t new_size);
  int add_end(uint8_t *data, size_t data_size);
  int add_middle(uint8_t *data, size_t data_size, size_t index);
  void copy(Buffer &other);

  // Removes data from buffer without checking if it's in the buffer
  void remove_fast(uint8_t *data, size_t data_size);

  void remove(size_t index, size_t data_size);

  int find(uint8_t *data, size_t data_size);

  void remove(uint8_t *data, size_t data_size);

  void clear() { taken = 0; }

  bool save_to_file();
  bool save_to_file(std::string file_path);

  bool load_from_file();
  bool load_from_file(std::string file_path);
};

Buffer::Buffer()
{
  buffer = nullptr;
  size = 0;
  taken = 0;
}

Buffer::Buffer(size_t size)
{
  this->size = size;
  this->buffer = new uint8_t[size];
}

Buffer::~Buffer()
{
  if (buffer)
    delete[] buffer;
}

bool Buffer::resize(size_t new_size)
{
  if (size >= new_size)
    return true;
  if (new_size < size * 2)
    new_size = size * 2;
  uint8_t *new_buffer = (uint8_t *)realloc(buffer, new_size);
  if (!new_buffer)
  {
    printf("Error resizing buffer\n");
    return false;
  }
  buffer = new_buffer;
  size = new_size;
  return true;
}

int Buffer::add_end(uint8_t *data, size_t data_size)
{
  if (taken + data_size > size)
    if (!resize(size + data_size))
      return -1;

  memcpy(buffer + taken, data, data_size);
  taken += data_size;
  return taken - data_size;
}

int Buffer::add_middle(uint8_t *data, size_t data_size, size_t index)
{
  if (taken + data_size > size)
    if (!resize(size + data_size))
      return -1;
  memmove(buffer + index + data_size, buffer + index, taken - index);
  memcpy(buffer + index, data, data_size);
  taken += data_size;
  return index;
}

void Buffer::copy(Buffer &other)
{
  if (other.size > size)
    if (!resize(other.size))
      return;
  memcpy(buffer, other.buffer, other.taken);
  taken = other.taken;
}

void Buffer::remove_fast(uint8_t *data, size_t data_size)
{
  int64_t index = data - buffer;
  if (index < 0 || index > taken || index + data_size > taken)
  {
    printf("Error removing from buffer\n");
    return;
  }
  memmove(buffer + index, buffer + index + data_size, taken - index - data_size);
  taken -= data_size;
}

void Buffer::remove(size_t index, size_t data_size)
{
  if (index + data_size > taken)
  {
    printf("Error removing from buffer\n");
    return;
  }
  memmove(buffer + index, buffer + index + data_size, taken - index - data_size);
  taken -= data_size;
}

int Buffer::find(uint8_t *data, size_t data_size)
{
  for (int i = 0; i < taken; i++)
    if (memcmp(buffer + i, data, data_size) == 0)
      return i;

  return -1;
}

void Buffer::remove(uint8_t *data, size_t data_size)
{
  int index = find(data, data_size);
  if (index == -1)
  {
    printf("Error removing from buffer\n");
    return;
  }
  remove(index, data_size);
}

bool Buffer::save_to_file()
{
  std::ofstream file(this->file_path, std::ios::binary | std::ios::out);
  if (!file.is_open())
  {
    printf("Error saving file\n");
    return false;
  }
  file.write((char *)this->buffer, this->taken);
  file.close();
  return true;
}

bool Buffer::save_to_file(std::string file_path)
{
  this->file_path = file_path;
  return save_to_file();
}

bool Buffer::load_from_file()
{
  std::ifstream file(this->file_path, std::ios::binary | std::ios::in);
  if (!file.is_open())
    return false;

  file.seekg(0, std::ios::end);

  size_t file_size = file.tellg();
  resize(file_size);

  file.seekg(0, std::ios::beg);
  file.read((char *)buffer, size);

  if (file)
    taken = file_size;
  else
    taken = file.gcount();

  file.close();
  return true;
}

bool Buffer::load_from_file(std::string file_path)
{
  this->file_path = file_path;
  return load_from_file();
}

#endif // BUFFER

#ifdef BIT_READER

class BitReader
{
public:
  int k = 8;
  Buffer buffer;
  char x = 0;
  size_t pos = 0;

  int readInt()
  {
    int ret = *(int *)&buffer.buffer[pos];
    pos += 4;
    return ret;
  }

  unsigned short int readShort()
  {
    short ret = *(unsigned short int *)&buffer.buffer[pos];
    pos += 2;
    return ret;
  }

  char readByte()
  {
    x = buffer.buffer[pos];
    pos++;
    return x;
  }

  bool readBit()
  {
    if (k == 8)
    {
      readByte();
      k = 0;
    }
    bool b = (x >> k) & 1;
    k++;
    return b;
  }
};
#endif // BIT_READER

#ifdef BIT_WRITER

class BitWriter
{
public:
  int k = 0;
  Buffer buffer;
  char x = 0;

  void writeByte(uint8_t v)
  {
    buffer.add_end(&v, sizeof(v));
  }

  void writeInt(int v)
  {
    buffer.add_end((uint8_t *)&v, sizeof(v));
  }

  void write16(unsigned short int v)
  {
    buffer.add_end((uint8_t *)&v, sizeof(v));
  }

  void writeBit(bool b)
  {
    if (k == 8)
    {
      writeByte(x);
      k = 0;
      x = 0;
    }
    x ^= (-b ^ x) & (1 << k);
    k++;
  }

  void finish()
  {
    if (k > 0)
      writeByte(x);
  }
};
#endif // BIT_WRITER

#ifdef VECTOR_UTILS
template <typename T>
void print_vec(std::vector<T> &vec, int wrap = 5)
{
  for (size_t i = 0; i < vec.size(); i++)
  {
    printf("%2d ", vec[i]);
    if (i % wrap == wrap - 1)
      printf("\n");
  }
  printf("\n");
}

template <typename T>
void print_matrix(std::vector<std::vector<T>> &vec)
{
  for (size_t i = 0; i < vec.size(); i++)
  {
    for (size_t j = 0; j < vec[i].size(); j++)
    {
      printf("%2d ", vec[i][j]);
    }
    printf("\n");
  }
  printf("\n");
}
#endif // VECTOR_UTILS