/*
	cloop -- compress loopback device support
*/
#include <fcntl.h>
#include <sys/stat.h>

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <zlib.h>
#include <err.h>
#include <sys/param.h>
#include <sys/endian.h>
#include <netinet/in.h>

#define CLOOP_HEADROOM 128
#define CLOOP_PREAMBLE "#!/bin/sh\n" "#V2.0 Format\n" "exit $?\n"

#define DEBUG 0

struct cloop_head {
	char preamble[CLOOP_HEADROOM];
	u_int32_t block_size;
	u_int32_t num_blocks;
};

struct obuf {
	int fd;
	char *buf;
	uint size;
	uint pos;
	off_t file_off;
};

struct cloop_sc {
	int fd;
	uint block_size;
	uint num_blocks;
	off_t zero_off; /* all zero data */
	char *zero_block;
	char *zero_block_c;
	unsigned long zero_block_c_size;
	off_t cur_off;  /* input */
	off_t data_off; /* out */
	struct obuf data, index;
	z_stream zstream;
	off_t zeros;
};


void
cloop_flush_buffer(struct cloop_sc *sc, struct obuf *b) 
{
	if (b->pos == 0)
		return;

	if (pwrite(b->fd, b->buf, (size_t) b->pos, b->file_off) < 0)
		err(1, "pwrite failed");
	b->file_off += b->pos;
	b->pos = 0;
}

__inline void
cloop_slide_index(struct cloop_sc *sc, uint size)
{
	sc->index.pos += size;
	if (sc->index.pos > sc->index.size) {
		err(1, "index buffer overflow");
	}
	if (sc->index.pos >= sc->index.size) {
		cloop_flush_buffer(sc, &sc->index);
	}
}

void
cloop_deflate_init(struct cloop_sc *sc)
{
	int error;

	sc->zstream.zalloc = (alloc_func)0;
	sc->zstream.zfree = (free_func)0;
	sc->zstream.opaque = (voidpf)0;

	sc->zstream.next_out = sc->data.buf;
	sc->zstream.avail_out = sc->data.size;

#if 1
	error = deflateInit(&sc->zstream, Z_DEFAULT_COMPRESSION);
#else
	error = deflateInit(&sc->zstream, Z_BEST_SPEED);
#endif
	if (error != Z_OK) {
		err(1, "deflateInit");
	}
}

void
cloop_deflate(struct cloop_sc *sc, char *buf, size_t size)
{
	int error;

	sc->zstream.next_in = buf;
	sc->zstream.avail_in = size;
	while (sc->zstream.avail_in > 0) {
		error = deflate(&sc->zstream, Z_NO_FLUSH);
		if (error != Z_OK && error != Z_STREAM_END)
			err(1, "deflate");
		sc->data.pos = sc->data.size - sc->zstream.avail_out;
		if (sc->zstream.avail_out == 0) {
			cloop_flush_buffer(sc, &sc->data) ;
			sc->zstream.next_out = sc->data.buf;
			sc->zstream.avail_out = sc->data.size;
		}
	}
}

void
cloop_deflate_finish(struct cloop_sc *sc)
{
	int error;

	for (;;) {
		error = deflate(&sc->zstream, Z_FINISH);
		sc->data.pos = sc->data.size - sc->zstream.avail_out;
		if (sc->zstream.avail_out == 0) {
			cloop_flush_buffer(sc, &sc->data) ;
			sc->zstream.next_out = sc->data.buf;
			sc->zstream.avail_out = sc->data.size;
		}
		if (error == Z_STREAM_END)
			break;
		if (error != Z_OK)
			err(1, "deflate() in cloop_deflate_finish()");
	}
#if DEBUG
	printf("total in: %lu out: %lu -- %lu%%\n",
		sc->zstream.total_in,
		sc->zstream.total_out,
		100 * sc->zstream.total_out/sc->zstream.total_in);
#endif

	sc->data_off += sc->zstream.total_out;
	error = deflateReset(&sc->zstream);
	if (error != Z_OK)
		err(1, "deflateReset");
}



struct cloop_sc *
cloop_create(int fd, off_t size, uint block_size)
{
	struct cloop_sc *sc; 
	struct cloop_head *ch;
	int error;
	size_t iosize;
	char *buf;

	
	/* check seekable */
	error = lseek(fd, 0, SEEK_SET);
	if (error != 0) {
		return (NULL);
	}

	sc = malloc(sizeof(struct cloop_sc));
	if (sc == NULL) {
		err(1, "malloc failed\n");
	}
	sc->block_size = block_size;
	sc->num_blocks = (size + block_size - 1) / block_size;

	iosize = 4*1024*1024;
	buf = malloc(iosize*2);
	if (buf == NULL) {
		err(1, "malloc failed\n");
	}

	/* buffer for header and index */
	sc->index.fd = fd;
	sc->index.buf = buf;
	sc->index.size = iosize;
	bzero(sc->index.buf, sc->index.size);
	sc->index.file_off = 0;
	ch = (struct cloop_head *) sc->index.buf;
	bcopy(CLOOP_PREAMBLE, ch->preamble, sizeof(CLOOP_PREAMBLE));
	ch->block_size = htonl(sc->block_size);
	ch->num_blocks = htonl(sc->num_blocks);
	sc->index.pos = sizeof(struct cloop_head);

	/* buffer for compressed data */
	sc->data.fd = fd;
	sc->data.buf = buf + iosize;
	sc->data.size = iosize;
	bzero(sc->data.buf, sc->data.size);
	sc->data.file_off = sizeof(struct cloop_head)
				+ sizeof(off_t) * (sc->num_blocks + 1);
	sc->data.pos = 0;

	sc->data_off = sc->zero_off = sc->data.file_off;
	sc->cur_off = 0;

	/* build all zero block */
	sc->zero_block = malloc(block_size * 2);
	if (sc->zero_block == NULL) {
		err(1, "malloc failed\n");
	}

	bzero(sc->zero_block, block_size);
	sc->zero_block_c = sc->zero_block + sc->block_size;
	sc->zero_block_c_size = sc->block_size;
	sc->zeros = 0;

	cloop_deflate_init(sc);
#if 0
	cloop_deflate(sc, sc->zero_block, sc->block_size);
	cloop_deflate_finish(sc);
#endif
	error = compress2(sc->zero_block_c, &sc->zero_block_c_size,
	    sc->zero_block, sc->block_size, Z_BEST_COMPRESSION);
	if (error != Z_OK)
		err(1, "compress2");

	return (sc);
}
void
cloop_write_index(struct cloop_sc *sc, off_t offset)
{
	off_t *index;

	index = (off_t *)(sc->index.buf + sc->index.pos);
	*index = htobe64(offset);
	cloop_slide_index(sc, sizeof(off_t));
#if DEBUG
	printf(" index %ju(block:%ju)", offset, sc->cur_off / sc->block_size);
#endif
}

void
cloop_write_data(struct cloop_sc *sc, char *buf, uint size)
{
	uint res;

	sc->data_off += size;
	while (size > 0) {
		res = sc->data.size - sc->data.pos;
		if (size < res)
			res = size;
		bcopy(buf, sc->data.buf + sc->data.pos, res);
		size -= res;
		buf += res;
		sc->data.pos += res;
		if (sc->data.pos == sc->data.size)
			cloop_flush_buffer(sc, &sc->data);
	}
	sc->zstream.next_out = sc->data.buf + sc->data.pos;
	sc->zstream.avail_out = sc->data.size - sc->data.pos;
}
void
cloop_flush_block(struct cloop_sc *sc)
{
	uint block_off, res;

	/* pad zero and flush current block */
	block_off = sc->cur_off % sc->block_size;
	res = sc->block_size - block_off;
	cloop_write_index(sc, sc->data_off);
	if (block_off == 0) {
#if 1 /* XXX do nothing if geom_uzip is patched */
		cloop_write_data(sc, sc->zero_block_c, sc->zero_block_c_size);
		sc->zeros += sc->zero_block_c_size;
#endif
#if DEBUG
		printf(" all-zero");
#endif
	} else {
#if DEBUG
		printf(" zeros(%d)", res);
#endif
		cloop_deflate(sc, sc->zero_block, res);
		cloop_deflate_finish(sc);
	}
	sc->cur_off += res;
#if DEBUG
	printf("\nblock: %ju", sc->cur_off / sc->block_size);
#endif
}

void
cloop_flush(struct cloop_sc *sc)
{
	/* flush  */
	while (sc->cur_off < (off_t)sc->block_size * sc->num_blocks) {
		cloop_flush_block(sc);
	}
	cloop_write_index(sc, sc->data_off);
	cloop_flush_buffer(sc, &sc->index);
	cloop_flush_buffer(sc, &sc->data);

	printf("cloop: %.3lf/%.3lf MBytes (%.1lf%%) written"
		" (index: %.3lf MB, zeros: %.3lf MB)\n",
		(double)sc->data_off/1024/1024,
		(double)sc->cur_off/1024/1024,
		(double)sc->data_off/sc->cur_off * 100,
		((double)sizeof(struct cloop_head)
			+  sizeof(off_t) * (sc->num_blocks + 1))/1024/1024,
		(double)sc->zeros/1024/1024
		);
	free(sc->index.buf);
	free(sc->zero_block);
	deflateEnd(&sc->zstream);
	free(sc);
}


void
cloop_write(struct cloop_sc *sc, off_t off, char *buf, size_t size)
{
	uint cur_block, new_block, block_off, res, pad;
	off_t data_off;

	if (off < sc->cur_off) {
		err(1, "offset number backward! %ju -> %ju",
		    sc->cur_off, off);
	}
	if (off + size > (off_t)sc->block_size * sc->num_blocks) {
		err(1, "data size overflow %ju > %ju",
		    off + size, (off_t)sc->block_size * sc->num_blocks);
	}
	cur_block = (uint)(sc->cur_off / sc->block_size);
	new_block = (uint)(off / sc->block_size);

	/* Padding zero blocks */
	while (cur_block < new_block) {
		cloop_flush_block(sc);
		cur_block = sc->cur_off / sc->block_size;
	}

	pad = off - sc->cur_off;
	if (pad > 0) {
#if DEBUG
		printf(" zeros(%d)", pad);
#endif
		cloop_deflate(sc, sc->zero_block, pad);
		sc->cur_off = off;
	}
	block_off = sc->cur_off % sc->block_size;
	while (size > 0) {
		res = sc->block_size - block_off;
		if (size >= res) {
			data_off = sc->data_off;
#if DEBUG
			printf(" data_res(%d)", res);
#endif
			cloop_deflate(sc, buf, res);
			cloop_deflate_finish(sc);
			cloop_write_index(sc, data_off);
			sc->cur_off += res;
			size -= res;
			buf += res;
#if DEBUG
			printf("\nblock: %ju", sc->cur_off / sc->block_size);
#endif
		} else {
#if DEBUG
			printf(" data_size(%d)", size);
#endif
			cloop_deflate(sc, buf, size);
			sc->cur_off += size;
			buf += size;
			size = 0;
		}
		block_off = 0;
	}
}

#if 0
int
main()
{
	int fd, i;
	char buf[1024];
	struct cloop_sc *sc;
	fd = open("hogeimage", O_CREAT | O_WRONLY, S_IRUSR |  S_IWUSR );
	if (fd < 0) {
		err(1, "open");
	}
	sc = cloop_create(fd, 16*1024*1024, 16*1024);
#if 0
	for (i = 1; i < 64; i ++) {
		memset(buf, i, 256);
		cloop_write(sc, (off_t)256*i*2 + i, buf, 128);
	}
#endif
	cloop_flush(sc);
	close(fd);
}
#endif