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

// Recover Agrippa ciphertext from Agrippa program data fork
// Works both on a pristine copy, and on one that has been "destroyed" by running once

// Output of this can be sent into agrippa-decrypt for plaintext of the poem

typedef unsigned char byte;
typedef unsigned short word;

const unsigned DNA_Offset = 0xA60E8;  // Location where Agrippa attempts (unsuccessfully) to destroy itself

// In the data fork: LZW compressed block containing the ciphertext
const unsigned Agrippa_Offset = 446932;
const unsigned Agrippa_Length = 102793;

const unsigned mem_start = 0x147000; // where it gets uncompressed into, in 68k memory

const unsigned cipher_start = 0x16b8e0; // where the cipher lives, in 68k memory
const unsigned cipher_length = 9918;


byte compressed[Agrippa_Length];
int nextbyte = 0;
int bits = 0;

byte uncompressed[400000];  // actual length 175104
byte *out = uncompressed;


// A variant of LZW (MSB)

const word RESET = 0x100;
const word FIN = 0xFFFF;
const unsigned max_codeword = 0x1000;

byte wordchar[max_codeword];
word wordprefix[max_codeword];
unsigned words = 256;


unsigned get(int bitlen) {
  unsigned ret = 0;

  while (--bitlen >= 0) {
    if (nextbyte >= Agrippa_Length) return 0xffff;

    ret <<= 1;
    ret |= (compressed[nextbyte] >> (7-bits))&1;

    bits = (bits + 1) % 8;
    if (bits == 0)
      nextbyte++;
  }
  return ret;
}


byte output( word w) {
  if (w < 256) {
    *out++ = w;
    return w;
  }

  byte ret = output(wordprefix[w]);
  output(wordchar[w]);
  return ret;    
}

int main(int argc, char*argv[]) {
  if (argc < 2) exit(printf("need the filename of the Agrippa data fork\n"));

  FILE* f = fopen(argv[1],"rb");
  if (!f) exit(printf("No such file\n"));

  fseek(f,DNA_Offset,SEEK_SET);
  byte x = fgetc(f);
  printf("This copy of agrippa is: %s\n",x == 0x54 ? "Destroyed" : "Pristine");

  fseek(f,Agrippa_Offset,SEEK_SET);
  fread(compressed,1,sizeof(compressed),f);

  for (int i = 0; i <= 255; i++) { wordchar[i] = i; wordprefix[i] = 0xffff;}
  unsigned prevword = 0;
  unsigned bitsize = 9;

  for ( ;; ) {
    unsigned codeword = get(bitsize);

    if (codeword == FIN) 
      break;

    if (codeword == RESET) {
      words=256; 
      bitsize=9; 
      prevword = 0;

      if (bits) get(8-bits); // flush extra bits
      switch (nextbyte % 4) {  // alignment or checksum or both at reset, ignore it
      case 0: get(32); break;
      case 1: get(24); break;
      case 2: get(48); break;
      case 3: get(0); break; // not encountered, unknown skip
      }

      continue;
    }

    byte first;
    if (codeword == words) {
      first = output(prevword);
      output(first);
    } else {
      first = output(codeword);
    }

    if (words < max_codeword) {
      wordchar[words] = first;
      wordprefix[words] = prevword;
      words++;
      if ( (1<<bitsize) == words && words < max_codeword)
	bitsize++;
    }

    prevword = codeword;
  }

  printf("Recovered ciphertext:\n");
  for (int i = 0 ; i < cipher_length; i++) {
    printf("%02x%c",uncompressed[cipher_start-mem_start+i],(i%16)==15 ? '\n' : ' ');
  }
  printf("\n");
  return 0;
}