#include "decode.h"

// Bit Packing Variables
static int code = 0;
static int usedBits = 0;
static int bitLen = 9;



static void cleanUp(DecodeEntry *table, Stack *stack)
{
  free(table);
  free(stack -> data);
  free(stack);
}



Stack *initStack(int stack_size) 
{
  Stack *stack = (Stack *)malloc(sizeof(Stack));
  if (stack == NULL)
  {
    fprintf(stderr, "Memory allocation for stack failed");
    return NULL;
  }  
  
  stack -> data = (unsigned char *)malloc(stack_size * sizeof(unsigned char));
  if (stack -> data == NULL)
  {
    fprintf(stderr, "Memory allocation for stack data failed");
    return NULL;
  }

  stack -> top = -1;
  stack -> max_size = stack_size;
  return stack;
}

void push(Stack *stack, unsigned char value) 
{
  if (stack -> top < stack -> max_size - 1) 
  {
    stack -> top++;
    stack -> data[stack -> top] = value;
  } else {
    fprintf(stderr, "Stack overflow: string too long (max_size=%d)\n", stack->max_size);
    exit(1);
  }
}

unsigned char pop(Stack *stack) 
{
  if (stack -> top >= 0) 
  {
    unsigned char value = stack -> data[stack -> top];
    stack -> top--;
    return value;
  }
  return 0;  
}

int isStackEmpty(Stack *stack) 
{
    return stack -> top == -1;
}



DecodeEntry *initDecodeTable(int table_size) 
{
  struct DecodeEntry* table = (struct DecodeEntry*)calloc(table_size, sizeof(struct DecodeEntry));
  if (table == NULL)
  {
    fprintf(stderr, "Failed to allocate table\n");
    return NULL;
  }
  for (int i = 0; i < 256; i++) 
  {
    table[i].prefix = EMPTY;
    table[i].character = (unsigned char)i;
  }
  return table;
}

int pruneDecodeTable(DecodeEntry **table, int table_size)
{
  bool *usedAsPrefix = calloc(table_size, sizeof(bool));
  if (usedAsPrefix == NULL)
  {
    fprintf(stderr, "Memory allocation for prefix tracker failed\n");
    return -1;
  }
  for (int code = 256; code < table_size; code++)
  {
    DecodeEntry item = (*table)[code];
    usedAsPrefix[item.prefix] = true;
  }

  int *codeMap = calloc(table_size, sizeof(int));
  if (codeMap == NULL)
  {
    fprintf(stderr, "Memory allocation for code mapping table failed\n");
    free(usedAsPrefix);
    return -1;
  }
  for (int i = 0; i < table_size; i++) 
  {
    codeMap[i] = EMPTY;
  }

  DecodeEntry *newTable = initDecodeTable(table_size);
  if (newTable == NULL)
  {
    fprintf(stderr, "Memory allocation for new array failed\n");
    free(usedAsPrefix);
    free(codeMap);
    return -1;
  }
  int nextCode = 256;

  for (int code = 256; code < table_size; code++)
  {
    if (usedAsPrefix[code])
    {
      DecodeEntry item = (*table)[code];

      int findPrefix = item.prefix;
      if (findPrefix > 255 && codeMap[findPrefix] != EMPTY)
      {
        findPrefix = codeMap[findPrefix];
      }

      codeMap[code] = nextCode;
      if (item.prefix > 255 && codeMap[item.prefix] != EMPTY) 
      {
        item.prefix = codeMap[item.prefix];
      }

      newTable[nextCode++] = item;
    }
  }

  free(usedAsPrefix);
  free(codeMap);
  free(*table);
  *table = newTable;
  return nextCode;
}



void decode(void) 
{
  int maxBits, pruning;
  char *stage = getenv("STAGE");
  if (!stage)
  {
    stage = "3";
  }

  if (!strcmp(stage, "3"))
  {
    unsigned char coderConfig;
    if (fread(&coderConfig, sizeof(unsigned char), 1, stdin) != 1)
    {
      fprintf(stderr, "Error: failed to read maxBits/pruning configuration");
      exit(1);
    }
    maxBits = coderConfig & 0x7F;
    pruning = coderConfig >> 7;
  } else if (!strcmp(stage, "2"))
  {
    if (fscanf(stdin, "%d", &maxBits) != 1 || 
        fscanf(stdin, "%d", &pruning) != 1)
    {
      return;
    }
  } else {
    maxBits = DEFAULT_BITS;
    pruning = 0;
  }

  int table_size = (1 << maxBits);  // 2^MAXBITS
  DecodeEntry *table = initDecodeTable(table_size);

  Stack *stack = initStack(table_size);
  
  int nextCode = 256;
  int newC = -1;
  int C, oldC;
  int finalK;
  
  if (strcmp(stage, "3"))
  {
    // Stage 1 and 2
    if (fscanf(stdin, "%d", &newC) != 1) 
    {
      cleanUp(table, stack);
      return;
    }

    oldC = newC;
    C = newC;
    finalK = table[C].character;
    fputc(finalK, stdout);

    while (fscanf(stdin, "%d", &newC) == 1) 
    {
      if (nextCode >= table_size && pruning)
      {
        nextCode = pruneDecodeTable(&table, table_size);
        if (nextCode == -1)
        {
          cleanUp(table, stack);
          exit(1);
        }
      }

      C = newC;
      
      if (C >= nextCode) 
      {
        C = oldC;
        push(stack, finalK);
      }
      
      while (C > 255) 
      {
        push(stack, table[C].character);
        C = table[C].prefix;
      }
      
      finalK = table[C].character;
      fputc(finalK, stdout);
      
      while (!isStackEmpty(stack)) 
      {
        fputc(pop(stack), stdout);
      }

      if (nextCode < table_size) 
      {
        table[nextCode].prefix = oldC;
        table[nextCode].character = finalK;

        nextCode++;
      }
      
      oldC = newC;
    }
  } else {
    // Stage 3
    unsigned char byte;
    newC = -1;

    while (newC == -1 && fread(&byte, sizeof(unsigned char), 1, stdin))
    {
      newC = bitUnpacker(byte);
    }

    oldC = newC;
    C = newC;
    finalK = table[C].character;
    fputc(finalK, stdout);

    while (fread(&byte, sizeof(unsigned char), 1, stdin))
    {
      newC = bitUnpacker(byte);
      if (newC == -1) continue;
      
      if (nextCode >= table_size && pruning)
      {
        nextCode = pruneDecodeTable(&table, table_size);
        if (nextCode == -1)
        {
          cleanUp(table, stack);
          exit(1);
        }
      }

      C = newC;
      
      if (C >= nextCode) 
      {
        C = oldC;
        push(stack, finalK);
      }
      
      while (C > 255) 
      {
        push(stack, table[C].character);
        C = table[C].prefix;
      }
      
      finalK = table[C].character;
      fputc(finalK, stdout);
      
      while (!isStackEmpty(stack)) 
      {
        fputc(pop(stack), stdout);
      }

      if (nextCode < table_size) 
      {
        table[nextCode].prefix = oldC;
        table[nextCode].character = finalK;
        nextCode++;

        if (nextCode + 1 >> bitLen)
        {
          if (usedBits > 0 && fread(&byte, sizeof(unsigned char), 1, stdin))
          {
            bitUnpacker(byte);
            code = 0;
            usedBits = 0;
          }

          bitLen++;
        }
      }
      
      oldC = newC;
    }
  }
  
  if (getenv("DBG") != NULL)
  {
    dumpDecodeTable(table, nextCode, table_size);
  }

  cleanUp(table, stack);
}

int bitUnpacker(unsigned char byte) 
{
  code = (code << BYTE_SIZE) | byte;
  usedBits += BYTE_SIZE;

  if (usedBits >= bitLen) 
  {
    int temp = code >> (usedBits - bitLen);

    code &= ((1 << (usedBits - bitLen)) - 1);
    usedBits -= bitLen;

    return temp;
  }

  return -1;
}



void dumpDecodeTable(struct DecodeEntry *table, int nextCode, int table_size) 
{
  char filename[256];
  snprintf(filename, sizeof(filename), "./DBG.%s", "decode");
  FILE* fp = fopen(filename, "w");
  if (!fp) 
  {
    fprintf(stderr, "Error opening file %s: %s (errno: %d)\n", 
            filename, strerror(errno), errno);
    return;
  }
  fprintf(fp, "CODE\tPREF\tCHAR\tString\n");

  for (int code = 0; code < nextCode; code++) 
  {
    char string[table_size];
    int len = 0;
    int curr = code;

    while (curr > 255 && len < table_size) 
    {
      string[len++] = table[curr].character;
      curr = table[curr].prefix;
    }
    if (len < table_size) {
      string[len++] = table[curr].character;
    }


    fprintf(fp, "%d\t%d\t%d\t", 
            code, table[code].prefix, table[code].character);
    for (int i = len - 1; i >= 0; i--) 
    {
      fputc(string[i], fp);
    }
    fputc('\n', fp);
  }
  fclose(fp);
}