/*
 *	gzfile 0.3
 *
 *	Copyright (C) 2006 Jan Bobrowski <jb@wizard.ae.krakow.pl>
 *
 *	This library is free software; you can redistribute and modify
 *	it under the terms of the GNU Lesser General Public License
 *	version 2.1 as published by the Free Software Foundation.
 */

#define _GNU_SOURCE
#include <stdio.h>

/* provides: */

/**
 * fopenz - open gzip file, or just fopen, if non-gzip.
 * @name: filename or NULL for stdin
 * @mode: should be "r"
 */
FILE *fopenz(const char *name, const char *mode);

/**
 * fgunzip - inserts ungzip filter (if gzip is identified).
 * @f: FILE to ungzip
 */
FILE *fgunzip(FILE *f);

/* Fclose the FILE before exit. */


#if !defined __GLIBC__ && !#system(bsd)
#warning System not supported
FILE *fgunzip(FILE *f) {return f;}
FILE *fopenz(const char *name, const char *mode)
{return fopen(name, mode);}
#else

#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <errno.h>
#include <err.h>
#include <assert.h>

#if !defined WANT_FGUNZIP && !defined WANT_FOPENZ
#define WANT_FGUNZIP
#define WANT_FOPENZ
#endif

#if defined __GNUC__ && !defined PROF
#define LOCAL __attribute__((regparm(2))) static
#else
#define LOCAL static
#endif

typedef unsigned char u8;
typedef unsigned int u32;
typedef signed char s8;

#define elemof(T) (sizeof T/sizeof*T)
#define endof(T) (T+elemof(T))

#ifndef NDEBUG
#define WARN warnx("bad data (%s:%u)",__FILE__,__LINE__)
#else
#define WARN
#endif

/* tree */

#define HTBITS 5
#define HTELEM (1<<HTBITS)
#define HTMASK (HTELEM-1)
#define CLENMAX 15

/* the idea from algorithm.txt (zlib) */
struct htree {
	union {
		unsigned value;
		struct htree *child;
	} u;
	u8 bits;
};

static struct htree *htree_free(struct htree *);
static struct htree *htree_gen(unsigned tl[], unsigned tv[], int n);

static inline struct htree *htree_alloc()
{
	struct htree *t = calloc(HTELEM, sizeof *t);
	return t;
}

static struct htree *htree_free(struct htree *t)
{
	int i;
	for(i=0; i<HTELEM; i++)
		if(!t[i].bits && t[i].u.child)
			htree_free(t[i].u.child);
	free(t);
	return 0;
}

#if 1
#define HTREE_LOOKUP(htab,t,m,b) \
		t = (htab);				\
		for(;;) {				\
			t = &t[m & HTMASK];		\
			if(t->bits) break;		\
			t = t->u.child;			\
			m >>= HTBITS; b += HTBITS;	\
		}					\
		b += t->bits;
#else
#define HTREE_LOOKUP(htab,t,m,b) \
		t = &(htab)[m & HTMASK];		\
		while(!t->bits) {			\
			t = t->u.child;			\
			m >>= HTBITS;			\
			t = &t[m & HTMASK];		\
			b += HTBITS;			\
		}					\
		b += t->bits;
#endif

static void htree_add(struct htree *t, unsigned c, int l, unsigned v)
{
	while(l > HTBITS) {
		struct htree *tp = &t[c & HTMASK];
		assert(!tp->bits);
		t = tp->u.child;
		if(!t)
			t = tp->u.child = htree_alloc();
		c >>= HTBITS;
		l -= HTBITS;
	}
	assert(c < HTELEM);
	do {
		assert(!t[c].bits);
		t[c].bits = l;
		t[c].u.value = v;
		c += 1<<l;
	} while(c < HTELEM);
}

#define CARRY(V) ((V)<<(8*sizeof(V)-1))

static unsigned inv_add(unsigned a, unsigned b)
{
	for(;;) {
		unsigned aa = a;
		a ^= b;
		b &= aa;
		if(!b) break;
		b = CARRY(b) | b>>1; /* rotate right */
	}
	return a;
}

static struct htree *htree_gen(unsigned tl[], unsigned tv[], int n)
{
	struct htree *t;
	unsigned ct[CLENMAX+1] = {0};
	unsigned c;
	int i;

	if(n==0) {
		WARN;
		return 0;
	}

	assert(tv);
	/* XXX TODO: special case when n==1 and tl[0]==1 */
	assert(n>1);

	for(i=0; i<n; i++) {
		unsigned l = tl[i];
		assert(l);
		assert(l<=CLENMAX);
		ct[l] = inv_add(ct[l], 1<<l-1);
	}

	c = 0;
	for(i=1;;) {
		c = inv_add(c, ct[i]);
		ct[i++] = c;
		if(c & CARRY(1u)) {
			if(c != CARRY(1u)) {
				WARN;
				return 0;
			}
			while(i<=CLENMAX) if(ct[i++]) {
				WARN;
				return 0;
			}
			break;
		}
		if(i > CLENMAX) {
			WARN;
			return 0;
		}
	}

	t = htree_alloc();
	for(i=0; i<n; i++) {
		unsigned l = tl[i] - 1;
		htree_add(t, ct[l], l+1, tv[i]);
		ct[l] = inv_add(ct[l], 1<<l);
	}
	return t;
}

/* input */

#define IBUFSZ 4096

struct in {
	u8 *iptr, *iend;
#ifdef WANT_FGUNZIP
	FILE *file;
#endif
#ifdef WANT_FOPENZ
	int fd;
#endif
	u8 *ibuf; // input buffer
};

static inline void in_init(struct in *in)
{
	in->ibuf = malloc(IBUFSZ);
	in->iptr = in->iend = in->ibuf;
#ifdef WANT_FGUNZIP
	in->file = 0;
#endif
#ifdef WANT_FOPENZ
	in->fd = -1;
#endif
}

static inline void in_close(struct in *in)
{
	free(in->ibuf), in->ibuf = 0;
#ifdef WANT_FGUNZIP
	if(in->file) fclose(in->file), in->file = 0;
#endif
#ifdef WANT_FOPENZ
	if(in->fd>=0) close(in->fd), in->fd = -1;
#endif
}

struct zf {
	u32 bits;
	int nbit;

	struct htree *lit, *off;
	int (*pump)(struct zf *zf, int n);

	unsigned end, rdp;

	unsigned final_block:1;
	short blen;
	int errn;

	struct in in;

	u8 hist[32768];
};

static inline void zf_destroy(struct zf *zf)
{
	if(zf->lit) zf->lit = htree_free(zf->lit);
	if(zf->off) zf->off = htree_free(zf->off);
	in_close(&zf->in);
	free(zf);
}

#define error(Z,E) ({WARN; ret_error((Z),(E));})

#define ZF_ERR_TRUNC ESPIPE
#define ZF_ERR_DATA EIO
#define ZF_ERR_TREE EIO

static int ret_error(struct zf *zf, int e)
{
	if(!zf->errn) zf->errn = e;
	zf->pump = 0;
	return -1;
}

static u8 *fill_buf(struct zf *zf)
{
	struct in *in = &zf->in;
	int v;
#if defined WANT_FGUNZIP && defined WANT_FOPENZ
	if(in->file)
#elif defined WANT_FGUNZIP
	if(1)
#else
	if(0)
#endif
	{
#ifdef WANT_FGUNZIP
		v = 0;
		if(!feof(in->file) && !ferror(in->file))
			v = fread(in->ibuf, 1, IBUFSZ, in->file);
#endif
	} else {
#ifdef WANT_FOPENZ
		v = read(in->fd, in->ibuf, IBUFSZ);
		if(v<0) {
			zf->errn = errno;
			v = 0;
		}
#endif
	}
	in->iend = in->ibuf + v;
	return in->ibuf;
}

LOCAL unsigned eatb(struct zf *zf, int ack)
{
	zf->nbit -= ack;
	if(zf->nbit < 0) {
		zf->nbit = 0;
		zf->bits = 0;
		return zf->bits;
	}
	while(zf->nbit <= 24) {
		u8 *p = zf->in.iptr;
		if(p >= zf->in.iend)
			zf->in.iptr = p = fill_buf(zf);
		if(p >= zf->in.iend)
			break;
		zf->bits = zf->bits&~0xFF | *p++;
		zf->bits = zf->bits>>8 | zf->bits<<24;
		zf->nbit += 8;
		zf->in.iptr = p;
	}
	return zf->bits >> (-zf->nbit & 7);
}

#define eof(Z,N) ((Z)->nbit < (N))

LOCAL unsigned peek(struct zf *zf)
{
	return zf->bits >> (-zf->nbit & 7);
}

static inline unsigned alignb(struct zf *zf)
{
	return eatb(zf, zf->nbit & 7);
}

static unsigned getb(struct zf *zf, int n)
{
	unsigned m = peek(zf);
	m &= (1<<n)-1;
	eatb(zf, n);
	return m;
}

static unsigned eatbb(struct zf *zf, int n)
{
	while(n >= 4) {
		unsigned m;

		n -= 4;
		zf->nbit = 0;
		zf->bits = 0;
		zf->in.iptr += n;
		n = zf->in.iptr - zf->in.iend;
		m = eatb(zf, 0);
		if(n < 0)
			return m;
	}
	return eatb(zf, 8*n);
}

/* pump */

static int pump_gzip(struct zf *zf, int n);
static int pump_block(struct zf *zf, int n);
static int pump_uncompr(struct zf *zf, int n);
static int pump_compressed(struct zf *zf, int n);
static int pump_block_end(struct zf *zf, int n);

static int pump_gzip(struct zf *zf, int n)
{
	unsigned f;

	f = eatb(zf, 0); /* dont peek(), buffer may be empty */
	if((f & 0xE0FFFFFFu) != 0x088B1F)
//		return error(zf, ZF_ERR_DATA);
		return ret_error(zf, ZF_ERR_DATA);

	eatbb(zf, 10);
	f >>= 24;
	if(f & 1<<2) {
		int n = getb(zf, 16);
		eatbb(zf, n);
	}
	if(f & 1<<3)
		while(getb(zf, 8));
	if(f & 1<<4)
		while(getb(zf, 8));
	if(f & 1<<1)
		eatb(zf, 16);

	if(eof(zf, 1))
		return error(zf, ZF_ERR_TRUNC);

	zf->pump = pump_block;
	return 0;
}

static int block0(struct zf *zf);
static int block1(struct zf *zf);
static int block2(struct zf *zf);

static int pump_block(struct zf *zf, int n)
{
	unsigned m = getb(zf, 3);

	zf->final_block = m&1;
	switch(m>>1) {
		case 0: return block0(zf);
		case 1: return block1(zf);
		case 2: return block2(zf);
		default: return error(zf, ZF_ERR_DATA);
	}
}

static int pump_block_end(struct zf *zf, int max)
{
	u32 m;

	if(zf->lit) zf->lit = htree_free(zf->lit);
	if(zf->off) zf->off = htree_free(zf->off);
	if(!zf->final_block) {
		zf->pump = pump_block;
		return 0;
	}

	/* trailer */

	m = alignb(zf); // crc
	m = eatbb(zf, 4); // isize
	if(eof(zf, 32))
		error(zf, ZF_ERR_TRUNC);
	if(m != (u32)zf->end)
		error(zf, ZF_ERR_DATA);
	eatbb(zf, 4);
	zf->pump = eof(zf, 1) ? 0 : pump_gzip;
	return 0;
}

static int pump_uncompr(struct zf *zf, int n)
{
	int l = n<zf->blen ? n : zf->blen;
	u8 m;

	zf->blen -= l;
	m = peek(zf);
	do {
		if(eof(zf, 1))
			return error(zf, ZF_ERR_TRUNC);
		zf->hist[zf->end++ & 0x7FFF] = m;
		m = eatb(zf, 8);
	} while(--l);
	if(!zf->blen)
		zf->pump = pump_block_end;
	return 0;
}

static int pump_compressed(struct zf *zf, int num)
{
	u8 b = 0;

	assert(num > 0);

	for(;;) {
		struct htree *t;
		unsigned m, l;

		m = eatb(zf, b), b = 0;

		HTREE_LOOKUP(zf->lit,t,m,b)

		if(eof(zf, b))
			return error(zf, ZF_ERR_TRUNC);

		l = t->u.value>>8;
		if(!l) {
#if 0
			zf->hist[zf->end++ & 0x7FFF] = t->u.value;
#else
			long a = zf->end;
			u8 v = t->u.value;
			a &= 0x7FFF;
			zf->end++;
			zf->hist[a] = v;
#endif
			if(!--num)
				break;
		} else {
			unsigned d;

			if(l == 1)
				goto end;

			if((u8)t->u.value) {
				m >>= t->bits;
				b += (u8)t->u.value;
				l += m & (1<<(u8)t->u.value)-1;
			}

			m = eatb(zf, b), b = 0;

			HTREE_LOOKUP(zf->off,t,m,b)

			d = t->u.value >> 8;
			if((u8)t->u.value) {
				m >>= t->bits;
				b += (u8)t->u.value;
				d += m & (1<<(u8)t->u.value)-1;
			}

			if(eof(zf, b))
				return error(zf, ZF_ERR_TRUNC);

			num -= l;
#if 0
			do {
				zf->hist[zf->end & 0x7FFF]
					= zf->hist[zf->end-d & 0x7FFF];
				zf->end++;
			} while(--l);
#else
			{
				long a = zf->end - d;
				long b = zf->end;
				zf->end += l;
				do {
					a &= 0x7FFF;
					b &= 0x7FFF;
					zf->hist[b++] = zf->hist[a++];
				} while(--l);
			}
#endif
			if(num <= 0)
				break;
		}
	}

	eatb(zf, b);
	return 0;

end:
	zf->pump = pump_block_end;
	eatb(zf, b);
	return 0;
}

/* codes */

static unsigned lit_codes[286], off_codes[30+2];

static unsigned *filltab(unsigned *p, unsigned v, int rr, int n)
{
	unsigned m = 1<<8;
	int r = 2*rr;
	do {
		do {
			*p++ = v;
			v += m;
		} while(--r);
		m <<= 1; v++;
		r = rr;
	} while(--n);
	return p;
}

__attribute__((constructor))
static void init()
{
	unsigned *p = lit_codes;
	int i;
	for(i=0;i<=256;i++) *p++ = i;
	p = filltab(p, 0x300, 4, 6);
	*p = 258<<8; // why not 259?
	assert(p-lit_codes == 285);

	p = filltab(off_codes, 0x100, 2, 14);
	assert(p-off_codes == 30);
	p[0] = p[1] = 1<<8; // block type 1 forbidden codes
}

static int block0(struct zf *zf)
{
	u32 l = alignb(zf);

	if((l>>16 ^ l&0xFFFF) != 0xFFFF)
		return error(zf, ZF_ERR_DATA);
	eatbb(zf, 4);
	l &= 0xFFFF;

	zf->blen = l;
	zf->pump = l ? pump_uncompr : pump_block_end;
	return 0;
}

static int block1(struct zf *zf)
{
	unsigned l[288];
	int i;

	assert(!zf->lit);
	assert(!zf->off);

	for(i=0;i<=143;i++) l[i] = 8;
	for(;i<=255;i++) l[i] = 9;
	for(;i<=279;i++) l[i] = 7;
	for(;i<=287;i++) l[i] = 8;
	zf->lit = htree_gen(l, lit_codes, 288);

	for(i=0;i<32;i++) l[i] = 5;
	zf->off = htree_gen(l, off_codes, 32);

	if(!zf->lit || !zf->off)
		return error(zf, ZF_ERR_TREE);
	zf->pump = pump_compressed;
	return 0;
}

static struct htree *readht(struct zf *zf, struct htree *ct, unsigned val[], int nl)
{
	unsigned lt[288];
	unsigned vt[288];
	int i, n;
	u8 b, pvl;

	pvl = 0;
	i = n = 0;
	b = 0;
	do {
		struct htree *t;
		u32 m;

		m = eatb(zf, b); b = 0;

		HTREE_LOOKUP(ct,t,m,b)
		m >>= t->bits;

		switch(t->u.value) {
		case 0:
			pvl = t->u.value;
			n++;
			break;
		default:
			lt[i] = pvl = t->u.value;
			vt[i++] = val[n++];
			break;
		case 16:
			m = (m&3)+3; b += 2;
			if(n+m > nl) {
				WARN;
				return 0;
			}
			if(!pvl) {
				n += m;
				break;
			}
			do {
				lt[i] = pvl;
				vt[i++] = val[n++];
			} while(--m);
			break;
		case 17:
			b += 3;
			n += (m&7)+3;
			pvl = 0;
			break;
		case 18:
			b += 7;
			n += (m&127)+11;
			pvl = 0;
			break;
		}
	} while(n < nl);

	eatb(zf, b);

	return htree_gen(lt, vt, i);
}

static unsigned tvl[] = {16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15};

static int block2(struct zf *zf)
{
	u32 hlens = peek(zf);
	struct htree *t;

	assert(!zf->lit);
	assert(!zf->off);

	if((hlens&0x1F)+257 > 286 || (hlens>>5&0x1F)+1 > 30)
		return error(zf, ZF_ERR_DATA);

	{
		unsigned lt[19], vt[19];
		u32 m;
		int i, n;

		m = eatb(zf, 14);
		i = 0;
		n = (hlens>>10 & 0xF) + 4;
		m = m&0xFFFFFF;
		do {
			lt[tvl[i++]] = m&7; m >>= 3;
			if(!(i & 7)) {
				m = eatb(zf, 24);
				m = m&0xFFFFFF;
			}
		} while(i<n);
		if(i&7) eatb(zf, 3*(i&7));
		while(i<19)
			lt[tvl[i++]] = 0;

		for(i=n=0;i<19;i++) if(lt[i]) {
			lt[n] = lt[i];
			vt[n++] = i;
		}

		t = htree_gen(lt, vt, n);
		if(!t)
			return error(zf, ZF_ERR_TREE);
	}

	zf->lit = readht(zf, t, lit_codes, (hlens&0x1F)+257);
	zf->off = readht(zf, t, off_codes, (hlens>>5&0x1F)+1);
	htree_free(t);

	if(!zf->lit || !zf->off)
		return error(zf, ZF_ERR_TREE);

	zf->pump = pump_compressed;
	return 0;
}

/* stdio */

#ifdef __GLIBC__
static __ssize_t zf_file_read(void *cookie, char *ptr, size_t len)
#else
static int zf_file_read(void *cookie, char *ptr, int len)
#endif
{
	struct zf *zf = (struct zf*)cookie;
	u8 *p, *s;

	if(len <= 0)
		return 0;

	p = s = (u8*)ptr;
	for(;;) {
		while(zf->end == zf->rdp) {
			if(!zf->pump)
				goto ret;
			zf->pump(zf, len<16384 ? len : 16384);
		}
		while(zf->end != zf->rdp) {
			*p++ = zf->hist[zf->rdp++ & 0x7FFF];
			if(!--len)
				goto ret;
		}
	}
ret:
	if(p-s || !zf->errn)
		return p - s;
	else {
		errno = zf->errn;
		return -1;
	}
}

static int zf_file_close(void *cookie)
{
	zf_destroy((struct zf*)cookie);
	return 0;
}

static inline FILE *fopen_zf(struct zf *zf)
{
	FILE *f;
#if #system(bsd)
	f = funopen(zf,
		zf_file_read, 0,
		0, zf_file_close);
#elif defined __GLIBC__
	f = fopencookie(zf, "r", (cookie_io_functions_t){
		read: zf_file_read,
		close: zf_file_close,
	});
#else
#error funopen() equivalent not known
#endif
	return f;
}

static FILE *plain_file(struct zf *zf, FILE *f, int n)
{
	while(--n >= 0) {
		int v = ungetc(zf->in.ibuf[n], f);
		if(v == EOF)
			err(1, "ungetc");
	}
	zf_destroy(zf);
	return f;
}

#ifdef WANT_FOPENZ
static inline int read_all(int fd, u8 *p, int n)
{
	int l = n;
	do {
		int v = read(fd, p, l);
		if(v<=0) break;
		p += v; l -= v;
	} while(l);
	return n - l;
}

FILE *fopenz(const char *name, const char *mode)
{
	struct zf *zf;
	FILE *f;
	int fd = 0;
	int v;

	if(name) {
		fd = open(name, O_RDONLY);
		if(fd<0)
			return 0;
	}

	zf = (struct zf*)calloc(1, sizeof *zf);
	in_init(&zf->in);
	zf->in.fd = fd;

	v = read_all(fd, zf->in.ibuf, 4);
	if(v != 4) {
		if(v < 0)
			v = 0;
		goto ret_orig2;
	}
	zf->in.iend += v;

	v = pump_gzip(zf, 0);
	if(v) {
ret_orig:
		v = 4;
ret_orig2:
		f = name ? fdopen(fd, mode) : stdin;
		zf->in.fd = -1; // prevent close()
		return plain_file(zf, f, v);
	}

	f = fopen_zf(zf);
	if(!f)
		goto ret_orig;
	return f;
}
#endif

#ifdef WANT_FGUNZIP
FILE *fgunzip(FILE *origf)
{
	struct zf *zf;
	FILE *f;
	int v;

	zf = (struct zf*)calloc(1, sizeof *zf);
	in_init(&zf->in);
	zf->in.file = origf;

	v = fread(zf->in.ibuf, 1, 4, origf);
	if(v != 4)
		goto ret_orig2;
	zf->in.iend += v;

	v = pump_gzip(zf, 0);
	if(v) {
ret_orig:
		v = 4;
ret_orig2:
		zf->in.file = 0; // prevent fclose()
		return plain_file(zf, origf, v);
	}

	f = fopen_zf(zf);
	if(!f)
		goto ret_orig;
	return f;
}
#endif
#endif

