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

#define BUFFER_SIZE 4096
#define MAX_SYMBOLS 257

// Макрос для обмена двух узлов
#define SWAP_NODES(a, b) { Node* temp = a; a = b; b = temp; }

// Определение структуры узла дерева
typedef struct Node {
    int symbol;                
    unsigned int frequency;     
    struct Node *left, *right;  
} Node;

// Функция для создания нового узла
Node* create_node(int symbol, unsigned int frequency) {
    Node* node = (Node*)malloc(sizeof(Node));
    node->symbol = symbol;
    node->frequency = frequency;
    node->left = node->right = NULL;
    return node;
}

// Функция для построения дерева Хаффмана
Node* build_huffman_tree(const unsigned int* frequencies) {
    Node* nodes[MAX_SYMBOLS];
    int node_count = 0;

    // Создаем узлы для всех символов с ненулевой частотой
    for (int i = 0; i < MAX_SYMBOLS; i++) {
        if (frequencies[i] > 0) {
            nodes[node_count++] = create_node(i, frequencies[i]);
        }
    }

    // Объединяем узлы в дерево
    while (node_count > 1) {
        // Сортируем узлы по частоте
        for (int i = 0; i < node_count - 1; i++) {
            for (int j = i + 1; j < node_count; j++) {
                if (nodes[i]->frequency > nodes[j]->frequency) {
                    SWAP_NODES(nodes[i], nodes[j]);
                }
            }
        }

        // Объединяем два узла с наименьшей частотой
        Node* left = nodes[0];
        Node* right = nodes[1];
        Node* parent = create_node(-1, left->frequency + right->frequency);
        parent->left = left;
        parent->right = right;

        // Заменяем объединенные узлы новым родительским узлом
        nodes[0] = parent;
        nodes[1] = nodes[--node_count];
    }

    return nodes[0];
}

// Рекурсивная функция для генерации кодов Хаффмана
void generate_huffman_codes(Node* root, char* code, int depth, char codes[MAX_SYMBOLS][MAX_SYMBOLS]) {
    if (!root->left && !root->right) {
        code[depth] = '\0'; // Завершаем код символа
        strcpy(codes[root->symbol], code);
        return;
    }
    if (root->left) {
        code[depth] = '0'; // Добавляем бит '0' для левого поддерева
        generate_huffman_codes(root->left, code, depth + 1, codes);
    }
    if (root->right) {
        code[depth] = '1'; // Добавляем бит '1' для правого поддерева
        generate_huffman_codes(root->right, code, depth + 1, codes);
    }
}

// Функция для освобождения памяти, выделенной под дерево Хаффмана
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, const char* output_file) {
    FILE* input = fopen(input_file, "rb");
    FILE* output = fopen(output_file, "wb");
    if (!input || !output) return -1;

    unsigned int frequencies[MAX_SYMBOLS] = {0};
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;

    // Подсчет частот символов
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, input)) > 0) {
        for (size_t i = 0; i < bytes_read; i++) {
            frequencies[buffer[i]]++;
        }
    }
    frequencies[256] = 1; // Добавляем маркер EOF

    Node* root = build_huffman_tree(frequencies);
    if (!root) return -1;

    // Генерация кодов Хаффмана
    char codes[MAX_SYMBOLS][MAX_SYMBOLS] = {{0}};
    char code[MAX_SYMBOLS] = {0};
    generate_huffman_codes(root, code, 0, codes);

    // Записываем частоты в выходной файл
    fwrite(frequencies, sizeof(frequencies[0]), MAX_SYMBOLS, output);

    // Сжимаем данные
    rewind(input);
    unsigned char current_byte = 0;
    int bit_count = 0;

    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, input)) > 0) {
        for (size_t i = 0; i < bytes_read; i++) {
            char* symbol_code = codes[buffer[i]];
            for (size_t j = 0; symbol_code[j] != '\0'; j++) {
                current_byte = (current_byte << 1) | (symbol_code[j] - '0');
                bit_count++;
                if (bit_count == 8) {
                    fwrite(&current_byte, 1, 1, output);
                    current_byte = 0;
                    bit_count = 0;
                }
            }
        }
    }

    // Записываем маркер EOF
    char* eof_code = codes[256];
    for (size_t j = 0; eof_code[j] != '\0'; j++) {
        current_byte = (current_byte << 1) | (eof_code[j] - '0');
        bit_count++;
        if (bit_count == 8) {
            fwrite(&current_byte, 1, 1, output);
            current_byte = 0;
            bit_count = 0;
        }
    }
    if (bit_count > 0) {
        current_byte <<= (8 - bit_count);
        fwrite(&current_byte, 1, 1, output);
    }

    fclose(input);
    fclose(output);
    free_huffman_tree(root);
    return 0;
}

// Функция декомпрессии данных с использованием алгоритма Хаффмана
int decompress_1(const char* input_file, const char* output_file) {
    FILE* input = fopen(input_file, "rb");
    FILE* output = fopen(output_file, "wb");
    if (!input || !output) return -1;

    unsigned int frequencies[MAX_SYMBOLS] = {0};
    fread(frequencies, sizeof(frequencies[0]), MAX_SYMBOLS, input);
    Node* root = build_huffman_tree(frequencies);
    if (!root) return -1;

    Node* current = root;
    unsigned char byte;
    int bit;

    // Читаем и декодируем символы
    while (fread(&byte, 1, 1, input) == 1) {
        for (bit = 7; bit >= 0; bit--) {
            current = (byte & (1 << bit)) ? current->right : current->left;

            if (!current->left && !current->right) {
                if (current->symbol == 256) { // Маркер EOF
                    fclose(input);
                    fclose(output);
                    free_huffman_tree(root);
                    return 0;
                }
                fwrite(&current->symbol, 1, 1, output);
                current = root;
            }
        }
    }

    fclose(input);
    fclose(output);
    free_huffman_tree(root);
    return 0;
}

int compress_2(const char* input_file_name, const char* output_file_name){
    return 0;
}

int decompress_2(const char* input_file_name, const char* output_file_name){
    return 0;
}