// TO DO: add suitable code to save mem[0:65535] after init.
// Currently only compressing a whole file actually writes out the compressed array.

// from https://raw.githubusercontent.com/andyherbert/lz1/refs/heads/master/lz.c
// Use this to save then restore initial mem[] array.
// Bug fixed applied.

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <math.h>
#include <string.h>
#include <error.h>
#include <errno.h>


static inline int imax(int a, int b) {
  if (a > b) return a; else return b;
}

uint32_t crc32(uint8_t *data, uint32_t len) {
  uint32_t crc = (uint32_t)-1;
  for (int i = 0; i < len; i++) {
    crc = crc ^ data[i]; for (int j = 7; j >= 0; j--) crc = (crc >> 1) ^ (0xEDB88320 & (-(crc & 1)));
  }
  return ~crc;
}

uint32_t lz77_compress(uint8_t *uncompressed_text,
                       uint32_t uncompressed_size,
                       char *filename_out) {
  uint8_t pointer_length_width = 5;
  uint16_t pointer_pos, temp_pointer_pos, output_pointer, pointer_length, temp_pointer_length;
  uint32_t compressed_pointer, coding_pos, output_lookahead_ref, look_behind, look_ahead;
  uint16_t pointer_pos_max, pointer_length_max;
  uint8_t *compressed_text;
  uint32_t compressed_size;
  size_t malloc_size = 1024*16;
  FILE *out = NULL;

  compressed_text = calloc(1, imax(uncompressed_size, malloc_size)); // not very good that it requires an array as big as the output :-(

  pointer_pos_max = (1 << (16 - pointer_length_width));
  pointer_length_max = (1 << pointer_length_width);

  *((uint32_t *)compressed_text) = uncompressed_size;
  *(compressed_text + 4) = pointer_length_width;
  compressed_pointer = compressed_size = 5;

  for (coding_pos = 0; coding_pos < uncompressed_size; ++coding_pos) {
    pointer_pos = 0;
    pointer_length = 0;
    for (temp_pointer_pos = 1; (temp_pointer_pos < pointer_pos_max) && (temp_pointer_pos <= coding_pos);
         ++temp_pointer_pos) {
      look_behind = coding_pos - temp_pointer_pos;
      look_ahead = coding_pos;
      //for(temp_pointer_length = 0; uncompressed_text[look_ahead++] == uncompressed_text[look_behind++]; ++temp_pointer_length)
      for (temp_pointer_length = 0;
           look_ahead < uncompressed_size && uncompressed_text[look_ahead++] == uncompressed_text[look_behind++];
           ++temp_pointer_length) {
        if (temp_pointer_length == pointer_length_max) {
          break;
        }
      }
      if (temp_pointer_length > pointer_length) {
        pointer_pos = temp_pointer_pos;
        pointer_length = temp_pointer_length;
        if (pointer_length == pointer_length_max) break;
      }
    }
    coding_pos += pointer_length;
    if ((coding_pos == uncompressed_size) && pointer_length) {
      output_pointer = (pointer_length == 1) ? 0 : ((pointer_pos << pointer_length_width) | (pointer_length - 2));
      output_lookahead_ref = coding_pos - 1;
    } else {
      output_pointer = (pointer_pos << pointer_length_width) | (pointer_length ? (pointer_length - 1) : 0);
      output_lookahead_ref = coding_pos;
    }
    *((uint16_t *)(compressed_text + compressed_pointer)) = output_pointer;
    compressed_pointer += 2;
    *(compressed_text + compressed_pointer++) = *(uncompressed_text + output_lookahead_ref);
    compressed_size += 3;
  }

  out = fopen(filename_out, "w");
  if (out == NULL) {
    fprintf(stderr, "lz77_compress: cannot open \"%s\" - %s\n", filename_out, strerror(errno));
    return 0;
  }
  
  uint32_t crc;
  crc = crc32(uncompressed_text, uncompressed_size);

  fprintf(out, "#include <stdint.h>\n");
  fprintf(out, "#include <math.h>\n");
  fprintf(out, "#include <stdio.h>\n\n");
  fprintf(out, "#include <stdlib.h>\n\n");

  fprintf(out, "const uint8_t d%0x[%d] = {\n  ", crc, compressed_size);
  for (int i = 0; i < compressed_size; i++) {
    fprintf(out, "%3d, ", compressed_text[i]);
    if ((i&15)==15) {
      fprintf(out, "\n  ");
    }
  }
  fprintf(out, "\n};\n\n");

  fprintf(out, "uint32_t lz77_decompress(const uint8_t *compressed_text,\n");
  fprintf(out, "                         uint8_t *uncompressed_text) {\n");
  fprintf(out, "  uint8_t pointer_length_width;\n");
  fprintf(out, "  uint16_t input_pointer, pointer_length, pointer_pos, pointer_length_mask;\n");
  fprintf(out, "  uint32_t compressed_pointer, coding_pos, pointer_offset, uncompressed_size;\n");
  fprintf(out, "  uncompressed_size = *((uint32_t *)compressed_text);\n");
  fprintf(out, "  pointer_length_width = *(compressed_text + 4);\n");
  fprintf(out, "  compressed_pointer = 5;\n");
  fprintf(out, "  pointer_length_mask = (1 << pointer_length_width) - 1;\n");
  fprintf(out, "  for (coding_pos = 0; coding_pos < uncompressed_size; ++coding_pos) {\n");
  fprintf(out, "    input_pointer = *((uint16_t *)(compressed_text + compressed_pointer));\n");
  fprintf(out, "    compressed_pointer += 2;\n");
  fprintf(out, "    pointer_pos = input_pointer >> pointer_length_width;\n");
  fprintf(out, "    pointer_length = pointer_pos ? ((input_pointer & pointer_length_mask) + 1) : 0;\n");
  fprintf(out, "    if (pointer_pos) {\n");
  fprintf(out, "      for (pointer_offset = coding_pos - pointer_pos; pointer_length > 0; --pointer_length) {\n");
  fprintf(out, "        uncompressed_text[coding_pos++] = uncompressed_text[pointer_offset++];\n");
  fprintf(out, "      }\n");
  fprintf(out, "    }\n");
  fprintf(out, "    *(uncompressed_text + coding_pos) = *(compressed_text + compressed_pointer++);\n");
  fprintf(out, "  }\n");
  fprintf(out, "  return coding_pos;\n");
  fprintf(out, "}\n\n");

  fprintf(out, "uint32_t crc32(uint8_t *data, uint32_t len) {\n");
  fprintf(out, "  uint32_t crc = (uint32_t)-1;\n");
  fprintf(out, "  for (int i = 0; i < len; i++) {\n");
  fprintf(out, "    crc = crc ^ data[i]; for (int j = 7; j >= 0; j--) crc = (crc >> 1) ^ (0xEDB88320 & (-(crc & 1)));\n");
  fprintf(out, "  }\n");
  fprintf(out, "  return ~crc;\n");
  fprintf(out, "}\n\n");

  fprintf(out, "int restore_%0x(uint8_t *output) {\n", crc);
  fprintf(out, "  lz77_decompress(d%0x, output);\n", crc);
  fprintf(out, "  if (crc32(output,%0d) == 0x%0xU) {\n", uncompressed_size, crc);
  fprintf(out, "    // fprintf(stderr, \"Decompressed OK!\\n\");\n");
  fprintf(out, "  } else {\n");
  fprintf(out, "    fprintf(stderr, \"Decompression failed.\\n\"); exit(1);\n");
  fprintf(out, "  }\n");
  fprintf(out, "  return 0;\n");
  fprintf(out, "}\n\n");

  fprintf(out, "#ifdef DBMAIN\n");
  fprintf(out, "int main(int argc, char **argv) {\n");
  fprintf(out, "  uint8_t output[%0d];\n", uncompressed_size);
  fprintf(out, "  restore_%0x(output);\n", crc);
  fprintf(out, "  return 0;\n");
  fprintf(out, "}\n");
  fprintf(out, "#endif\n\n");
  (void) fclose(out);
  free(compressed_text);
  return compressed_size;
}

uint32_t lz77_decompress(const uint8_t *compressed_text,
                         uint8_t *uncompressed_text) {
  uint8_t pointer_length_width;
  uint16_t input_pointer, pointer_length, pointer_pos, pointer_length_mask;
  uint32_t compressed_pointer, coding_pos, pointer_offset, uncompressed_size;
  uncompressed_size = *((uint32_t *)compressed_text);
  pointer_length_width = *(compressed_text + 4);
  compressed_pointer = 5;
  pointer_length_mask = (1 << pointer_length_width) - 1;
  for (coding_pos = 0; coding_pos < uncompressed_size; ++coding_pos) {
    input_pointer = *((uint16_t *)(compressed_text + compressed_pointer));
    compressed_pointer += 2;
    pointer_pos = input_pointer >> pointer_length_width;
    pointer_length = pointer_pos ? ((input_pointer & pointer_length_mask) + 1) : 0;
    if (pointer_pos) {
      for (pointer_offset = coding_pos - pointer_pos; pointer_length > 0; --pointer_length) {
        uncompressed_text[coding_pos++] = uncompressed_text[pointer_offset++];
      }
    }
    *(uncompressed_text + coding_pos) = *(compressed_text + compressed_pointer++);
  }
  return coding_pos;
}

long fsize(FILE *in) {
  long pos, length;
  pos = ftell(in);
  fseek(in, 0L, SEEK_END);
  length = ftell(in);
  fseek(in, pos, SEEK_SET);
  return length;
}

uint32_t file_lz77_compress(char *filename_in, char *filename_out) {
  FILE *in;
  uint8_t *uncompressed_text;
  uint32_t uncompressed_size, compressed_size, uc;

  in = fopen(filename_in, "rb");
  if (in == NULL) {
    fprintf(stderr, "Cannot open \"%s\" - %s", filename_in, strerror(errno));
    return 0;
  }
  uncompressed_size = fsize(in);
  uncompressed_text = calloc(1, uncompressed_size);
  if ((uncompressed_size != (uc = fread(uncompressed_text, 1, uncompressed_size, in)))) {
    fprintf(stderr, "Read of uncompressed data returned %d bytes\n", uc);
    // should really be in a loop if uc > 0
    return 0;
  }
  fclose(in);

  compressed_size = lz77_compress(uncompressed_text, uncompressed_size, filename_out);

  free(uncompressed_text);

  return compressed_size;
}

uint32_t file_lz77_decompress(char *filename_in, char *filename_out) {
  FILE *in, *out;
  uint32_t compressed_size, uncompressed_size, uc;
  uint8_t *compressed_text, *uncompressed_text;

  in = fopen(filename_in, "rb");
  if (in == NULL) {
    fprintf(stderr, "file_lz77_decompress: cannot open \"%s\" - %s", filename_in, strerror(errno));
    return 0;
  }
  compressed_size = fsize(in);
  compressed_text = calloc(1, compressed_size);
  if ((uc = fread(compressed_text, 1, compressed_size, in)) != compressed_size) {
    fprintf(stderr, "file_lz77_decompress: read of compressed data returned %d bytes\n", uc);
    return 0;
  }
  fclose(in);

  uncompressed_size = *((uint32_t *)compressed_text);
  uncompressed_text = calloc(1, uncompressed_size);

  if ((uc = lz77_decompress(compressed_text, uncompressed_text)) != uncompressed_size) {
    fprintf(stderr, "file_lz77_decompress: decompressed size was %d bytes\n", uc);
    return 0;
  }

  out = fopen(filename_out, "wb");
  if (out == NULL) {
    fprintf(stderr, "file_lz77_decompress: cannot open \"%s\" - %s", filename_out, strerror(errno));
    return 0;
  }
  if (fwrite(uncompressed_text, 1, uncompressed_size, out) != uncompressed_size) {
    fprintf(stderr, "file_lz77_decompress: write of uncompressed data only wrote %d bytes\n", uc);
    return 0;
  }
  fclose(out);
  free(compressed_text);
  free(uncompressed_text);
  return uncompressed_size;
}

void savemem(uint8_t *mem) {
  uint32_t size = lz77_compress(mem, 65536, "mem.c");
  if (size == 0) exit(1);
}

#ifdef STANDALONE
int main(int argc, char **argv) {
  if (argc != 2) {
    fprintf(stderr, "syntax: lz file\n");
    return 0;
  }
  FILE *in = fopen(argv[1], "rb");
  if (in == NULL) {
    fprintf(stderr, "lz: cannot open \"%s\" - %s", argv[1], strerror(errno));
    return 0;
  }
  fclose(in);
  file_lz77_compress(argv[1], "/dev/stdout");
  // possible tweak to consider adding is RLE for large blocks of 0's.
  // since a 64K file of 0's takes up around 6K, and the point of this is
  // to save a 64K memory image, at least 32K of which may be 0's.

  // Could encode the array using strings and a more dense encoding rather
  // than single bytes as decimal numbers.  Can use 32..126 excluding '"'.
  // It depends on whether we want the C source to be smaller, or the exe.
  return 0;
}
#endif