/* encode.c - Don Yang (uguu.org)

   Huffman encode file to base-64 representation.

   11/21/01
                                                                           */
/*@ -compdef -compmempass -mustfree -nullpass -nullstate -shiftsigned @*/
/*@ -temptrans -type -usedef -usereleased @*/

#include<stdio.h>
#include<stdlib.h>
#include<string.h>

#define MAX_LINE_SIZE      256
#define MAX_CODE_LENGTH    1024
#define BASE64_OFFSET      '0'

typedef struct _node
{
   int c, len;
   char *code;
   struct _node *left, *right;

} Node;


static int AddNode(Node **bst, int c, int len, char *code);
static void DelTree(Node *bst);
static void FlushCode(int byte, FILE *outfile);
static /*@null@*/char *GetCode(Node *bst, int c, int len);
static int LoadDictionary(Node **bst, FILE *dict, FILE *hint);
static void WriteCode(Node *bst, int c, int len,
                      int *byte, int *bit, FILE *outfile);

/******************************************************************** main */
int main(int argc, char **argv)
{
   FILE *infile, *outfile, *dict, *hint;
   int i, c, len, bit, byte;
   Node *dictionary;

   if( argc < 3 )
      return printf("%s <dictionary> <hint> [infile] [outfile]\n", *argv);

   infile = stdin;
   outfile = stdout;
   dict = fopen(argv[1], "rt");
   hint = fopen(argv[2], "rt");
   if( argc > 3 )
   {
      infile = fopen(argv[3], "rt");
      if( argc > 4 )
         outfile = fopen(argv[4], "wt+");
   }
   if( infile == NULL || outfile == NULL || dict == NULL || hint == NULL )
   {
      (void)puts("Error opening files");
      if( infile != stdin && infile != NULL )      (void)fclose(infile);
      if( outfile != stdout && outfile != NULL )   (void)fclose(outfile);
      if( dict != NULL )                           (void)fclose(dict);
      if( hint != NULL )                           (void)fclose(hint);
      return 1;
   }


   if( LoadDictionary(&dictionary, dict, hint) != 0 )
   {
      (void)puts("Error loading dictionary");
      (void)fclose(infile);
      (void)fclose(outfile);
      (void)fclose(dict);
      (void)fclose(hint);
      DelTree(dictionary);
      return 0;
   }

   c = EOF;
   len = bit = byte = 0;
   (void)fputc('\"', outfile);
   do
   {
      i = fgetc(infile);
      if( i != c )
      {
         if( c != EOF )
            WriteCode(dictionary, c, len, &byte, &bit, outfile);
         c = i;
         len = 1;
      }
      else
      {
         len++;
      }
   } while( i != EOF );

   if( bit > 0 )
      FlushCode(byte, outfile);
   (void)fputs("\"\n", outfile);


   (void)fclose(infile);
   (void)fclose(outfile);
   (void)fclose(dict);
   (void)fclose(hint);
   DelTree(dictionary);
   return 0;

} /* main() */

/***************************************************************** AddNode */
static int AddNode(Node **bst, int c, int len, char *code)
{
   Node **cursor;
   Node *node;

   for(cursor = bst;
      *cursor != NULL && ((*cursor)->c != c || (*cursor)->len != len);)
   {
      if( (*cursor)->c == c )
      {
         if( (*cursor)->len == len )
            break;
         cursor = (*cursor)->len < len ? &((*cursor)->left)
                                       : &((*cursor)->right);
      }
      else
      {
         cursor = (*cursor)->c < c ? &((*cursor)->left)
                                   : &((*cursor)->right);
      }
   }

   if( *cursor != NULL )
      return 1;

   if( (node = (Node*)malloc(sizeof(Node))) == NULL )
      return 1;
   if( (node->code = (char*)malloc(strlen(code) + 1)) == NULL )
   {
      free(node);
      return 1;
   }

   strcpy(node->code, code);
   node->c = c;
   node->len = len;
   node->left = node->right = NULL;
   *cursor = node;
   return 0;

} /* AddNode() */

/***************************************************************** DelTree */
static void DelTree(Node *bst)
{
   Node *cursor, *next;

   for(cursor = bst; cursor != NULL; cursor = next)
   {
      DelTree(cursor->left);
      next = cursor->right;
      free(cursor->code);
      free(cursor);
   }

} /* DelTree() */

/*************************************************************** FlushCode */
static void FlushCode(int byte, FILE *outfile)
{
   if( (byte += BASE64_OFFSET) == '\\' )
      (void)fputc('\\', outfile);
   (void)fputc(byte, outfile);

} /* FlushCode() */

/***************************************************************** GetCode */
static char *GetCode(Node *bst, int c, int len)
{
   Node *cursor;

   for(cursor = bst; cursor != NULL;)
   {
      if( cursor->c == c )
      {
         if( cursor->len == len )
            break;
         cursor = cursor->len < len ? cursor->left : cursor->right;
      }
      else
      {
         cursor = cursor->c < c ? cursor->left : cursor->right;
      }
   }

   if( cursor != NULL )
      return cursor->code;
   return NULL;

} /* GetCode() */

/********************************************************** LoadDictionary */
static int LoadDictionary(Node **bst, FILE *dict, FILE *hint)
{
   static char buffer[MAX_LINE_SIZE + 1];
   static char code[MAX_CODE_LENGTH + 1];
   char *p, *q;
   int i, c, len;

   *bst = NULL;
   while( fgets(buffer, MAX_LINE_SIZE, dict) != NULL )
   {
      i = 0;
      for(p = buffer; *p == '0' || *p == '1'; code[i++] = *(p++));
      code[i] = '\0';
      if( sscanf(p, " (x %d) -> char %d, len %d", &i, &c, &len) != 3 )
         return 1;
      if( AddNode(bst, c, len, code) != 0 )
         return 1;
   }

   while( fgets(buffer, MAX_LINE_SIZE, hint) != NULL )
   {
      if( sscanf(buffer, "char %d, len %d", &c, &len) != 2 )
         return 1;
      code[0] = '\0';
      for(p = strchr(buffer, '=') + 1; p != NULL; p = strchr(p + 1, ' '))
      {
         i = atoi(p);
         if( (q = GetCode(*bst, c, i)) != NULL )
            strcat(code, q);
      }
      if( AddNode(bst, c, len, code) != 0 )
         return 1;
   }

   return 0;

} /* LoadDictionary() */

/*************************************************************** WriteCode */
static void WriteCode(Node *bst, int c, int len,
                      int *byte, int *bit, FILE *outfile)
{
   char *p;

   if( (p = GetCode(bst, c, len)) == NULL )
      return;
   for(; *p != '\0'; p++)
   {
      *byte |= (*p == '0' ? 0 : 1) << (*bit)++;
      if( *bit > 5 )
      {
         FlushCode(*byte, outfile);
         *byte = *bit = 0;
      }
   }

} /* WriteCode() */