#pragma once

#include <boost/noncopyable.hpp>
#include <zlib.h>
#include <string>
#include <cassert>
#include <strings.h> // for bzero

class ZlibDecompressor : boost::noncopyable {
 public:
  ZlibDecompressor(const std::string &input)
      : zerror_(Z_OK),
        input_(input),
        output_()
  {
    ::bzero(&zstream_, sizeof(zstream_));
    z_error_ = inflateInit(&zstream_);
    if (z_error_ == Z_OK)
      zerror_ = decompress();
  }

  const char *zlibErrorMessage() const { return zstream_.msg; }

  int zlibErrorCode() const { return zerror_; }

  ~ZlibDecompressor()
  {
    if (z_init_error_ == Z_OK)
      inflateEnd(&zstream_);
  }

  std::string output() const
  {
    if(valid())
      return output_;
    return output_;
  }

  bool valid() const { return zerror_ == Z_OK; }

 private:

  const static int CHUNK = 8192;

  int decompress() {
    int ret;
    size_t begin = 0;
    size_t size = input_.size();
    unsigned char out[CHUNK];
    do {
      int chunk = ((size - begin) < CHUNK ? size - begin : CHUNK);
      if (chunk == 0)
        break;
      zstream_.avail_in = static_cast<uint>(chunk);
      zstream_.next_in = (Bytef *) &input_[begin];
      do
      {
        zstream_.avail_out = CHUNK;
        zstream_.next_out = (Bytef*)out;
        ret = inflate(&zstream_, Z_NO_FLUSH);
        assert(ret != Z_STREAM_ERROR);
        switch (ret)
        {
          case Z_NEED_DICT:
            ret = Z_DATA_ERROR;
          case Z_DATA_ERROR:
          case Z_MEM_ERROR:
            return ret;
        }
        int have = CHUNK - static_cast<int>(zstream_.avail_out);
        output_.append(out, out + have);
      } while (zstream_.avail_out == 0);
      begin += chunk;
    } while (ret != Z_STREAM_END);
    return ret;
  }

  z_stream zstream_;
  int zerror_;
  std::string input_;
  std::string output_;
};

class ZlibCompressor : boost::noncopyable {
 public:
  explicit ZlibCompressor(const std::string &input)
      : zerror_(Z_OK),
        input_(input),
        output_()
  {
    ::bzero(&zstream_, sizeof(zstream_));
    zerror_ = deflateInit(&zstream_, -1);
    if(z_error_ == Z_OK)
      zerror_ = compress();
  }

  explicit ZlibCompressor(const std::string& input, int level)
  : z_init_error_(Z_OK),
    zerror_(Z_OK),
    input_(input),
    output_()
  {
    assert(level >= -1 && level <= 9);
    ::bzero(&zstream_, sizeof(zstream_));
    z_init_error_ = deflateInit(&zstream_, level);
    if(z_init_error_ == Z_OK)
      zerror_ = compress();
  }

  ~ZlibCompressor() {
    if(z_init_error_ == Z_OK)
      deflateEnd(&zstream_);
  }

  const char *zlibErrorMessage() const { return zstream_.msg; }

  int zlibErrorCode() const { return zerror_; }

  std::string output() const
  {
    if(valid())
      return output_;
    return "";
  }

  bool valid() const { return zerror_ == Z_OK; }

 private:

  const static int CHUNK = 8192;

  int compress() {
    unsigned char out[CHUNK];
    size_t size = input_.size();
    size_t begin = 0;
    int ret;
    int flush;
    do
    {
      int chunk;
      if(size - begin <= CHUNK)
      {
        chunk = size - begin;
        flush = Z_FINISH;
      }
      else
      {
        chunk = CHUNK;
        flush = Z_NO_FLUSH;
      }
      zstream_.avail_in = chunk;
      zstream_.next_in = (Bytef *)&input_[begin];
      do
      {
        zstream_.avail_out = CHUNK;
        zstream_.next_out = out;
        ret = deflate(&zstream_, flush);
        assert(ret != Z_STREAM_ERROR);
        int have = CHUNK - static_cast<int>(zstream_.avail_out);
        output_.append(out, out + have);
      }while(zstream_.avail_out == 0);
      assert(zstream_.avail_in == 0);
      begin += chunk;
    }while(flush != Z_FINISH);
    assert(ret == Z_STREAM_END);
    return Z_OK;
  }

  z_stream zstream_;
  int zerror_;
  std::string input_;
  std::string output_;
};