/*
 * Copyright (c) 2005-2009 Nominet UK (www.nic.uk)
 * All rights reserved.
 * Contributors: Ben Laurie, Rachel Willmer, Alasdair Mackintosh.
 * The Contributors have asserted their moral rights under the
 * UK Copyright Design and Patents Act 1988 to
 * be recorded as the authors of this copyright work.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. 
 * 
 * You may obtain a copy of the License at 
 *     http://www.apache.org/licenses/LICENSE-2.0 
 * 
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * 
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/** \file
 */

#include <zlib.h>
#include <bzlib.h>
#include <string.h>

#include <openpgpsdk/compress.h>
#include <openpgpsdk/packet-parse.h>
#include <openpgpsdk/crypto.h>
#include <openpgpsdk/errors.h>
#include "parse_local.h"
#include <openpgpsdk/final.h>
#include <openpgpsdk/partial.h>

static const int debug = 0;

#define DECOMPRESS_BUFFER	1024
#define COMPRESS_BUFFER	        32768

typedef struct
    {
    ops_compression_type_t type;
    ops_region_t *region;
    unsigned char in[DECOMPRESS_BUFFER];
    unsigned char out[DECOMPRESS_BUFFER];
    z_stream zstream; // ZIP and ZLIB
    size_t offset;
    int inflate_ret;
    } z_decompress_arg_t;

typedef struct
    {
    ops_compression_type_t type;
    ops_region_t *region;
    char in[DECOMPRESS_BUFFER];
    char out[DECOMPRESS_BUFFER];
    bz_stream bzstream; // BZIP2
    size_t offset;
    int inflate_ret;
    } bz_decompress_arg_t;

typedef struct
    {
    z_stream stream;
    unsigned char *src;
    unsigned char *dst;
    size_t bytes_in;
    size_t bytes_out;
    } compress_arg_t;

// \todo remove code duplication between this and bzip2_compressed_data_reader
static int zlib_compressed_data_reader(void *dest,size_t length,
				       ops_error_t **errors,
				       ops_reader_info_t *rinfo,
				       ops_parse_cb_info_t *cbinfo)
{
	z_decompress_arg_t *arg=ops_reader_get_arg(rinfo);

	if(! (arg->type==OPS_C_ZIP || arg->type==OPS_C_ZLIB) ) // ASSERT(arg->type==OPS_C_ZIP || arg->type==OPS_C_ZLIB);
	{
		fprintf(stderr,"zlib_compressed_data_reader: unexpected compress format %d\n",arg->type) ;
		return -1 ;
	}

	//ops_parser_content_t content;
	int saved=length;

	if(/*arg->region->indeterminate && */ arg->inflate_ret == Z_STREAM_END
			&& arg->zstream.next_out == &arg->out[arg->offset])
		return 0;

	if(arg->region->length_read == arg->region->length)
	{
		if(arg->inflate_ret != Z_STREAM_END)
			OPS_ERROR(cbinfo->errors, OPS_E_P_DECOMPRESSION_ERROR,
					"Compressed data didn't end when region ended.");
		/*
			else
			return 0;
			www.zlib.org
		 */
	}

	while(length > 0)
	{
		unsigned len;

		if(&arg->out[arg->offset] == arg->zstream.next_out)
		{
			int ret;

			arg->zstream.next_out=arg->out;
			arg->zstream.avail_out=sizeof arg->out;
			arg->offset=0;
			if(arg->zstream.avail_in == 0)
			{
				unsigned n=arg->region->length;

				if(!arg->region->indeterminate)
				{
					n-=arg->region->length_read;
					if(n > sizeof arg->in)
						n=sizeof arg->in;
				}
				else
					n=sizeof arg->in;

				if(!ops_stacked_limited_read(arg->in,n,arg->region,
							errors,rinfo,cbinfo))
					return -1;

				arg->zstream.next_in=arg->in;
				arg->zstream.avail_in=arg->region->indeterminate
					? arg->region->last_read : n;
			}

			ret=inflate(&arg->zstream,Z_SYNC_FLUSH);
			if(ret == Z_STREAM_END)
			{
				if(!arg->region->indeterminate
						&& arg->region->length_read != arg->region->length)
					OPS_ERROR(cbinfo->errors,OPS_E_P_DECOMPRESSION_ERROR,
							"Compressed stream ended before packet end.");
			}
			else if(ret != Z_OK)
			{
				fprintf(stderr,"ret=%d\n",ret);
				OPS_ERROR(cbinfo->errors,OPS_E_P_DECOMPRESSION_ERROR, arg->zstream.msg);
			}
			arg->inflate_ret=ret;
		}

		if(!(arg->zstream.next_out > &arg->out[arg->offset])) // ASSERT(arg->zstream.next_out > &arg->out[arg->offset]);
		{
			fprintf(stderr,"decompression error: next packet is too large.\n");
			return -1 ;
		}
		len=arg->zstream.next_out-&arg->out[arg->offset];
		if(len > length)
			len=length;
		memcpy(dest,&arg->out[arg->offset],len);
		arg->offset+=len;
		length-=len;
	}

	return saved;
}

// \todo remove code duplication between this and zlib_compressed_data_reader
static int bzip2_compressed_data_reader(void *dest,size_t length,
					ops_error_t **errors,
					ops_reader_info_t *rinfo,
					ops_parse_cb_info_t *cbinfo)
{
	bz_decompress_arg_t *arg=ops_reader_get_arg(rinfo);

	if(!(arg->type==OPS_C_BZIP2)) // ASSERT(arg->type==OPS_C_BZIP2);
	{
		fprintf(stderr,"bzip2_compressed_data_reader: unexpected compress format %d\n",arg->type) ;
		return -1 ;
	}

	//ops_parser_content_t content;
	int saved=length;

	if(arg->inflate_ret == BZ_STREAM_END
			&& arg->bzstream.next_out == &arg->out[arg->offset])
		return 0;

	if(arg->region->length_read == arg->region->length)
	{
		if(arg->inflate_ret != BZ_STREAM_END)
			OPS_ERROR(cbinfo->errors, OPS_E_P_DECOMPRESSION_ERROR,
					"Compressed data didn't end when region ended.");
	}

	while(length > 0)
	{
		unsigned len;

		if(&arg->out[arg->offset] == arg->bzstream.next_out)
		{
			int ret;

			arg->bzstream.next_out=(char *) arg->out;
			arg->bzstream.avail_out=sizeof arg->out;
			arg->offset=0;
			if(arg->bzstream.avail_in == 0)
			{
				unsigned n=arg->region->length;

				if(!arg->region->indeterminate)
				{
					n-=arg->region->length_read;
					if(n > sizeof arg->in)
						n=sizeof arg->in;
				}
				else
					n=sizeof arg->in;

				if(!ops_stacked_limited_read(arg->in, n, arg->region, errors,
							rinfo, cbinfo))
					return -1;

				arg->bzstream.next_in=arg->in;
				arg->bzstream.avail_in=arg->region->indeterminate
					? arg->region->last_read : n;
			}

			ret=BZ2_bzDecompress(&arg->bzstream);
			if(ret == BZ_STREAM_END)
			{
				if(!arg->region->indeterminate
						&& arg->region->length_read != arg->region->length)
					OPS_ERROR(cbinfo->errors, OPS_E_P_DECOMPRESSION_ERROR,
							"Compressed stream ended before packet end.");
			}
			else if(ret != BZ_OK)
			{
				OPS_ERROR_1(cbinfo->errors, OPS_E_P_DECOMPRESSION_ERROR,
						"Invalid return %d from BZ2_bzDecompress", ret);
			}
			arg->inflate_ret=ret;
		}
		if(!(arg->bzstream.next_out > &arg->out[arg->offset])) // ASSERT(arg->bzstream.next_out > &arg->out[arg->offset]);
		{
			fprintf(stderr,"bzip2_compressed_data_reader: unexpected size in compressed packet.\n") ;
			return -1 ;
		}
		len=arg->bzstream.next_out-&arg->out[arg->offset];
		if(len > length)
			len=length;
		memcpy(dest,&arg->out[arg->offset],len);
		arg->offset+=len;
		length-=len;
	}

	return saved;
}

/**
 * \ingroup Core_Compress
 * 
 * \param *region 	Pointer to a region
 * \param *parse_info 	How to parse
 * \param type Which compression type to expect
*/

int ops_decompress(ops_region_t *region,ops_parse_info_t *parse_info,
		   ops_compression_type_t type)
    {
    z_decompress_arg_t z_arg;
    bz_decompress_arg_t bz_arg;
    int ret;

    switch (type)
        {
    case OPS_C_ZIP:
    case OPS_C_ZLIB:
        memset(&z_arg,'\0',sizeof z_arg);

        z_arg.region=region;
        z_arg.offset=0;
        z_arg.type=type;

        z_arg.zstream.next_in=Z_NULL;
        z_arg.zstream.avail_in=0;
        z_arg.zstream.next_out=z_arg.out;
        z_arg.zstream.zalloc=Z_NULL;
        z_arg.zstream.zfree=Z_NULL;
        z_arg.zstream.opaque=Z_NULL;
        break;

    case OPS_C_BZIP2:
        memset(&bz_arg,'\0',sizeof bz_arg);

        bz_arg.region=region;
        bz_arg.offset=0;
        bz_arg.type=type;

        bz_arg.bzstream.next_in=NULL;
        bz_arg.bzstream.avail_in=0;
        bz_arg.bzstream.next_out=bz_arg.out;
        bz_arg.bzstream.bzalloc=NULL;
        bz_arg.bzstream.bzfree=NULL;
        bz_arg.bzstream.opaque=NULL;
        break;

    default:
        OPS_ERROR_1(&parse_info->errors, OPS_E_ALG_UNSUPPORTED_COMPRESS_ALG,
		    "Compression algorithm %d is not yet supported", type);
        return 0;
        }

    switch(type)
        {
    case OPS_C_ZIP:
        ret=inflateInit2(&z_arg.zstream,-15);
        break;

    case OPS_C_ZLIB:
        ret=inflateInit(&z_arg.zstream);
        break;

    case OPS_C_BZIP2:
        /*
        OPS_ERROR_1(&parse_info->errors, OPS_E_ALG_UNSUPPORTED_COMPRESS_ALG, "Compression algorithm %s is not yet supported", "BZIP2");
        return 0;
        */
        ret=BZ2_bzDecompressInit(&bz_arg.bzstream, 1, 0);
        break;

    default:
        OPS_ERROR_1(&parse_info->errors, OPS_E_ALG_UNSUPPORTED_COMPRESS_ALG,
		    "Compression algorithm %d is not yet supported", type);
        return 0;
        }

    switch (type)
        {
    case OPS_C_ZIP:
    case OPS_C_ZLIB:
        if(ret != Z_OK)
            {
            OPS_ERROR_1(&parse_info->errors, OPS_E_P_DECOMPRESSION_ERROR,
			"Cannot initialise ZIP or ZLIB stream "
			"for decompression: error=%d", ret);
            return 0;
            }
        ops_reader_push(parse_info,zlib_compressed_data_reader,NULL,&z_arg);
        break;

    case OPS_C_BZIP2:
        if (ret != BZ_OK)
            {
            OPS_ERROR_1(&parse_info->errors, OPS_E_P_DECOMPRESSION_ERROR,
			"Cannot initialise BZIP2 stream "
			"for decompression: error=%d", ret);
            return 0;
            }
        ops_reader_push(parse_info,bzip2_compressed_data_reader,NULL,&bz_arg);
        break;

    default:
        OPS_ERROR_1(&parse_info->errors, OPS_E_ALG_UNSUPPORTED_COMPRESS_ALG,
		    "Compression algorithm %d is not yet supported", type);
        return 0;
        }

    ret=ops_parse(parse_info,ops_true);

    ops_reader_pop(parse_info);

    return ret;
    }

/**
\ingroup Core_WritePackets
\brief Writes Compressed packet
\param data Data to write out
\param len Length of data
\param cinfo Write settings
\return ops_true if OK; else ops_false
*/

ops_boolean_t ops_write_compressed(const unsigned char *data,
                                   const unsigned int len,
                                   ops_create_info_t *cinfo)
{
	int r=0;
	int sz_in=0;
	int sz_out=0;
	compress_arg_t* compress=ops_mallocz(sizeof *compress);

	// compress the data
	const int level=Z_DEFAULT_COMPRESSION; // \todo allow varying levels
	compress->stream.zalloc=Z_NULL;
	compress->stream.zfree=Z_NULL;
	compress->stream.opaque=NULL;

	// all other fields set to zero by use of ops_mallocz

	if (deflateInit(&compress->stream,level) != Z_OK)
	{
		// can't initialise

		fprintf(stderr,"ops_write_compressed: can't initialise compression engine\n") ;
		return ops_false ;
	}

	// do necessary transformation
	// copy input to maintain const'ness of src

	if(!(compress->src==NULL)) // ASSERT(compress->src==NULL);
	{
		fprintf(stderr,"ops_write_compressed: compress->src should be NULL\n") ;
		return ops_false ;
	}
	if(!(compress->dst==NULL)) // ASSERT(compress->dst==NULL);
	{
		fprintf(stderr,"ops_write_compressed: compress->dst should be NULL\n") ;
		return ops_false ;
	}

	sz_in=len * sizeof (unsigned char);
	sz_out= (sz_in * 1.01) + 12; // from zlib webpage
	compress->src=ops_mallocz(sz_in);
	compress->dst=ops_mallocz(sz_out);
	memcpy(compress->src,data,len);

	// setup stream
	compress->stream.next_in=compress->src;
	compress->stream.avail_in=sz_in;
	compress->stream.total_in=0;

	compress->stream.next_out=compress->dst;
	compress->stream.avail_out=sz_out;
	compress->stream.total_out=0;

	r=deflate(&compress->stream, Z_FINISH);

	if(!(r==Z_STREAM_END)) // ASSERT(r==Z_STREAM_END); // need to loop if not
	{
		fprintf(stderr,"ops_write_compressed: compression failed.\n") ;
		return ops_false ;
	}

	// write it out
	return (ops_write_ptag(OPS_PTAG_CT_COMPRESSED, cinfo)
			&& ops_write_length(1+compress->stream.total_out, cinfo)
			&& ops_write_scalar(OPS_C_ZLIB,1,cinfo)
			&& ops_write(compress->dst, compress->stream.total_out,cinfo));
}


// Writes out the header for the compressed packet. Invoked by the
// partial stream writer. Note that writing the packet tag and the
// packet length is handled by the partial stream writer.
static ops_boolean_t write_compressed_header(ops_create_info_t *info,
                                             void *header_data)
    {
    OPS_USED(header_data);
    // Write the compression type. Currently we just use ZLIB
    ops_write_scalar(OPS_C_ZLIB, 1, info);
    return ops_true;
    }

static void zlib_error(ops_error_t **errors, z_stream *stream, int error)
    {
    OPS_ERROR_2(errors,OPS_E_FAIL,
		"Error from compression stream %d (%s)", error,
		stream->msg == NULL ? "Unknown" :  stream->msg);
    }

static ops_boolean_t stream_compress_writer(const unsigned char *src,
                                            unsigned length,
                                            ops_error_t **errors,
                                            ops_writer_info_t *winfo)
    {
    // ZLib doesn't like being asked to compress nothing, so return if
    // we are given no input.
    if (length == 0)
	return ops_true;
    if (debug)
	fprintf(stderr, "Compressing %u bytes\n", length);
    compress_arg_t* compress = ops_writer_get_arg(winfo);
    compress->bytes_in += length;
    compress->stream.next_in = (void*) src;
    compress->stream.avail_in = length;
    ops_boolean_t result = ops_true;
    do
	{
	compress->stream.next_out = compress->dst;
	compress->stream.avail_out = COMPRESS_BUFFER;
	int retcode = deflate(&compress->stream, Z_NO_FLUSH);
	if (retcode != Z_OK)
	    {
	    zlib_error(errors, &compress->stream, retcode);
	    deflateEnd(&compress->stream);
	    return ops_false;
	    }
	unsigned bytes_to_write = COMPRESS_BUFFER - compress->stream.avail_out;
	if (debug)
	    fprintf(stderr, "bytes_to_write = %u\n", bytes_to_write);
	compress->bytes_out += bytes_to_write;
	result = ops_stacked_write(compress->dst, bytes_to_write, errors,
				   winfo);
	}
    while (result && compress->stream.avail_out == 0);

    return result;
    }

static ops_boolean_t stream_compress_finaliser(ops_error_t **errors,
                                               ops_writer_info_t *winfo)
    {
    compress_arg_t* compress = ops_writer_get_arg(winfo);
    compress->stream.next_in = NULL;
    compress->stream.avail_in = 0;
    int retcode = Z_OK;
    int output_size = COMPRESS_BUFFER;
    ops_boolean_t result = ops_true;
    do
	{
	compress->stream.next_out = compress->dst;
	compress->stream.avail_out = output_size;
	retcode = deflate(&compress->stream, Z_FINISH);
	if (retcode != Z_STREAM_END && retcode != Z_OK)
	    {
	    zlib_error(errors, &compress->stream, retcode);
	    deflateEnd(&compress->stream);
	    return ops_false;
	    }
	int bytes_to_write = output_size - compress->stream.avail_out;
	if (debug)
	    fprintf(stderr, "At end, bytes_to_write = %u\n", bytes_to_write);
	compress->bytes_out += bytes_to_write;
	result = ops_stacked_write(compress->dst, bytes_to_write, errors,
				   winfo);

	// If deflate returns Z_OK after we have asked to flush, it means
	// that there was not enough space in the output buffer. Increase
	// the buffer size and try again.
	if (retcode != Z_STREAM_END)
	    {
	    if (debug)
		fprintf(stderr, "Reallocating %u\n", output_size * 2);
	    output_size *= 2;
	    compress->dst = realloc(compress->dst, output_size);
	    }
	}
    while (result && retcode != Z_STREAM_END);
  
    int error = deflateEnd(&compress->stream);
    if (error != Z_OK)
	{
	zlib_error(errors, &compress->stream, error);
	return ops_false;
	}
    return result;
    }

static void stream_compress_destroyer(ops_writer_info_t *winfo)
    {
    compress_arg_t* compress = ops_writer_get_arg(winfo);
    if (debug)
	fprintf(stderr, "Compressed %zu to %zu\n", compress->bytes_in,
		compress->bytes_out);
    free(compress->dst);
    free(compress);
    }

/**
\ingroup Core_WritePackets
\brief Pushes a compressed writer onto the stack. Data written
       will be encoded as a compressed packet.
\param cinfo Write settings
*/
ops_boolean_t ops_writer_push_compressed(ops_create_info_t *cinfo)
{
	// This is a streaming writer, so we don't know the length in
	// advance. Use a partial writer to handle the partial body
	// packet lengths.
	ops_writer_push_partial(COMPRESS_BUFFER,
			cinfo, OPS_PTAG_CT_COMPRESSED,
			write_compressed_header, NULL);

	// Create arg to be used with this writer
	// Remember to free this in the destroyer
	compress_arg_t *compress = ops_mallocz(sizeof *compress);

	compress->dst = malloc(COMPRESS_BUFFER);
	const int level=Z_DEFAULT_COMPRESSION; // \todo allow varying levels
	compress->stream.zalloc=Z_NULL;
	compress->stream.zfree=Z_NULL;
	compress->stream.opaque=NULL;
	compress->stream.avail_out = COMPRESS_BUFFER;
	// all other fields set to zero by use of ops_mallocz

	if (deflateInit(&compress->stream, level) != Z_OK)
		// can't initialise. Is there a better way to handle this?
		return ops_false ;

	// And push writer on stack
	ops_writer_push(cinfo, stream_compress_writer, stream_compress_finaliser,
			stream_compress_destroyer, compress);

	return ops_true ;
}

// EOF
