#include "decompress.h"  
#include "compress.h"  
#include <stdio.h>  
#include <stdlib.h>  
#include <string.h>  
#include <stdint.h>  
#include <time.h>  

// FNV-1a 64位哈希算法常量  
#define FNV1A_64_INIT 0xcbf29ce484222325ULL  
#define FNV1A_64_PRIME 0x100000001b3ULL  

// 哈夫曼树节点结构  
typedef struct HuffmanNode {  
    unsigned char byte;  
    struct HuffmanNode* left;  
    struct HuffmanNode* right;  
    int is_leaf;  
} HuffmanNode;  

// 解码表结构  
typedef struct {  
    size_t original_size;  
    uint32_t original_hash;  
    HuffmanNode* root;  
} DecodingTable;  

// 性能统计结构体  
typedef struct {  
    clock_t start_time;  
    clock_t end_time;  
    size_t processed_bytes;  
    double compression_ratio;  
} PerformanceStats;  

// 哈希表项结构  
typedef struct {  
    uint32_t code;  
    unsigned char code_len;  
    unsigned char byte;  
} HashEntry;  

// 哈希表结构  
typedef struct {  
    HashEntry* entries;  
    size_t size;  
    size_t capacity;  
} HashTable;  

// 哈希计算函数  
static uint64_t calculate_hash(const unsigned char* data, size_t size) {  
    uint64_t hash = FNV1A_64_INIT;  
    for (size_t i = 0; i < size; i++) {  
        hash ^= data[i];  
        hash *= FNV1A_64_PRIME;  
    }  
    return hash;  
}  

// 创建哈夫曼树节点  
static HuffmanNode* create_node(unsigned char byte, int is_leaf) {  
    HuffmanNode* node = (HuffmanNode*)malloc(sizeof(HuffmanNode));  
    if (node) {  
        node->byte = byte;  
        node->left = node->right = NULL;  
        node->is_leaf = is_leaf;  
    }  
    return node;  
}  

// 初始化哈希表  
static HashTable* init_hash_table(size_t initial_capacity) {  
    HashTable* table = (HashTable*)malloc(sizeof(HashTable));  
    if (!table) return NULL;  

    table->entries = (HashEntry*)calloc(initial_capacity, sizeof(HashEntry));  
    if (!table->entries) {  
        free(table);  
        return NULL;  
    }  

    table->size = 0;  
    table->capacity = initial_capacity;  
    
    // 初始化所有条目  
    for (size_t i = 0; i < initial_capacity; i++) {  
        table->entries[i].code_len = 0;  
    }  
    
    return table;  
}  

// 向哈希表中添加项  
static void hash_table_add(HashTable* table, uint32_t code,  
    unsigned char code_len, unsigned char byte) {  
if (!table || table->size >= table->capacity) return;  

size_t index = code % table->capacity;  
size_t original_index = index;  

// 线性探测  
while (table->entries[index].code_len != 0) {  
index = (index + 1) % table->capacity;  
if (index == original_index) return; // 表已满  
}  

table->entries[index].code = code;  
table->entries[index].code_len = code_len;  
table->entries[index].byte = byte;  
table->size++;  
} 
// 在哈希表中查找  
static unsigned char hash_table_find(HashTable* table, uint32_t code,  
    unsigned char code_len) {  
if (!table) return 0;  

size_t index = code % table->capacity;  
size_t original_index = index;  

do {  
if (table->entries[index].code_len == code_len &&  
table->entries[index].code == code) {  
return table->entries[index].byte;  
}  
if (table->entries[index].code_len == 0) {  
break; // 找到空位置，说明不存在  
}  
index = (index + 1) % table->capacity;  
} while (index != original_index);  

return 0;  
}

// 释放哈希表  
static void free_hash_table(HashTable* table) {  
    if (table) {  
        free(table->entries);  
        free(table);  
    }  
}  

// 基于哈希的解码方法  
static size_t decode_with_hash(const unsigned char* compressed_data,  
    size_t compressed_size,  
    unsigned char* decoded_data,  
    size_t original_size,  
    const CodeEntry* codes,  
    size_t unique_chars) {  
// 调试信息  
printf("压缩数据大小: %zu 字节\n", compressed_size);  
printf("原始数据大小: %zu 字节\n", original_size);  
printf("唯一字符数: %zu\n", unique_chars);  

HashTable* table = init_hash_table(unique_chars * 4); // 增加哈希表大小  
if (!table) {  
printf("哈希表初始化失败\n");  
return 0;  
}  

// 打印编码表信息进行调试  
printf("编码表内容:\n");  
for (size_t i = 0; i < unique_chars; i++) {  
printf("字节: 0x%02X, 长度: %d, 编码: 0x%X\n",  
codes[i].byte, codes[i].code_len, codes[i].code);  
hash_table_add(table, codes[i].code, codes[i].code_len, codes[i].byte);  
}  

BitReader reader;  
init_bit_reader(&reader, compressed_data, compressed_size * 8);  
size_t bytes_written = 0;  
uint32_t buffer = 0;  
size_t bits_in_buffer = 0;  

while (bytes_written < original_size && reader.current_bit < reader.total_bits) {  
// 填充缓冲区  
while (bits_in_buffer < 32 && reader.current_bit < reader.total_bits) {  
uint32_t bit = (compressed_data[reader.current_bit / 8] >> (7 - (reader.current_bit % 8))) & 1;  
buffer = (buffer << 1) | bit;  
bits_in_buffer++;  
reader.current_bit++;  
}  

// 尝试匹配编码  
int found = 0;  
uint32_t temp_buffer = buffer;  
for (unsigned char len = 1; len <= bits_in_buffer && len <= 32; len++) {  
uint32_t code = temp_buffer >> (bits_in_buffer - len);  
unsigned char byte = hash_table_find(table, code, len);  

if (byte) {  
decoded_data[bytes_written++] = byte;  
// 从缓冲区移除已处理的位  
buffer = (buffer << len) >> len;  
bits_in_buffer -= len;  
found = 1;  
break;  
}  
}  

if (!found) {  
if (bits_in_buffer >= 32) {  
bits_in_buffer--; // 丢弃一位并继续  
buffer = buffer << 1;  
}  
}  
}  

// 输出解码统计信息  
printf("解码统计:\n");  
printf("处理的位数: %zu\n", reader.current_bit);  
printf("写入的字节数: %zu\n", bytes_written);  

free_hash_table(table);  
return bytes_written;  
} 

// 初始化位流读取器  
void init_bit_reader(BitReader* reader, const unsigned char* data, size_t size) {  
    reader->data = data;  
    reader->total_bits = size * 8;  
    reader->current_bit = 0;  
}  

// 从位流中读取指定数量的位  
uint32_t read_bits(BitReader* reader, int num_bits) {  
    uint32_t result = 0;  
    for (int i = 0; i < num_bits; i++) {  
        size_t byte_pos = reader->current_bit / 8;  
        size_t bit_pos = 7 - (reader->current_bit % 8);  
        int bit = (reader->data[byte_pos] >> bit_pos) & 1;  
        result = (result << 1) | bit;  
        reader->current_bit++;  
    }  
    return result;  
}  

// 从编码表构建解码信息（方法1）  
static DecodingTable* build_decoding_table(const char* code_filename) {  
    FILE* file = fopen(code_filename, "r");  
    if (!file) {  
        printf("无法打开编码表文件: %s\n", code_filename);  
        return NULL;  
    }  

    DecodingTable* table = (DecodingTable*)malloc(sizeof(DecodingTable));  
    if (!table) {  
        fclose(file);  
        return NULL;  
    }  

    if (fscanf(file, "%zu\n", &table->original_size) != 1) {  
        printf("读取原始文件大小失败\n");  
        free(table);  
        fclose(file);  
        return NULL;  
    }  

    table->root = create_node(0, 0);  
    if (!table->root) {  
        free(table);  
        fclose(file);  
        return NULL;  
    }  

    char line[1024];  
    while (fgets(line, sizeof(line), file)) {  
        char* token = strtok(line, " \t\n");  
        if (!token) continue;  
        unsigned int byte_val = (unsigned int)strtoul(token, NULL, 16);  

        token = strtok(NULL, " \t\n");  
        if (!token) continue;  
        unsigned int len_val = (unsigned int)strtoul(token, NULL, 16);  

        HuffmanNode* current = table->root;  
        unsigned char byte = (unsigned char)byte_val;  
        size_t code_len = (size_t)len_val;  
        size_t code_bytes = (code_len + 7) / 8;  

        unsigned char code[32] = {0};  
        for (size_t i = 0; i < code_bytes; i++) {  
            token = strtok(NULL, " \t\n");  
            if (!token) break;  
            unsigned int code_byte = (unsigned int)strtoul(token, NULL, 16);  
            code[i] = (unsigned char)code_byte;  
        }  

        for (size_t i = 0; i < code_len; i++) {  
            int bit = (code[i / 8] >> (7 - (i % 8))) & 1;  
            if (i == code_len - 1) {  
                HuffmanNode* leaf = create_node(byte, 1);  
                if (bit == 0) {  
                    current->left = leaf;  
                } else {  
                    current->right = leaf;  
                }  
            } else {  
                if (bit == 0) {  
                    if (!current->left) {  
                        current->left = create_node(0, 0);  
                    }  
                    current = current->left;  
                } else {  
                    if (!current->right) {  
                        current->right = create_node(0, 0);  
                    }  
                    current = current->right;  
                }  
            }  
        }  
    }  

    fclose(file);  
    return table;  
}  

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

// 释放解码表  
static void free_decoding_table(DecodingTable* table) {  
    if (table) {  
        free_huffman_tree(table->root);  
        free(table);  
    }  
}  

// 读取编码表（方法2）  
CodeEntry* load_code_table(const char* code_filename,  
                                size_t* unique_chars,  
                                size_t* original_size) {  
    FILE* file = fopen(code_filename, "r");  
    if (!file) {  
        printf("无法打开编码表文件: %s\n", code_filename);  
        return NULL;  
    }  

    if (fscanf(file, "%zu\n", original_size) != 1) {  
        printf("读取原始文件大小失败\n");  
        fclose(file);  
        return NULL;  
    }  

    CodeEntry* codes = (CodeEntry*)malloc(256 * sizeof(CodeEntry));  
    if (!codes) {  
        fclose(file);  
        return NULL;  
    }  

    *unique_chars = 0;  
    char line[1024];  
    while (fgets(line, sizeof(line), file)) {  
        unsigned int byte_val, len_val;  
        if (sscanf(line, "0x%x 0x%x", &byte_val, &len_val) != 2) {  
            continue;  
        }  

        codes[*unique_chars].byte = (unsigned char)byte_val;  
        codes[*unique_chars].code_len = (unsigned char)len_val;  

        uint32_t code = 0;  
        char* ptr = line;  
        size_t code_bytes_count = (len_val + 7) / 8;  
        size_t bytes_read = 0;  

        for (int i = 0; i < 2; i++) {  
            while (*ptr && *ptr != ' ') ptr++;  
            while (*ptr && *ptr == ' ') ptr++;  
        }  

        while (*ptr && bytes_read < code_bytes_count) {  
            unsigned int byte;  
            if (sscanf(ptr, "0x%x", &byte) == 1) {  
                code = (code << 8) | (byte & 0xFF);  
                bytes_read++;  
            }  
            while (*ptr && *ptr != ' ') ptr++;  
            while (*ptr && *ptr == ' ') ptr++;  
        }  

        if (len_val < bytes_read * 8) {  
            code >>= (bytes_read * 8 - len_val);  
        }  
        codes[*unique_chars].code = code;  

        (*unique_chars)++;  
    }  

    fclose(file);  
    return codes;  
}  

// 性能统计打印函数  
static void print_performance_stats(PerformanceStats stats, double decode_time, size_t bytes_written) {  
    double overall_time = (double)(stats.end_time - stats.start_time) / CLOCKS_PER_SEC;  
    double overall_speed = stats.processed_bytes / overall_time / 1024.0;  
    double decode_speed = (double)bytes_written / decode_time / 1024.0;  

    printf("\n=== 性能统计 ===\n");  
    printf("压缩率: %.2f%%\n", stats.compression_ratio);  
    printf("解码操作耗时: %.3f 秒\n", decode_time);  
    printf("解码部分处理速度: %.2f KB/s\n", decode_speed);  
    printf("===============\n");  
}  

// 用户提示函数  
static int ui_prompt_decoding_method() {  
    int method;  
    printf("\n请选择解码方式:\n");  
    printf("1. 使用解码树解码\n");  
    printf("2. 使用编码表直接匹配\n");  
    printf("3. 使用哈希表解码\n");  
    printf("请输入选择 (1-3): ");  
    scanf("%d", &method);  
    return method;  
}  

// 解压缩文件的实现  
int decompress_file(const char* filename,  
                      const char* sender_info,  
                      const char* receiver_info) {  

    PerformanceStats stats;  
    stats.start_time = clock();  // 记录整体起始时间  

    char code_filename[256];  
    char output_filename[256];  
    char original_filename[256];  
    strncpy(original_filename, filename, sizeof(original_filename));  
    char* ext = strstr(original_filename, ".hfm");  
    if (ext)  
        *ext = '\0';  
    snprintf(code_filename, sizeof(code_filename), "%s.code", filename);  
    snprintf(output_filename, sizeof(output_filename), "%s_recovered.txt", original_filename);  

    printf("正在构建解码信息...\n");  
    // 构造解码树（方法1）  
    DecodingTable* table = build_decoding_table(code_filename);  
    if (!table) {  
        printf("构建解码树失败\n");  
        return -1;  
    }  

    // 打开压缩文件  
    FILE* input_file = fopen(filename, "rb");  
    if (!input_file) {  
        printf("无法打开压缩文件: %s\n", filename);  
        free_decoding_table(table);  
        return -1;  
    }  

    char enc_method_line[16];  
    char file_sender_info[128], file_receiver_info[128];  
    // 读取加密方式标志行  
    if (!fgets(enc_method_line, sizeof(enc_method_line), input_file)) {  
        printf("读取加密方式失败\n");  
        fclose(input_file);  
        free_decoding_table(table);  
        return -1;  
    }  
    int enc_method = atoi(enc_method_line);  // 0: 偏移加密, 1: XOR 加密  

    // 读取发送人和接收人信息  
    if (!fgets(file_sender_info, sizeof(file_sender_info), input_file) ||  
        !fgets(file_receiver_info, sizeof(file_receiver_info), input_file)) {  
        printf("读取头信息失败\n");  
        fclose(input_file);  
        free_decoding_table(table);  
        return -1;  
    }  
    // 去除行末换行符  
    file_sender_info[strcspn(file_sender_info, "\n")] = '\0';  
    file_receiver_info[strcspn(file_receiver_info, "\n")] = '\0';  

    // 验证输入信息与文件头信息是否一致  
    if (strcmp(file_sender_info, sender_info) != 0 ||  
        strcmp(file_receiver_info, receiver_info) != 0) {  
        printf("收发人信息不匹配，解压拒绝！\n");  
        fclose(input_file);  
        free_decoding_table(table);  
        return -1;  
    }  

    // 定位压缩数据开始位置  
    long data_start = ftell(input_file);  
    fseek(input_file, 0, SEEK_END);  
    size_t file_total_size = ftell(input_file);  
    fseek(input_file, data_start, SEEK_SET);  
    size_t compressed_size = file_total_size - data_start;  

    unsigned char* compressed_data = (unsigned char*)malloc(compressed_size);  
    if (!compressed_data) {  
        printf("内存分配失败\n");  
        fclose(input_file);  
        free_decoding_table(table);  
        return -1;  
    }  
    if (fread(compressed_data, 1, compressed_size, input_file) != compressed_size) {  
        printf("读取压缩数据失败\n");  
        free(compressed_data);  
        fclose(input_file);  
        free_decoding_table(table);  
        return -1;  
    }  
    fclose(input_file);  

    // 提示用户选择解码方式  
    printf("请选择解码方式：\n");  
    printf("  1 - 基于解码树的逐位解码\n");  
    printf("  2 - 基于编码表直接匹配\n");  
    printf("  3 - 基于哈希的快速解码\n");  
    printf("请输入解码方式选择(1或2或3): ");  
    int decode_method = 0;  
    if (scanf("%d", &decode_method) != 1) {  
        printf("读取选择失败\n");  
        free(compressed_data);  
        free_decoding_table(table);  
        return -1;  
    }  
    while (getchar() != '\n');  // 清空输入缓冲区  

    // 为解码后的数据分配缓冲区：大小为原始文件大小  
    unsigned char* decoded_data = (unsigned char*)malloc(table->original_size);  
    if (!decoded_data) {  
        printf("内存分配失败\n");  
        free(compressed_data);  
        free_decoding_table(table);  
        return -1;  
    }  
    size_t bytes_written = 0;  

    // 记录核心解码操作的开始时间  
    clock_t decode_start = clock();  

    if (decode_method == 1) {  
        // 方法1：使用解码树进行逐位解码  
        printf("使用解码树进行解码...\n");  
        size_t total_bits = compressed_size * 8;  
        size_t current_bit = 0;  
        HuffmanNode* current = table->root;  
        while (current_bit < total_bits && bytes_written < table->original_size) {  
            size_t byte_pos = current_bit / 8;  
            size_t bit_pos = 7 - (current_bit % 8);  
            int bit = (compressed_data[byte_pos] >> bit_pos) & 1;  
            current = (bit == 0) ? current->left : current->right;  
            if (current && current->is_leaf) {  
                decoded_data[bytes_written++] = current->byte;  
                current = table->root;  
            }  
            current_bit++;  
        }  
    } 
    else if (decode_method == 2) {  
        printf("使用编码表直接匹配进行解码...\n");  
        size_t unique_chars = 0, orig_size = 0;  
        CodeEntry* codes = load_code_table(code_filename, &unique_chars, &orig_size);  
        if (!codes) {  
            printf("加载编码表失败\n");  
            free(compressed_data);  
            free(decoded_data);  
            free_decoding_table(table);  
            return -1;  
        }  

        BitReader reader;  
        init_bit_reader(&reader, compressed_data, compressed_size);  
        
        // 主解码循环  
        while (reader.current_bit < reader.total_bits && bytes_written < orig_size) {  
            int found = 0;  
            
            // 尝试匹配每个编码  
            for (size_t i = 0; i < unique_chars; i++) {  
                uint32_t bits = read_bits(&reader, codes[i].code_len);  
                reader.current_bit -= codes[i].code_len;  
                
                if (bits == codes[i].code) {  
                    decoded_data[bytes_written++] = codes[i].byte;  
                    reader.current_bit += codes[i].code_len;  
                    found = 1;  
                    break;  
                }  
            }  
            
            if (!found) {  
                reader.current_bit++;  
            }  
        }  
        
        free(codes);  
    } 
    else if (decode_method == 3) {  
        printf("使用哈希表进行解码...\n");  
        
        // 构造编码表文件名  
        char code_filename[256];  
        snprintf(code_filename, sizeof(code_filename), "%s.code", filename);  
        
        // 加载编码表信息  
        size_t unique_chars = 0;  
        size_t orig_size = 0;  
        CodeEntry* codes = NULL;  
        
        // 从编码表文件加载数据  
        FILE* code_file = fopen(code_filename, "r");  
        if (!code_file) {  
            printf("无法打开编码表文件\n");  
            free(compressed_data);  
            free(decoded_data);  
            free_decoding_table(table);  
            return -1;  
        }  
    
        // 读取原始文件大小  
        if (fscanf(code_file, "%zu\n", &orig_size) != 1) {  
            printf("读取原始文件大小失败\n");  
            fclose(code_file);  
            free(compressed_data);  
            free(decoded_data);  
            free_decoding_table(table);  
            return -1;  
        }  
    
        // 先计算编码表项数  
        char line[256];  
        unique_chars = 0;  
        while (fgets(line, sizeof(line), code_file)) {  
            if (strlen(line) > 0) {  
                unique_chars++;  
            }  
        }  
        rewind(code_file); // 重置文件指针  
        fscanf(code_file, "%zu\n", &orig_size); // 重新跳过第一行  
    
        // 分配编码表空间  
        codes = (CodeEntry*)malloc(unique_chars * sizeof(CodeEntry));  
        if (!codes) {  
            printf("分配编码表内存失败\n");  
            fclose(code_file);  
            free(compressed_data);  
            free(decoded_data);  
            free_decoding_table(table);  
            return -1;  
        }  
    
        // 读取编码表内容  
        size_t index = 0;  
        unsigned int byte_val, code_len;  
        while (index < unique_chars && fscanf(code_file, "0x%x 0x%x", &byte_val, &code_len) == 2) {  
            codes[index].byte = (unsigned char)byte_val;  
            codes[index].code_len = (unsigned char)code_len;  
            
            // 读取编码字节  
            uint32_t code = 0;  
            unsigned int code_byte;  
            size_t code_bytes = (code_len + 7) / 8; // 计算需要读取的字节数  
            
            for (size_t i = 0; i < code_bytes; i++) {  
                if (fscanf(code_file, " 0x%x", &code_byte) != 1) {  
                    printf("读取编码字节失败\n");  
                    fclose(code_file);  
                    free(codes);  
                    free(compressed_data);  
                    free(decoded_data);  
                    free_decoding_table(table);  
                    return -1;  
                }  
                code = (code << 8) | code_byte;  
            }  
            
            // 如果编码长度不是8的倍数，需要右移对齐  
            if (code_len % 8 != 0) {  
                code >>= (8 - (code_len % 8));  
            }  
            
            codes[index].code = code;  
            index++;  
            
            // 读取换行符  
            while (fgetc(code_file) != '\n' && !feof(code_file));  
        }  
    
        fclose(code_file);  
    
        // 执行哈希解码  
        // 添加进度显示  
        printf("开始解码处理...\n");  
        bytes_written = decode_with_hash(compressed_data, compressed_size,  
                                   decoded_data, table->original_size,  
                                   codes, unique_chars);  
        printf("解码完成，共处理 %zu 字节\n", bytes_written);
      


    
        // 检查解码结果  
        if (bytes_written != table->original_size) {  
            printf("解码失败：解码后的数据大小不匹配\n");  
            free(codes);  
            free(compressed_data);  
            free(decoded_data);  
            free_decoding_table(table);  
            return -1;  
        }  
    
        // 释放编码表内存  
        free(codes);  
    }  


    // 记录核心解码操作的结束时间  
    clock_t decode_end = clock();  
    double decode_time = (double)(decode_end - decode_start) / CLOCKS_PER_SEC;  

    // --- 对解码后的数据进行解密处理 ---  
    unsigned char key = 0x55;  // 与压缩时使用的密钥保持一致  
    clock_t decrypt_start = clock();  
    if (enc_method == ENCRYPT_METHOD_OFFSET) {  
        // 对偏移加密进行逆操作：减去 key  
        for (size_t i = 0; i < bytes_written; i++) {  
            decoded_data[i] = (unsigned char)((decoded_data[i] - key) & 0xFF);  
        }  
    } else if (enc_method == ENCRYPT_METHOD_XOR) {  
        // XOR 加密的逆操作和加密相同  
        for (size_t i = 0; i < bytes_written; i++) {  
            decoded_data[i] ^= key;  
        }  
    }  
    clock_t decrypt_end = clock();  
    double decrypt_time = (double)(decrypt_end - decrypt_start) / CLOCKS_PER_SEC;  

    // 将还原的原始数据写入输出文件  
    FILE* output_file = fopen(output_filename, "wb");  
    if (!output_file) {  
        printf("无法创建输出文件: %s\n", output_filename);  
        free(compressed_data);  
        free(decoded_data);  
        free_decoding_table(table);  
        return -1;  
    }  
    if (fwrite(decoded_data, 1, bytes_written, output_file) != bytes_written) {  
        printf("写入输出文件失败\n");  
        fclose(output_file);  
        free(compressed_data);  
        free(decoded_data);  
        free_decoding_table(table);  
        return -1;  
    }  
    fclose(output_file);  

    uint64_t hash = calculate_hash(decoded_data, bytes_written);  
    printf("解压完成!\n");  
    printf("输出文件: %s\n", output_filename);  
    printf("解压后大小: %zu 字节\n", bytes_written);  
    printf("文件哈希值: 0x%016llX\n", hash);  

    stats.end_time = clock();  
    stats.processed_bytes = bytes_written;  
    stats.compression_ratio = (double)compressed_size / table->original_size * 100;  
    print_performance_stats(stats, decode_time, bytes_written);  

    free(compressed_data);  
    free(decoded_data);  
    free_decoding_table(table);  
    return 0;  
}