/*
 *     Copyright (c) 2017, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 *
 */

/*
 * This tool uses the zlib general purpose compression library
 * written by Jean-loup Gailly and Mark Adler.  This tool is based
 * on the excellent and well-documented zpipe.c example Version 1.2
 * (November 2004) written by Mark Adler.  See www.zlib.net for
 * more information on zlib.
 * zlib is Copyright (C) 1995-2005 Jean-loup Gailly and Mark Adler
 */
/*
 * pgzip.c - compress, using zlib, a file
 */

#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "zlib.h"

static int debug = 0;

#define BUFSIZE 1024

/*
 * compress infile to outfile.
 * return 0 on success, or a zlib error if the library returns an error.
 */
static int
zipcompress(FILE *infile, FILE *outfile, int level)
{
  int zret, zflush;
  unsigned int bytesout, b;
  z_stream ztream;
  unsigned char inbuf[BUFSIZE];
  unsigned char outbuf[BUFSIZE];

  /* initialize */
  ztream.zalloc = Z_NULL;
  ztream.zfree = Z_NULL;
  ztream.opaque = Z_NULL;
  zret = deflateInit(&ztream, level);
  if (zret != Z_OK)
    return zret;

  /* read a chunk at a time */
  do {
    ztream.avail_in = fread(inbuf, 1, BUFSIZE, infile);
    if (ferror(infile)) {
      deflateEnd(&ztream);
      return Z_ERRNO;
    }
    /* if this is the end of file, flush all output */
    zflush = feof(infile) ? Z_FINISH : Z_NO_FLUSH;
    ztream.next_in = inbuf;

    /* compress until the last output buffer chunk */
    do {
      ztream.avail_out = BUFSIZE;
      ztream.next_out = outbuf;
      zret = deflate(&ztream, zflush);
      if (zret == Z_STREAM_ERROR) {
        deflateEnd(&ztream);
        return zret;
      }
      bytesout = BUFSIZE - ztream.avail_out;
      if (bytesout) {
        b = fwrite(outbuf, sizeof(char), bytesout, outfile);
        if (b != bytesout || ferror(outfile)) {
          deflateEnd(&ztream);
          return Z_ERRNO;
        }
      }
      /* continue if this output buffer was full */
    } while (ztream.avail_out == 0);
    /* continue until EOF on infile */
  } while (zflush != Z_FINISH);

  deflateEnd(&ztream);
  return 0;
} /* zipcompress */

/*
 * report error from zlib
 */
static void
puterr(int err, FILE *infile, FILE *outfile, char *infilename,
       char *outfilename)
{
  switch (err) {
  case Z_ERRNO:
    if (ferror(infile)) {
      fprintf(stderr, "pgzip: error reading input file: %s\n", infilename);
      err = 0;
    }
    if (ferror(outfile)) {
      fprintf(stderr, "pgzip: error reading output file: %s\n", outfilename);
      err = 0;
    }
    if (err) {
      fprintf(stderr, "pgzip: unknown file error\n");
    }
    break;
  case Z_STREAM_ERROR:
    fprintf(stderr, "pgzip: invalid compression level\n");
    break;
  case Z_DATA_ERROR:
    fprintf(stderr, "pgzip: invalid or incomplete data\n");
    break;
  case Z_MEM_ERROR:
    fprintf(stderr, "pgzip: out of memory\n");
    break;
  case Z_VERSION_ERROR:
    fprintf(stderr, "pgzip: zlib version mismatch\n");
    break;
  }
} /* puterr */

/*
 * compress a file
 */
int
main(int argc, char *argv[])
{
  int a, dofile = 0, level = 9;
  char *infilename = NULL, *outfilename = NULL;
  FILE *infile = NULL, *outfile = NULL;
  int err = 0;
  for (a = 1; a < argc; ++a) {
    if (!dofile && argv[a][0] == '-') {
      if (argv[a][1] >= '0' && argv[a][1] <= '9') {
        level = argv[a][1] - '0';
      } else if (strcmp(argv[a], "--") == 0) {
        dofile = 1;
      } else if (strcmp(argv[a], "-d") == 0) {
        debug = 1;
      } else {
        if (!err)
          fprintf(stderr, "pgzip infile output\n");
        fprintf(stderr, "   unexpected argument: %s\n", argv[a]);
        err = 1;
      }
    } else {
      if (infilename == NULL) {
        infilename = argv[a];
        if (!dofile && strcmp(infilename, "-") == 0) {
          infile = stdin;
        } else {
          infile = fopen(infilename, "rb");
          if (infile == NULL) {
            if (!err)
              fprintf(stderr, "pgzip infile output\n");
            fprintf(stderr, "  could not open input file: %s\n", argv[a]);
            err = 1;
          }
        }
      } else if (outfilename == NULL) {
        outfilename = argv[a];
        if (!dofile && strcmp(outfilename, "-") == 0) {
          outfile = stdout;
        } else {
          outfile = fopen(outfilename, "wb");
          if (outfile == NULL) {
            if (!err)
              fprintf(stderr, "pgzip infile output\n");
            fprintf(stderr, "  could not open output file: %s\n", argv[a]);
            err = 1;
          }
        }
      } else {
        if (!err)
          fprintf(stderr, "pgzip infile output\n");
        fprintf(stderr, "   unexpected argument: %s\n", argv[a]);
        err = 1;
      }
    }
  }
  if (err)
    return err;

  if (infile == NULL)
    infile = stdin;
  if (outfile == NULL)
    outfile = stdout;

  err = zipcompress(infile, outfile, level);
  if (err)
    puterr(err, infile, outfile, infilename, outfilename);
  return err;
} /* main */
