// Copyright 2011 Scott Worley <sworley@chkno.net>
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include <memory>
#include <iostream>
#include <fstream>
#include <cassert>
#include "snappy.h"

// A simple command-line tool for streaming snappy compression.
// Snappy is http://code.google.com/p/snappy , formerly known as Zippy.
//
// The frame protocol used here is just length-prefix.  Lengths are
// encoded as varints.  A more complicated frame protocol is being
// worked out here: http://code.google.com/p/snappy/issues/detail?id=34 .
// This is not that, and is not compatible with any of that.

int compress(std::istream &in, std::ostream &out) {
	const unsigned inbufsize = snappy::kBlockSize;
	char inbuf[inbufsize];
	std::string outbuf;
	while (1) {
		in.read(inbuf, inbufsize);
		size_t actually_read = in.gcount();
		if (actually_read > 0) {
			snappy::Compress(inbuf, actually_read, &outbuf);
			assert(!outbuf.empty());
			
			// Write size prefix as a varint
			size_t varint = outbuf.size();
			while (varint != 0) {
				unsigned char b = varint & 0x7f;
				varint >>= 7;
				if (varint != 0) {
					b |= 0x80;
				}
				out.write((char*)(&b), 1);
				if (out.bad()) {
					std::cerr << "Error writing" << std::endl;
					return 1;
				}
			}
			
			// Write data
			out.write(&*outbuf.begin(), static_cast<std::streamsize>(outbuf.size()));
			if (out.bad()) {
				std::cerr << "Error writing" << std::endl;
				return 1;
			}
		}
		if (actually_read != inbufsize || in.eof() || in.fail()) {
			assert(actually_read != inbufsize);
			assert(in.fail());
			if (in.eof()) {
				return 0;
			} else {
				std::cerr << "Error reading" << std::endl;
				return 1;
			}
		}
	}
}

int uncompress(std::istream &in, std::ostream &out) {
	std::string outbuf;
	const size_t max_block_size = snappy::MaxCompressedLength(snappy::kBlockSize);
	std::auto_ptr<char> inbuf(new char[max_block_size]);
	while (1) {
		// Read a varint that will be the size of the compressed block
		size_t varint = 0;
		size_t multiplicand = 1;
		bool first_byte = true;
		while (1) {
			unsigned char b;
			in.read((char*)(&b), 1);
			size_t actually_read = in.gcount();
			if (actually_read != 1 || in.eof() || in.fail()) {
				assert(actually_read == 0);
				assert(in.fail());
				if (in.eof()) {
					if (first_byte) {
						return 0;
					} else {
						std::cerr << "EOF during frame header" << std::endl;
						return 1;
					}
				} else {
					std::cerr << "Error reading" << std::endl;
					return 1;
				}
			}
			first_byte = false;
			varint += multiplicand * (b & 0x7f); 
			multiplicand <<= 7;
			if ((b & 0x80) == 0) {
				break;
			}
			if (multiplicand == 0) {
				std::cerr << "Overflow while reading frame header.  Corrupted stream?" << std::endl;
				return 1;
			}
		}
		if (varint > max_block_size) {
			std::cerr << "Block size " << varint << " exceeds max block size " << max_block_size << ".  Corrupted stream?" << std::endl;
			return 1;
		}

		// Read data
		in.read(inbuf.get(), varint);
		size_t actually_read = in.gcount();
		if (actually_read != varint || in.eof() || in.fail()) {
			assert(actually_read != varint);
			assert(in.fail());
			if (in.eof()) {
				std::cerr << "Unexpected EOF" << std::endl;
				return 1;
			} else {
				std::cerr << "Error reading" << std::endl;
				return 1;
			}
		}
		snappy::Uncompress(inbuf.get(), actually_read, &outbuf);
		assert(!outbuf.empty());
		out.write(&*outbuf.begin(), static_cast<std::streamsize>(outbuf.size()));
		if (out.bad()) {
			std::cerr << "Error writing" << std::endl;
			return 1;
		}
	}
}

int main(int argc, char *argv[]) {
	const std::string decode_flag("-d");
	if (argc == 1) {
		return compress(std::cin, std::cout);
	} else if (argc == 2) {
		if (argv[1] == decode_flag) {
			return uncompress(std::cin, std::cout);
		}
		std::ifstream in(argv[1]);
		if (!in.fail()) {
			return compress(in, std::cout);
		}
	} else if (argc == 3 && argv[1] == decode_flag) {
		std::ifstream in(argv[1]);
		if (!in.fail()) {
			return uncompress(in, std::cout);
		}
	}
	std::cerr << "usage: snappy [-d] [filename]" << std::endl;
	return 1;
}