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

// --- Алгоритм Run-Length Encoding (RLE) ---

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

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

    unsigned char current_byte, previous_byte;
    size_t count = 0;

    if (fread(&previous_byte, 1, 1, infile) != 1) {
        fclose(infile);
        fclose(outfile);
        return 0;  // Порожній файл
    }

    count = 1;

    while (fread(&current_byte, 1, 1, infile) == 1) {
        if (current_byte == previous_byte && count < 255) {
            count++;
        } else {
            fwrite(&previous_byte, 1, 1, outfile);
            fwrite(&count, 1, 1, outfile);
            previous_byte = current_byte;
            count = 1;
        }
    }

    fwrite(&previous_byte, 1, 1, outfile);
    fwrite(&count, 1, 1, outfile);

    fclose(infile);
    fclose(outfile);

    return 1;
}

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

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

    unsigned char current_byte;
    unsigned char count;

    while (fread(&current_byte, 1, 1, infile) == 1) {
        if (fread(&count, 1, 1, infile) != 1) {
            perror("Malformed input file");
            fclose(infile);
            fclose(outfile);
            return -1;
        }

        for (size_t i = 0; i < count; i++) {
            fwrite(&current_byte, 1, 1, outfile);
        }
    }

    fclose(infile);
    fclose(outfile);

    return 1;
}

// --- Алгоритм Хаффмана (Huffman Coding) ---

typedef struct Node {
    unsigned char symbol;
    size_t frequency;
    struct Node *left, *right;
} Node;

typedef struct {
    unsigned char symbol;
    char *code;
} HuffmanCode;

int compare_nodes(const void *a, const void *b) {
    return (*(Node**)a)->frequency - (*(Node**)b)->frequency;
}

Node* create_huffman_tree(unsigned char *data, size_t size) {
    size_t freq[256] = {0};
    for (size_t i = 0; i < size; i++) {
        freq[data[i]]++;
    }

    Node *nodes[256];
    size_t node_count = 0;

    for (int i = 0; i < 256; i++) {
        if (freq[i] > 0) {
            nodes[node_count] = malloc(sizeof(Node));
            if (!nodes[node_count]) {
                perror("Memory allocation failed");
                return NULL;
            }
            nodes[node_count]->symbol = (unsigned char)i;
            nodes[node_count]->frequency = freq[i];
            nodes[node_count]->left = nodes[node_count]->right = NULL;
            node_count++;
        }
    }

    while (node_count > 1) {
        qsort(nodes, node_count, sizeof(Node*), compare_nodes);

        Node* left = nodes[0];
        Node* right = nodes[1];

        Node* parent = malloc(sizeof(Node));
        if (!parent) {
            perror("Memory allocation failed");
            return NULL;
        }

        parent->symbol = 0;
        parent->frequency = left->frequency + right->frequency;
        parent->left = left;
        parent->right = right;

        memmove(nodes, nodes + 2, (node_count - 2) * sizeof(Node*));
        nodes[node_count - 2] = parent;
        node_count--;
    }

    return nodes[0];
}

void generate_huffman_codes(Node* root, HuffmanCode* codes, char* current_code, int depth) {
    if (!root) return;

    if (root->left == NULL && root->right == NULL) {
        current_code[depth] = '\0';
        codes[root->symbol].symbol = root->symbol;
        codes[root->symbol].code = strdup(current_code);
        return;
    }

    current_code[depth] = '0';
    generate_huffman_codes(root->left, codes, current_code, depth + 1);

    current_code[depth] = '1';
    generate_huffman_codes(root->right, codes, current_code, depth + 1);
}

void serialize_huffman_tree(Node* root, FILE* outfile) {
    if (!root) return;

    if (root->left == NULL && root->right == NULL) {
        fputc('L', outfile);
        fputc(root->symbol, outfile);
    } else {
        fputc('I', outfile);
        serialize_huffman_tree(root->left, outfile);
        serialize_huffman_tree(root->right, outfile);
    }
}

Node* rebuild_huffman_tree(unsigned char* tree_data, size_t size) {
    (void)size;  // Позначаємо параметр як тимчасово невикористаний
    size_t index = 0;

    Node* build_tree_recursively(unsigned char* data, size_t* index) {
        if (data[*index] == 'L') { // Лист (Leaf)
            (*index)++;
            Node* leaf = malloc(sizeof(Node));
            if (!leaf) {
                perror("Memory allocation failed");
                return NULL;
            }
            leaf->symbol = data[*index];
            leaf->frequency = 0; // частота не потрібна для декомпресії
            leaf->left = leaf->right = NULL;
            (*index)++;
            return leaf;
        } else if (data[*index] == 'I') { // Вузол (Internal)
            (*index)++;
            Node* internal = malloc(sizeof(Node));
            if (!internal) {
                perror("Memory allocation failed");
                return NULL;
            }
            internal->symbol = 0; // внутрішні вузли не мають символів
            internal->frequency = 0;
            internal->left = build_tree_recursively(data, index);
            internal->right = build_tree_recursively(data, index);
            return internal;
        }
        return NULL;
    }

    return build_tree_recursively(tree_data, &index);
}

void free_huffman_tree(Node* root) {
    if (!root) return;
    free_huffman_tree(root->left);
    free_huffman_tree(root->right);
    free(root);
}

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

    fseek(infile, 0, SEEK_END);
    size_t file_size = ftell(infile);
    fseek(infile, 0, SEEK_SET);

    unsigned char* data = malloc(file_size);
    if (!data) {
        fclose(infile);
        perror("Memory allocation failed");
        return -1;
    }

    fread(data, 1, file_size, infile);
    fclose(infile);

    Node* root = create_huffman_tree(data, file_size);
    if (!root) {
        free(data);
        return -1;
    }

    HuffmanCode codes[256] = {0};
    char current_code[256];
    generate_huffman_codes(root, codes, current_code, 0);

    FILE *outfile = fopen(output_file_name, "wb");
    if (!outfile) {
        free(data);
        free_huffman_tree(root);
        return -1;
    }

    serialize_huffman_tree(root, outfile);

    for (size_t i = 0; i < file_size; i++) {
        const char* code = codes[data[i]].code;
        for (size_t j = 0; code[j] != '\0'; j++) {
            fputc(code[j], outfile);
        }
    }

    fclose(outfile);
    free(data);
    free_huffman_tree(root);

    for (int i = 0; i < 256; i++) {
        free(codes[i].code);
    }

    return 1;
}

// Декомпресія (приклад потребує уточнень залежно від специфіки формату)
int decompress_1(const char* input_file_name, const char* output_file_name) {
    (void)input_file_name;
    (void)output_file_name;
    return -1;  // Поки що не реалізовано.
}