#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/types.h>
#include "zrio.h"

#define Z_BUFSIZE 16384

Access * load_access(FILE *in);

zr_stream * zropen(const char *filename){
	zr_stream *zr;
	int fid;
	char * accname;
	int ret;

	fid = open(filename, O_RDONLY);
	if(fid == -1){ return NULL; }

	zr = (zr_stream*)malloc(sizeof(zr_stream));
	zr->stream  = (z_stream*)malloc(sizeof(z_stream));
	zr->fid    = fid;
	zr->acc_file= NULL;
	zr->access  = NULL;
	zr->in      = 0;
	zr->last_in = 0;
	zr->pos     = 0;
	zr->z_eof   = 0;
	zr->z_err   = 0;
	zr->inbuf   = (void*)malloc(WINSIZE);
	zr->stream->next_in = Z_NULL;
	zr->stream->avail_in = 0;
	zr->outbuf   = (void*)malloc(WINSIZE);
	zr->stream->next_out = Z_NULL;
	zr->stream->avail_out = 0;
    zr->stream->zalloc = Z_NULL;
    zr->stream->zfree = Z_NULL;
    zr->stream->opaque = Z_NULL;

    ret = inflateInit2(zr->stream, 47);      /* automatic zlib or gzip decoding */
	if(ret != Z_OK){
		zrclose(zr);
		return NULL;
	}

	accname = (char*)malloc(strlen(filename) + 4);
	strcpy(accname, filename);
	strcat(accname, ".zr");
	
	zr->acc_file = fopen(accname, "r");

	free(accname);

	return zr;
}

int64_t zrread(void *buf, int64_t len, zr_stream *zr){
	int ret;
	if(zr->z_err || zr->z_eof) return 0;
	zr->stream->next_out = buf;
	zr->stream->avail_out = len;
	zr->out += len;
	while(zr->stream->avail_out){
		if(!zr->stream->avail_in){
			zr->stream->avail_in = read(zr->fid, zr->inbuf, WINSIZE);
			zr->in += zr->stream->avail_in;
			zr->stream->next_in = zr->inbuf;
		}
		if(zr->stream->avail_in == 0){
			inflateEnd(zr->stream);
			zr->z_eof = 1;
			goto END;
		}
		ret = inflate(zr->stream, Z_NO_FLUSH);
		if (ret == Z_NEED_DICT){
			break;
		}
		if (ret == Z_MEM_ERROR || ret == Z_DATA_ERROR){
			fprintf(stderr, "** Error in inflating data at %s:%d **\n", __FILE__, __LINE__);
			zr->z_err = 1;
			goto END;
		}
		if (ret == Z_STREAM_END) {
			zr->z_eof = 1;
			break;
		}
	}
	END:
	zr->out -= zr->stream->avail_out;
	zr->pos += len - zr->stream->avail_out;
	return len - zr->stream->avail_out;
}

int zreof(zr_stream *zr){
	return zr->z_eof;
}

int zrerror(zr_stream *zr){
	return zr->z_err;
}

int64_t zrtell(zr_stream *zr){
	return zr->pos;
}

int zrskip(zr_stream *zr, int bytes){
	int m, n;
	unsigned char input[CHUNK];
	if(bytes < 0) return 0;
	n = 0;
	while(bytes > n){
		if(bytes - n >= CHUNK){
			m = CHUNK;
		} else {
			m = bytes - n;
		}
		m = zrread(input, m, zr);
		n += m;
		if(zr->z_eof || zr->z_err) return n;
	}
	return n;
}

int64_t set_z_pos(zr_stream *zr, Point *here, int64_t offset){
	int ret;
	int64_t off;
	z_stream *strm;
    ret = lseek(zr->fid, here->in - (here->bits ? 1 : 0), SEEK_SET);
	if(ret == -1) return -1;
	strm = zr->stream;
	if(!zr->z_eof) inflateEnd(zr->stream);
    strm->zalloc = Z_NULL;
    strm->zfree = Z_NULL;
    strm->opaque = Z_NULL;
    strm->avail_in = 0;
    strm->next_in = Z_NULL;
    ret = inflateInit2(strm, -15);
    if (ret != Z_OK)
        return ret;
	if(here->bits){
		ret = 0;
		if(read(zr->fid, &ret, 1) == 0){
			return -1;
        }
        (void)inflatePrime(strm, here->bits, ret >> (8 - here->bits));
	}
    (void)inflateSetDictionary(strm, here->window, WINSIZE);
	off = offset - here->out;
	strm->avail_in = 0;
	#ifdef ZRSKIP
	zr->pos = here->out;
	zr->z_eof = 0;
	zr->z_err = 0;
	zrskip(zr, (int)off);
	#else
	do {
        if (off > WINSIZE) {             /* skip WINSIZE bytes */
            strm->avail_out = WINSIZE;
            strm->next_out = zr->outbuf;
            off -= WINSIZE;
        }
        else if (off != 0) {             /* last skip */
            strm->avail_out = (unsigned)off;
            strm->next_out = zr->outbuf;
            off = 0;
        }

        /* uncompress until avail_out filled, or end of stream */
        do {
            if (strm->avail_in == 0) {
                strm->avail_in = read(zr->fid, zr->inbuf, WINSIZE);
                if (strm->avail_in == 0) {
					return -1;
                }
                strm->next_in = zr->inbuf;
            }
            ret = inflate(strm, Z_NO_FLUSH);       /* normal inflate */
            if (ret == Z_NEED_DICT)
                ret = Z_DATA_ERROR;
            if (ret == Z_MEM_ERROR || ret == Z_DATA_ERROR){
				printf("Inflate Error: %s:%d\n", __FILE__, __LINE__);
				return -1;
			}
            if (ret == Z_STREAM_END){
                break;
			}
        } while (strm->avail_out != 0);

        /* if reach end of stream, then don't keep trying to get more */
        if (ret == Z_STREAM_END){
			return (off == 0)? offset:-1; 
		}

	} while(off);
	zr->pos = offset;
	zr->z_eof = 0;
	zr->z_err = 0;
	#endif
	return offset;
}

int64_t zrseek(zr_stream *zr, int64_t offset){
	int i;
	Point *here;

	if(offset < 0) return -1;
	if(offset >= zr->pos && offset <= zr->pos + WINSIZE){
		zrskip(zr, (int)(offset - zr->pos));
		return offset;
	}

	if(zr->access == NULL){
		if(zr->acc_file != NULL){
			if((zr->access = load_access(zr->acc_file)) == NULL) return -1;
		} else return -1;
	}

	for(i=0;i<zr->access->have;i+=128){
		if(zr->access->list[i].out > offset){
			break;
		}
	}

	if(i >= zr->access->have) i = zr->access->have - 1;
	for(;i>0;i--){
		if(zr->access->list[i].out <= offset) break;
	}

	here = zr->access->list + i;
	
	if(offset > zr->pos && here->out < zr->pos + WINSIZE / 64){
		zrskip(zr, (int)(offset - zr->pos));
		return offset;
	}

	if(here->window == NULL){
		if(fseek(zr->acc_file, sizeof(int) + (sizeof(int) + 2 * sizeof(int64_t)) * zr->access->have + WINSIZE * i, SEEK_SET)) return -1;
		here->window = (unsigned char*)malloc(WINSIZE);
		if((i = fread(here->window, 1, WINSIZE, zr->acc_file)) != WINSIZE){
			return -1;
		}
	}

	return set_z_pos(zr, here, offset);
}

int64_t zrgets(void *buf, int64_t len, int64_t offset, zr_stream *zr){
	int i;
	int64_t l;
	Point *here;

	if(offset < 0) return -1;
	if(offset == zr->pos){
		return zrread(buf, len, zr);
	}
	if(offset > zr->pos && offset <= zr->pos + WINSIZE){
		zrskip(zr, (int)(offset - zr->pos));
		return zrread(buf, len, zr);
	}
	
	if(zr->access == NULL){
		if(zr->acc_file != NULL){
			if((zr->access = load_access(zr->acc_file)) == NULL) return -1;
		} else return -1;
	}
	
	for(i=0;i<zr->access->have;i+=128){
		if(zr->access->list[i].out > offset){
			break;
		}
	}

	if(i >= zr->access->have) i = zr->access->have - 1;
	for(;i>0;i--){
		if(zr->access->list[i].out <= offset) break;
		if(zr->access->list[i].out - WINSIZE <= offset) goto SPEC_WAY;
	}

	here = zr->access->list + i;

	if(offset > zr->pos && here->out < zr->pos + WINSIZE/64){
		zrskip(zr, (int)(offset - zr->pos));
		return zrread(buf, len, zr);
	}

	if(here->window == NULL){
		if(fseek(zr->acc_file, sizeof(int) + (sizeof(int) + 2 * sizeof(int64_t)) * zr->access->have + WINSIZE * i, SEEK_SET)) return -1;
		here->window = (unsigned char*)malloc(WINSIZE);
		if(fread(here->window, 1, WINSIZE, zr->acc_file) != WINSIZE) return -1;
	}

	if(set_z_pos(zr, here, offset) < 0) return -1;
	return zrread(buf, len, zr);

	SPEC_WAY:

	here = zr->access->list + i;
	if(here->window == NULL){
		if(fseek(zr->acc_file, sizeof(int) + (sizeof(int) + 2 * sizeof(int64_t)) * zr->access->have + WINSIZE * i, SEEK_SET)) return -1;
		here->window = (unsigned char*)malloc(WINSIZE);
		if(fread(here->window, 1, WINSIZE, zr->acc_file) != WINSIZE) return -1;
	}
	l = offset - zr->access->list[i].out;
	if(l > len) l = len;
	memcpy(buf, here->window + WINSIZE - l, l);
	if(l == len) return len;
	len -= l;
	if(set_z_pos(zr, here, here->out) < 0) return -1;
	return zrread(buf + l, len, zr) + l;
}

void zrclose(zr_stream *zr){
	if(zr->fid != -1) close(zr->fid);
	if(zr->acc_file) fclose(zr->acc_file);
	if(zr->access){
		free(zr->access->list);
		free(zr->access);
	}
	if(zr->inbuf) free(zr->inbuf);
	if(zr->outbuf) free(zr->outbuf);
	(void)inflateEnd(zr->stream);
	free(zr->stream);
	free(zr);
}


/* Save the access in a file
 */
#define my_write(file, buf, len) fwrite(buf, 1, len, file)
int save_access(FILE *out, Access *index){
	int i;
	Point *p;
	my_write(out, &(index->have), sizeof(int));
	for(i=0;i<index->have;i++){
		p = index->list + i;
		my_write(out, &(p->out), sizeof(int64_t));
		my_write(out, &(p->in), sizeof(int64_t));
		my_write(out, &(p->bits), sizeof(int));
	}
	for(i=0;i<index->have;i++){
		p = index->list + i;
		my_write(out, p->window, WINSIZE);
	}
	return i;
}

/* Load an access from file
 */
#define my_read(file, buf, len) fread(buf, 1, len, file)
Access * load_access(FILE *in){
	int i, size;
	Access *index;
	Point *p;
	my_read(in, &size, sizeof(int));
	if(size < 0){ return NULL; }
	index = (Access*)malloc(sizeof(Access));
	index->have = index->size = size;
	index->list = (Point*)malloc(sizeof(Point) * size);
	for(i=0;i<size;i++){
		p = index->list + i;
		p->window = NULL;
		my_read(in, &(p->out), sizeof(int64_t));
		my_read(in, &(p->in), sizeof(int64_t));
		my_read(in, &(p->bits), sizeof(int));
	}
	/*
	for(i=0;i<size;i++){
		p = index->list + i;
		p->window = (unsigned char *)malloc(WINSIZE);
		my_read(in, p->window, WINSIZE);
	}
	*/
	return index;
}

/* Deallocate an index built by build_index() */
void free_index(struct access *index)
{
    if (index != NULL) {
        free(index->list);
        free(index);
    }
}

/* Add an entry to the access point list.  If out of memory, deallocate the
   existing list and return NULL. */
struct access *addpoint(struct access *index, int bits,
    int64_t in, int64_t out, unsigned left, unsigned char *window)
{
    struct point *next;

    /* if list is empty, create it (start with eight points) */
    if (index == NULL) {
        index = malloc(sizeof(struct access));
        if (index == NULL) return NULL;
        index->list = malloc(sizeof(struct point) << 3);
        if (index->list == NULL) {
            free(index);
            return NULL;
        }
        index->size = 8;
        index->have = 0;
    }

    /* if list is full, make it bigger */
    else if (index->have == index->size) {
        index->size <<= 1;
        next = realloc(index->list, sizeof(struct point) * index->size);
        if (next == NULL) {
            free_index(index);
            return NULL;
        }
        index->list = next;
    }

    /* fill in entry and increment how many we have */
    next = index->list + index->have;
    next->bits = bits;
    next->in = in;
    next->out = out;
	next->window = (unsigned char*)malloc(WINSIZE);
    if (left)
        memcpy(next->window, window + WINSIZE - left, left);
    if (left < WINSIZE)
        memcpy(next->window + left, window, WINSIZE - left);
    index->have++;

    /* return list, possibly reallocated */
    return index;
}

unsigned short get_max_dist(struct inflate_state *state){
	int i, dmax;
	dmax = 0;
	for(i=state->dist_window_tail;i<state->dmax;i++){
		if(dmax + (i - state->dist_window_tail) < state->dist_window[i]) 
			dmax = state->dist_window[i] - (i - state->dist_window_tail);
	}
	for(i=0;i<state->dist_window_tail;i++){
		if(dmax + (i + state->dmax - state->dist_window_tail) < state->dist_window[i])
			dmax = state->dist_window[i] - (i + state->dmax - state->dist_window_tail);
	}
	return (unsigned short)dmax;
}

/* Make one entire pass through the compressed stream and build an index, with
   access points about every span bytes of uncompressed output -- span is
   chosen to balance the speed of random access against the memory requirements
   of the list, about 32K bytes per access point.  Note that data after the end
   of the first zlib or gzip stream in the file is ignored.  build_index()
   returns the number of access points on success (>= 1), Z_MEM_ERROR for out
   of memory, Z_DATA_ERROR for an error in the input file, or Z_ERRNO for a
   file read error.  On success, *built points to the resulting index. */
int build_index(int in, int64_t span, struct access **built, void (*notify)(void *obj, unsigned char *buf, int buf_size, int64_t pos), void *obj)
{
    int ret;
    int64_t totin, totout, tmp;        /* our own total counters to avoid 4GB limit */
    int64_t last;                 /* totout value of last access point */
    struct access *index;       /* access points being generated */
    z_stream strm;
	unsigned char input[CHUNK];
	unsigned char discard[WINSIZE];
	unsigned char *window = discard;
	struct inflate_state *state;

	memset(window, 0, WINSIZE);
    /* initialize inflate */
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    strm.avail_in = 0;
    strm.next_in = Z_NULL;
    ret = inflateInit2(&strm, 47);      /* automatic zlib or gzip decoding */
    if (ret != Z_OK)
        return ret;

    /* inflate the input, maintain a sliding window, and build an index -- this
       also validates the integrity of the compressed data using the check
       information at the end of the gzip or zlib stream */
    totin = totout = last = 0;
    index = NULL;               /* will be allocated by first addpoint() */
    strm.avail_out = 0;
	state = (struct inflate_state *)strm.state;
    do {
        /* get some compressed data from input file */
        strm.avail_in = read(in, input, CHUNK);
        if (strm.avail_in == 0) {
            ret = Z_DATA_ERROR;
            goto build_index_error;
        }
        strm.next_in = input;

        /* process all of that, or until end of stream */
        do {
            /* reset sliding window if necessary */
            if (strm.avail_out == 0) {
                strm.avail_out = WINSIZE;
                strm.next_out = window;
//				printf("MARK: dist = %d / %d %s:%d\n", get_max_dist(state), state->dmax, __FILE__, __LINE__);
            }

            /* inflate until out of input, output, or at end of block --
               update the total input and output counters */
            totin += strm.avail_in;
            totout += strm.avail_out;
			tmp = strm.avail_out;
            ret = inflate_zr(&strm, Z_BLOCK);      /* return at end of block */
            totin -= strm.avail_in;
            totout -= strm.avail_out;
			if(notify) notify(obj, window + WINSIZE - tmp, tmp - strm.avail_out, totout);
            if (ret == Z_NEED_DICT)
                ret = Z_DATA_ERROR;
            if (ret == Z_MEM_ERROR || ret == Z_DATA_ERROR)
                goto build_index_error;
            if (ret == Z_STREAM_END)
                break;

            /* if at end of block, consider adding an index entry (note that if
               data_type indicates an end-of-block, then all of the
               uncompressed data from that block has been delivered, and none
               of the compressed data after that block has been consumed,
               except for up to seven bits) -- the totout == 0 provides an
               entry point after the zlib or gzip header, and assures that the
               index always has at least one access point; we avoid creating an
               access point after the last block by checking bit 6 of data_type
             */
            if ((strm.data_type & 128) && !(strm.data_type & 64) &&
                (totout == 0 || totout - last > span)) {
                index = addpoint(index, strm.data_type & 7, totin,
                                 totout, strm.avail_out, window);
                if (index == NULL) {
                    ret = Z_MEM_ERROR;
                    goto build_index_error;
                }
                last = totout;
            }
        } while (strm.avail_in != 0);
    } while (ret != Z_STREAM_END);

    /* clean up and return index (release unused entries in list) */
    (void)inflateEnd(&strm);
    index = realloc(index, sizeof(struct point) * index->have);
    index->size = index->have;
    *built = index;
    return index->size;

    /* return error */
  build_index_error:
    (void)inflateEnd(&strm);
    if (index != NULL)
        free_index(index);
    return 0;
}

int zrmkindex(const char *filename, int level,
			void (*notify)(void *obj, unsigned char *buf, int buf_size, int64_t pos), void *obj){
	int fid;
	FILE *file;
	Access *index;
	int64_t span;
	char *accname;
	if(level <= 1) return 0;
	span = level * WINSIZE;
	if((fid = open(filename, O_RDONLY)) == -1) return 0;
	if(!build_index(fid, span, &index, notify, obj)){
		close(fid);
		return 0;
	}
	close(fid);
	accname = (char*)malloc(strlen(filename) + 4);
	strcpy(accname, filename);
	strcat(accname, ".zr");
	file = fopen(accname, "w");
	free(accname);
	if(file == NULL){
		free_index(index);
		return 0;
	}
	save_access(file, index);
	free_index(index);
	fclose(file);
	return 1;
}

