/*
 * Huffman Compression Decoder using static precalculated tree
 * Aaron Logue 1999
 * 
 * --------- file format --------
 * The decoding tree
 *    First byte is number of nodes
 *    Array of nodes; one word per node
 *    Bitstream of flags; two bits per node
 *       Each bit indicates whether a byte in a node is (0) a pointer
 *       to another node or (1) a leaf indicating the output 8-bit value.
 * 3 bit count of the number of bytes encoded mod 8
 * The bitstream
 * ------------------------------
 * 
 * Each node in the decoding tree is a word, with a left byte and a right
 * byte.  A byte is either a pointer to the next node or a value to emit
 * as decompression output.  There are two flag bits per node to tell us
 * which is which.
 *
 * To decode, we process sequential bits from the bitstream which tell us
 * to follow a left branch or right branch in the tree.  We look at a
 * flag associated with the node in the tree where that took us, and that
 * tells us whether we're at a leaf and can emit a byte or whether we have
 * to process another bit.
 */
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <string.h>
#include "huffman.h"

#define BUFSIZE 1024

HUFFMAN_DECODING_HANDLE * huffman_init_decoder(void) {
   HUFFMAN_DECODING_HANDLE * d;
   d = malloc(sizeof(HUFFMAN_DECODING_HANDLE));
   if (d) {
      d->flags = 0;
      d->numnodes = 0;
      d->tree = NULL;
      d->tflags = NULL;
      d->totalbytes = 0;
      d->totalmod8 = 0;
      d->head = 0;
      d->tail = 0;
      d->bitstream = malloc(BUFSIZE);
	  if (d->bitstream == NULL) {
         free(d);
         d = NULL;
      }
	  d->full = 0;
	  d->outbuf = malloc(BUFSIZE);
	  if (d->outbuf == NULL) {
         free(d->bitstream);
         free(d);
         d = NULL;
      }
   }
   return(d);
}

void huffman_deinit_decoder(HUFFMAN_DECODING_HANDLE * d) {
   if (d->tree) {
      free(d->tree);
   }
   if (d->tflags) {
      free(d->tflags);
   }
   if (d->bitstream) {
      free(d->bitstream);
   }
   if (d->outbuf) {
      free(d->outbuf);
   }
   if (d) {
      free(d);
   }
   return;
}

int read_decoder_tree(FILE * fp, HUFFMAN_DECODING_HANDLE * d) {
   int bytes;
   int expected;
   unsigned char ibuf[2];

   expected = 1;
   bytes = fread(&d->numnodes, 1, expected, fp);
   if (bytes != expected) {
      return 0;
   }

   expected = d->numnodes * sizeof(HUFFMAN_NODE);
   d->tree = malloc(expected);
   if (d->tree == NULL) {
      return 0;
   }
   bytes = fread(d->tree, 1, expected, fp);
   if (bytes != expected) {
      free(d->tree);
	  d->tree = NULL;
      return 0;
   }

   expected = ((d->numnodes-1)>>2)+1;
   d->tflags = malloc(expected);
   if (d->tflags == NULL) {
      free(d->tree);
	  d->tree = NULL;
      return;
   }
   bytes = fread(d->tflags, 1, expected, fp);
   if (bytes != expected) {
	  free(d->tflags);
	  d->tflags = NULL;
      free(d->tree);
	  d->tree = NULL;
      return;
   }
   return;
}

FILE * huffman_load_header(HUFFMAN_DECODING_HANDLE * d, char * infile) {
   FILE * fp;
   int bytes;

   fp = fopen(infile, "rb");
   if (fp == 0) {
      printf("Unable to open %s.\n", infile);
      return 0;
   }
   bytes = read_decoder_tree(fp, d);
   if (!bytes) {
      fclose(fp);
	  fp = NULL;
   }
   return fp;
}

void emit_byte(FILE * ofp, HUFFMAN_DECODING_HANDLE * d, unsigned char value) {
   d->outbuf[d->full] = value;
   d->totalbytes++;
   d->full++;
   if (d->full == BUFSIZE) {
      fwrite(d->outbuf, 1, BUFSIZE, ofp);
      d->full = 0;
   }
   return;
}

void huffman_decode_bitstream(FILE * fp, HUFFMAN_DECODING_HANDLE * d, char * ofile) {
   FILE * ofp;
   int t;
   int lastbuf;

   ofp = fopen(ofile, "wb");
   if (ofp == 0) {
      printf("Unable to open %s.\n", ofile);
      return;
   }
   /* to make a modular decoder, make this d->t and add a finish function */
   t = 0;
   d->tail = fread(d->bitstream, 1, BUFSIZE, fp) << 3;
   if (IS_HBIT(d->bitstream, 0)) {
      d->totalmod8 |= 1;
   }
   if (IS_HBIT(d->bitstream, 1)) {
      d->totalmod8 |= 2;
   }
   if (IS_HBIT(d->bitstream, 2)) {
      d->totalmod8 |= 4;
   }
   d->head = 3;
   lastbuf = feof(fp);

   /* loop until there are no more bits in the input stream */
   while (d->tail) {
	  /* loop until we've processed all the bits in this input buffer */
      while (d->head < d->tail) {
         /* the input bit takes us to the right or the left */
         if (IS_HBIT(d->bitstream, d->head)) {
            /* follow the right branch */
            if (IS_HBIT(d->tflags, (t<<1)+1)) {
               /* we've descended to a leaf - output a byte */
               emit_byte(ofp, d, d->tree[t].right);
               /* pop back up to root node for next byte */
               t = 0;
            } else {
               /* descend to the next node */
               t = d->tree[t].right;
            }
         } else {
            /* follow the left branch */
            if (IS_HBIT(d->tflags, t<<1)) {
               /* we've descended to a leaf - output a byte */
               emit_byte(ofp, d, d->tree[t].left);
               /* pop back up to root node for next byte */
               t = 0;
            } else {
               /* descend to the next node */
               t = d->tree[t].left;
            }
         }
         d->head++;
		 if (lastbuf && (d->head+8 >= d->tail) && 
               ((d->totalbytes & 0x07) == d->totalmod8)) {
            break;
         }
	  }
      /* get another input buffer */
      d->head = 0;
      d->tail = fread(d->bitstream, 1, BUFSIZE, fp) << 3;
      lastbuf = feof(fp);
   }
   if (d->full) {
      fwrite(d->outbuf, 1, d->full, ofp);
   }
   fclose(ofp);
   return;
}

int main(int argc, char **argv) {
   char * p;
   char ifile[256];
   char ofile[256];
   HUFFMAN_DECODING_HANDLE * d;
   FILE * fp;

   printf("Minimum redundancy variable length character decoding.\n\n");
   ifile[0] = 0;
   ofile[0] = 0;
   if (argc > 2) {
      while (argc > 1) {
         argc--;
         if (*argv[argc] == '-' || *argv[argc] == '/') {
            p = argv[argc];
            while (*p) {
               switch (*p++) {
                  case '-':
                  case '/':
                  default:
                     break;
               }
            }
         } else {
            if (ofile[0] == 0) {
               strcpy(ofile, argv[argc]);
            } else {
               strcpy(ifile, argv[argc]);
            }
         }
      }

      d = huffman_init_decoder();
      fp = huffman_load_header(d, ifile);
      if (fp) {
         huffman_decode_bitstream(fp, d, ofile);
         fclose(fp);
	  } else {
         printf("Error: bad header in file\n");
      }
      huffman_deinit_decoder(d);

   } else {
      printf("Usage: huff_dec <source> <destination>\n\n");
      exit(-1);
   }
}