#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>

#include <sys/mman.h>

#include "lzma.h"

ssize_t zpm_uncompress_cb(void *buf, size_t bufsize, void *cbdata,
		int (*cb)(void *ud, void *buf, size_t bufsize)) {
	lzma_stream s = LZMA_STREAM_INIT;
	lzma_stream *strm;
	ssize_t bytes = 0;
	
	uint8_t outbuf[BUFSIZ];

	int ret;

	strm = &s;

	ret = lzma_stream_decoder(strm, UINT64_MAX, 0);
	/* The only reasonable error here is LZMA_MEM_ERROR. */
	if (ret != LZMA_OK) {
		fprintf(stderr, "%s", ret == LZMA_MEM_ERROR ? strerror(ENOMEM)
				: "Internal error (bug)");
		return -1;
	}

	strm->avail_in = bufsize;
	strm->next_in = buf;
	strm->avail_out = BUFSIZ;
	strm->next_out = outbuf;

        lzma_action action = LZMA_RUN;

	while (1) {
		ret = lzma_code(strm, action);

		// Write and check write error before checking decoder error.
		// This way as much data as possible gets written to output
		// even if decoder detected an error.
		if (strm->avail_out == 0 || ret != LZMA_OK) {
			size_t avail = BUFSIZ - strm->avail_out;
			ssize_t written = 0;
			uint8_t *start;

			start = outbuf;

			while (avail > 0) {
				written = cb(cbdata, outbuf, avail);
				if (written == -1) {
					/* Wouldn't be a surprise if writing to
					 * stderr would fail too but at least
					 * try to show an error message.
					 */
					return -1;
				}
				avail -= written;
				start += written;
				bytes += written;
			}

			strm->next_out = outbuf;
			strm->avail_out = BUFSIZ;
		}

		if (ret != LZMA_OK) {
			if (ret == LZMA_STREAM_END) {
				// lzma_stream_decoder() already guarantees
				// that there's no trailing garbage.
				assert(strm->avail_in == 0);
				//assert(action == LZMA_FINISH);
				lzma_end(strm);
				return bytes;
			}

			lzma_end(strm);
			const char *msg;
			switch (ret) {
				case LZMA_MEM_ERROR:
					msg = strerror(ENOMEM);
					break;

				case LZMA_FORMAT_ERROR:
					msg = "File format not recognized";
					break;

				case LZMA_OPTIONS_ERROR:
					// FIXME: Better message?
					msg = "Unsupported compression options";
					break;

				case LZMA_DATA_ERROR:
					msg = "File is corrupt";
					break;

				case LZMA_BUF_ERROR:
					msg = "Unexpected end of input";
					break;

				default:
					msg = "Internal error (bug)";
					break;
			}

			fprintf(stderr, "zpmuncompress: %s\n", msg);
			return -1;
		}
	}

}

ssize_t uncompresslzma(void *buf, size_t bufsize, int out) {
	lzma_stream s = LZMA_STREAM_INIT;
	lzma_stream *strm;
	ssize_t bytes = 0;
	
	uint8_t outbuf[BUFSIZ];

	int ret;

	strm = &s;

	ret = lzma_stream_decoder(strm, UINT64_MAX, 0);
	/* The only reasonable error here is LZMA_MEM_ERROR. */
	if (ret != LZMA_OK) {
		fprintf(stderr, "%s", ret == LZMA_MEM_ERROR ? strerror(ENOMEM)
				: "Internal error (bug)");
		return -1;
	}

	strm->avail_in = bufsize;
	strm->next_in = buf;
	strm->avail_out = BUFSIZ;
	strm->next_out = outbuf;

        lzma_action action = LZMA_RUN;

	while (1) {
		ret = lzma_code(strm, action);

		// Write and check write error before checking decoder error.
		// This way as much data as possible gets written to output
		// even if decoder detected an error.
		if (strm->avail_out == 0 || ret != LZMA_OK) {
			size_t avail = BUFSIZ - strm->avail_out;
			ssize_t written = 0;
			uint8_t *start;

			start = outbuf;

			while (avail > 0) {
				written = write(out, outbuf, avail);
				if (written == -1) {
					/* Wouldn't be a surprise if writing to
					 * stderr would fail too but at least
					 * try to show an error message.
					 */
					fprintf(stderr, "Cannot write to output"
							" file stream: %s",
							strerror(errno));
					return -1;
				}
				avail -= written;
				start += written;
				bytes += written;
			}

			strm->next_out = outbuf;
			strm->avail_out = BUFSIZ;
		}

		if (ret != LZMA_OK) {
			if (ret == LZMA_STREAM_END) {
				// lzma_stream_decoder() already guarantees
				// that there's no trailing garbage.
				assert(strm->avail_in == 0);
				//assert(action == LZMA_FINISH);
				lzma_end(strm);
				return bytes;
			}

			lzma_end(strm);
			const char *msg;
			switch (ret) {
				case LZMA_MEM_ERROR:
					msg = strerror(ENOMEM);
					break;

				case LZMA_FORMAT_ERROR:
					msg = "File format not recognized";
					break;

				case LZMA_OPTIONS_ERROR:
					// FIXME: Better message?
					msg = "Unsupported compression options";
					break;

				case LZMA_DATA_ERROR:
					msg = "File is corrupt";
					break;

				case LZMA_BUF_ERROR:
					msg = "Unexpected end of input";
					break;

				default:
					msg = "Internal error (bug)";
					break;
			}

			fprintf(stderr, "zpmuncompress: %s\n", msg);
			return -1;
		}
	}
}