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

// 内部结构体定义  
typedef struct HuffmanNode_ {  
    unsigned char byte;  
    size_t frequency;  
    struct HuffmanNode_* left;  
    struct HuffmanNode_* right;  
} HuffmanNode_;  

typedef struct {  
    unsigned char byte;  
    size_t frequency;  
} FrequencyItem_;  

// 比较函数  
static int compare_nodes(const HuffmanNode_* a, const HuffmanNode_* b) {  
    if (a->frequency != b->frequency) {  
        return a->frequency < b->frequency ? -1 : 1;  
    }  
    return a->byte < b->byte ? -1 : (a->byte > b->byte ? 1 : 0);  
}  

// 创建新节点  
static HuffmanNode_* create_node(unsigned char byte, size_t frequency) {  
    HuffmanNode_* node = (HuffmanNode_*)malloc(sizeof(HuffmanNode_));  
    if (node) {  
        node->byte = byte;  
        node->frequency = frequency;  
        node->left = node->right = NULL;  
    }  
    return node;  
}  

// 统计字节频率  
static FrequencyItem_* count_frequencies(const unsigned char* data, size_t size, size_t* unique_count) {  
    FrequencyItem_* freq = (FrequencyItem_*)calloc(256, sizeof(FrequencyItem_));  
    *unique_count = 0;  

    if (!freq) return NULL;  

    // 初始化频率数组  
    for (int i = 0; i < 256; i++) {  
        freq[i].byte = (unsigned char)i;  
        freq[i].frequency = 0;  
    }  

    // 统计频率  
    for (size_t i = 0; i < size; i++) {  
        if (freq[data[i]].frequency == 0) {  
            (*unique_count)++;  
        }  
        freq[data[i]].frequency++;  
    }  

    return freq;  
}  

// 构建哈夫曼树  
static HuffmanNode_* build_huffman_tree(FrequencyItem_* freq, size_t unique_count) {  
    if (unique_count == 0) return NULL;  

    // 创建节点数组  
    HuffmanNode_** nodes = (HuffmanNode_**)malloc(unique_count * sizeof(HuffmanNode_*));  
    if (!nodes) return NULL;  

    size_t node_count = 0;  

    // 创建叶子节点  
    for (int i = 0; i < 256 && node_count < unique_count; i++) {  
        if (freq[i].frequency > 0) {  
            nodes[node_count] = create_node(freq[i].byte, freq[i].frequency);  
            if (!nodes[node_count]) {  
                for (size_t j = 0; j < node_count; j++) {  
                    free(nodes[j]);  
                }  
                free(nodes);  
                return NULL;  
            }  
            node_count++;  
        }  
    }  

    // 构建树  
    while (node_count > 1) {  
        // 找到最小的两个节点  
        int min1 = 0, min2 = 1;  
        if (compare_nodes(nodes[min1], nodes[min2]) > 0) {  
            int temp = min1;  
            min1 = min2;  
            min2 = temp;  
        }  

        for (size_t i = 2; i < node_count; i++) {  
            if (compare_nodes(nodes[i], nodes[min1]) < 0) {  
                min2 = min1;  
                min1 = i;  
            } else if (compare_nodes(nodes[i], nodes[min2]) < 0) {  
                min2 = i;  
            }  
        }  

        // 创建新的父节点  
        HuffmanNode_* parent = create_node(0, nodes[min1]->frequency + nodes[min2]->frequency);  
        if (!parent) {  
            for (size_t i = 0; i < node_count; i++) {  
                free(nodes[i]);  
            }  
            free(nodes);  
            return NULL;  
        }  

        parent->left = nodes[min1];  
        parent->right = nodes[min2];  

        // 更新节点数组  
        nodes[min1] = parent;  
        nodes[min2] = nodes[--node_count];  
    }  

    HuffmanNode_* root = nodes[0];  
    free(nodes);  
    return root;  
}  

// 生成哈夫曼编码  
static void generate_huffman_codes(HuffmanNode_* root, unsigned char* code, size_t depth,  
                                 HuffmanCode* codes, size_t* code_count) {  
    if (!root) return;  

    if (!root->left && !root->right) {  
        codes[*code_count].byte = root->byte;  
        codes[*code_count].code_len = depth;  
        
        size_t code_bytes = (depth + 7) / 8;  
        codes[*code_count].code = (unsigned char*)malloc(code_bytes);  
        if (codes[*code_count].code) {  
            memset(codes[*code_count].code, 0, code_bytes);  
            for (size_t i = 0; i < depth; i++) {  
                if (code[i]) {  
                    codes[*code_count].code[i / 8] |= (1 << (7 - (i % 8)));  
                }  
            }  
            (*code_count)++;  
        }  
        return;  
    }  

    if (depth < 256) {  
        code[depth] = 0;  
        generate_huffman_codes(root->left, code, depth + 1, codes, code_count);  
        code[depth] = 1;  
        generate_huffman_codes(root->right, code, depth + 1, codes, code_count);  
    }  
}  

// 释放哈夫曼树  
static void free_huffman_tree(HuffmanNode_* root) {  
    if (root) {  
        free_huffman_tree(root->left);  
        free_huffman_tree(root->right);  
        free(root);  
    }  
}  

// 压缩数据  
static CompressResult* compress_data(const unsigned char* content, size_t size) {  
    CompressResult* result = (CompressResult*)malloc(sizeof(CompressResult));  
    if (!result) return NULL;  

    // 初始化结果  
    result->original_size = size;  
    result->data = NULL;  
    result->codes = NULL;  
    result->size = 0;  
    result->unique_chars = 0;  

    // 统计频率  
    FrequencyItem_* freq = count_frequencies(content, size, &result->unique_chars);  
    if (!freq) {  
        free(result);  
        return NULL;  
    }  

    // 构建哈夫曼树  
    HuffmanNode_* root = build_huffman_tree(freq, result->unique_chars);  
    free(freq);  

    if (!root) {  
        free(result);  
        return NULL;  
    }  

    // 生成编码表  
    result->codes = (HuffmanCode*)malloc(result->unique_chars * sizeof(HuffmanCode));  
    if (!result->codes) {  
        free_huffman_tree(root);  
        free(result);  
        return NULL;  
    }  

    unsigned char* temp_code = (unsigned char*)malloc(256);  
    if (!temp_code) {  
        free_huffman_tree(root);  
        free(result->codes);  
        free(result);  
        return NULL;  
    }  

    size_t code_count = 0;  
    generate_huffman_codes(root, temp_code, 0, result->codes, &code_count);  
    free(temp_code);  
    free_huffman_tree(root);  

    // 计算压缩后的大小  
    size_t total_bits = 0;  
    for (size_t i = 0; i < size; i++) {  
        for (size_t j = 0; j < result->unique_chars; j++) {  
            if (result->codes[j].byte == content[i]) {  
                total_bits += result->codes[j].code_len;  
                break;  
            }  
        }  
    }  

    // 分配压缩数据空间  
    result->size = (total_bits + 7) / 8;  
    result->data = (unsigned char*)malloc(result->size);  
    if (!result->data) {  
        for (size_t i = 0; i < result->unique_chars; i++) {  
            free(result->codes[i].code);  
        }  
        free(result->codes);  
        free(result);  
        return NULL;  
    }  

    memset(result->data, 0, result->size);  

    // 写入压缩数据  
    size_t bit_pos = 0;  
    for (size_t i = 0; i < size; i++) {  
        for (size_t j = 0; j < result->unique_chars; j++) {  
            if (result->codes[j].byte == content[i]) {  
                for (size_t k = 0; k < result->codes[j].code_len; k++) {  
                    if (result->codes[j].code[k / 8] & (1 << (7 - (k % 8)))) {  
                        result->data[bit_pos / 8] |= (1 << (7 - (bit_pos % 8)));  
                    }  
                    bit_pos++;  
                }  
                break;  
            }  
        }  
    }  

    return result;  
}  


int compress_file(const char* filename,   
    const char* sender_info,   
    const char* receiver_info,  
    EncryptMethod method) {  
CompressResult* result = NULL;  
int ret = -1;  
unsigned char* content = NULL;  

// 打开并读取原文件  
FILE* file = fopen(filename, "rb");  
if (!file) {  
printf("无法打开文件: %s\n", filename);  
return -1;  
}  
fseek(file, 0, SEEK_END);  
size_t file_size = ftell(file);  
fseek(file, 0, SEEK_SET);  

content = (unsigned char*)malloc(file_size);  
if (!content) {  
printf("内存分配失败\n");  
fclose(file);  
return -1;  
}  
if (fread(content, 1, file_size, file) != file_size) {  
printf("文件读取失败\n");  
free(content);  
fclose(file);  
return -1;  
}  
fclose(file);  

// --- 加密处理 ---  
unsigned char key = 0x55;  
if (method == ENCRYPT_METHOD_OFFSET) {  
for (size_t i = 0; i < file_size; i++) {  
content[i] = (unsigned char)((content[i] + key) & 0xFF);  
}  
} else if (method == ENCRYPT_METHOD_XOR) {  
for (size_t i = 0; i < file_size; i++) {  
content[i] ^= key;  
}  
}  

// 基于加密后的数据构建哈夫曼树、统计词频，生成编码表，并压缩文件数据  
result = compress_data(content, file_size);  
if (!result) {  
printf("压缩失败\n");  
free(content);  
return -1;  
}  

char compressed_filename[256];  
char code_filename[256];  
snprintf(compressed_filename, sizeof(compressed_filename), "%s.hfm", filename);  
snprintf(code_filename, sizeof(code_filename), "%s.hfm.code", filename);  

FILE* compressed_file = fopen(compressed_filename, "wb");  
if (!compressed_file) {  
printf("无法创建压缩文件: %s\n", compressed_filename);  
goto cleanup;  
}  

// 写入头信息：加密方式标志、发送人信息与接收人信息（各一行）  
{  
char header[256];  
int header_len = snprintf(header, sizeof(header), "%d\n%s\n%s\n", (int)method, sender_info, receiver_info);  
if ((size_t)header_len != strlen(header)) {  
printf("头信息生成失败\n");  
fclose(compressed_file);  
goto cleanup;  
}  
if (fwrite(header, 1, (size_t)header_len, compressed_file) != (size_t)header_len) {  
printf("写入头信息失败\n");  
fclose(compressed_file);  
goto cleanup;  
}  
}  

// 写入压缩二进制数据（注意，此时数据为加密后的结果）  
if (fwrite(result->data, 1, result->size, compressed_file) != result->size) {  
printf("写入压缩数据失败\n");  
fclose(compressed_file);  
goto cleanup;  
}  
fclose(compressed_file);  
printf("压缩文件已保存: %s\n", compressed_filename);  

// 编码表文件保存
FILE* code_file = fopen(code_filename, "w");  
if (!code_file) {  
printf("无法创建编码表文件: %s\n", code_filename);  
goto cleanup;  
}  
fprintf(code_file, "%zu\n", result->original_size);  
for (size_t i = 0; i < result->unique_chars; i++) {  
fprintf(code_file, "0x%02X ", result->codes[i].byte);  
fprintf(code_file, "0x%02X", (unsigned char)result->codes[i].code_len);  
size_t code_bytes = (result->codes[i].code_len + 7) / 8;  
for (size_t j = 0; j < code_bytes; j++) {  
fprintf(code_file, " 0x%02X", result->codes[i].code[j]);  
}  
fprintf(code_file, "\n");  
}  
fclose(code_file);  
printf("编码表文件已保存: %s\n", code_filename);  
ret = 0;  

cleanup:  
if (content) free(content);  
if (result) {  
if (result->data) free(result->data);  
if (result->codes) {  
for (size_t i = 0; i < result->unique_chars; i++) {  
 if (result->codes[i].code) free(result->codes[i].code);  
}  
free(result->codes);  
}  
free(result);  
}  
return ret;  
}

// 释放压缩结果  
void free_compress_result(CompressResult* result) {  
    if (!result) return;  

    if (result->data) {  
        free(result->data);  
    }  

    if (result->codes) {  
        for (size_t i = 0; i < result->unique_chars; i++) {  
            if (result->codes[i].code) {  
                free(result->codes[i].code);  
            }  
        }  
        free(result->codes);  
    }  

    free(result);  
}