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

#define BUFFER_SIZE 1024
#define LZ78_DICT_SIZE 256

// Function prototypes
int compress(const char* input_file_name, const char* output_file_name);
int decompress(const char* input_file_name, const char* output_file_name);
void lz78_compress(FILE* in, FILE* out);
void lz78_decompress(FILE* in, FILE* out);

void print_usage() {
    printf("Usage:\n");
    printf("  ./compressor -c infile outfile   Compress infile into outfile\n");
    printf("  ./compressor -d compressed uncompressed   Decompress compressed into uncompressed\n");
}

int main(int argc, char* argv[]) {
    if (argc != 4) {
        print_usage();
        return 1;
    }

    const char* mode = argv[1];
    const char* infile = argv[2];
    const char* outfile = argv[3];

    if (strcmp(mode, "-c") == 0) {
        return compress(infile, outfile);
    } else if (strcmp(mode, "-d") == 0) {
        return decompress(infile, outfile);
    } else {
        print_usage();
        return 1;
    }
}

int compress(const char* input_file_name, const char* output_file_name) {
    FILE* input_file = fopen(input_file_name, "rb");
    if (!input_file) {
        perror("Error opening input file");
        return -1;
    }

    FILE* output_file = fopen(output_file_name, "wb");
    if (!output_file) {
        perror("Error opening output file");
        fclose(input_file);
        return -1;
    }

    lz78_compress(input_file, output_file);

    fclose(input_file);
    fclose(output_file);

    printf("File compressed successfully.\n");
    return 0;
}

int decompress(const char* input_file_name, const char* output_file_name) {
    FILE* input_file = fopen(input_file_name, "rb");
    if (!input_file) {
        perror("Error opening input file");
        return -1;
    }

    FILE* output_file = fopen(output_file_name, "wb");
    if (!output_file) {
        perror("Error opening output file");
        fclose(input_file);
        return -1;
    }

    lz78_decompress(input_file, output_file);

    fclose(input_file);
    fclose(output_file);

    printf("File decompressed successfully.\n");
    return 0;
}

// LZ78 compression
struct dict {
    int index;
    struct dict* characters[LZ78_DICT_SIZE];
};

void delete_dict(struct dict* tree) {
    if (!tree) {
        return;
    }
    for (int i = 0; i < LZ78_DICT_SIZE; i++) {
        delete_dict(tree->characters[i]);
    }
    free(tree);
}

void lz78_compress(FILE* in, FILE* out) {
    struct dict* root = calloc(1, sizeof(struct dict));
    struct dict* current = root;
    int index = 1;
    int c;

    while ((c = fgetc(in)) != EOF) {
        if (current->characters[c]) {
            current = current->characters[c];
        } else {
            current->characters[c] = calloc(1, sizeof(struct dict));
            current->characters[c]->index = index++;
            fprintf(out, "%d-%c\n", current->index, c);
            current = root;
        }
    }

    delete_dict(root);
}

// LZ78 decompression
void lz78_decompress(FILE* in, FILE* out) {
    char buffer[LZ78_DICT_SIZE][BUFFER_SIZE] = {""};
    int index;
    char c;

    while (fscanf(in, "%d-%c\n", &index, &c) == 2) {
        if (index > 0) {
            fprintf(out, "%s", buffer[index]);
        }
        fputc(c, out);
        snprintf(buffer[index], BUFFER_SIZE, "%s%c", buffer[index], c);
    }
}