#include <stdio.h>  
#include <stdlib.h>  
#include <string.h>  
#include <stdint.h>  
#include "ui.h"
#include "compress.h"   // 包含 compress_file 函数及相关数据结构  
#include "decompress.h" // 包含 decompress_file 函数  

#ifdef _WIN32  
#include <windows.h>  
#endif  

// 设置控制台UTF-8编码  
static void set_console_utf8(void) {  
#ifdef _WIN32  
    // 设置控制台输出和输入代码页为 UTF-8  
    SetConsoleOutputCP(CP_UTF8);  
    SetConsoleCP(CP_UTF8);  

    HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);  
    if (hOut != INVALID_HANDLE_VALUE) {  
        DWORD mode = 0;  
        if (GetConsoleMode(hOut, &mode)) {  
            mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;  
            SetConsoleMode(hOut, mode);  
        }  
    }  
#endif  
}  

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

/**  
 * FNV-1a 64位哈希函数  
 *  
 * @param data   输入数据指针  
 * @param length 数据长度（字节数）  
 * @return 返回计算得到的64位哈希值  
 */  
static uint64_t fnv1a_64(const void *data, size_t length) {  
    uint64_t hash = FNV1A_64_INIT;  
    const unsigned char *byte_data = (const unsigned char*)data;  
    for (size_t i = 0; i < length; i++) {  
        hash ^= byte_data[i];  
        hash *= FNV1A_64_PRIME;  
    }  
    return hash;  
}  

// -------------------- 辅助数据结构与函数 ----------------------  

// 用于词频统计的基本结构体  
typedef struct {  
    unsigned char byte;  
    size_t frequency;  
} FreqItem;  

// Huffman树节点结构体  
typedef struct HuffmanNode {  
    unsigned char byte;  
    size_t frequency;  
    struct HuffmanNode* left;  
    struct HuffmanNode* right;  
} HuffmanNode;  

// 释放 Huffman树内存  
static void free_huffman_tree(HuffmanNode* root) {  
    if (root) {  
        free_huffman_tree(root->left);  
        free_huffman_tree(root->right);  
        free(root);  
    }  
}  

// 比较函数：按频率降序排序（用于词频列表显示）  
static int cmp_freq_desc(const void* a, const void* b) {  
    const FreqItem* item1 = (const FreqItem*)a;  
    const FreqItem* item2 = (const FreqItem*)b;  
    if (item2->frequency > item1->frequency)  
        return 1;  
    else if (item2->frequency < item1->frequency)  
        return -1;  
    return 0;  
}  

// 比较函数：按频率升序用于 Huffman节点排序  
static int cmp_huffnode_asc(const void* a, const void* b) {  
    HuffmanNode* const* pa = (HuffmanNode* const*)a;  
    HuffmanNode* const* pb = (HuffmanNode* const*)b;  
    if ((*pa)->frequency > (*pb)->frequency)  
        return 1;  
    else if ((*pa)->frequency < (*pb)->frequency)  
        return -1;  
    return 0;  
}  

// 构造 Huffman树，参数：频率列表items（长度为256），返回根节点  
static HuffmanNode* build_huffman_tree(FreqItem* items, size_t count) {  
    // 分配节点数组，筛选非0频率的项作为叶子节点  
    HuffmanNode** nodes = (HuffmanNode**)malloc(count * sizeof(HuffmanNode*));  
    if (!nodes) return NULL;  
    size_t nodeCount = 0;  
    for (size_t i = 0; i < 256; i++) {  
        if (items[i].frequency > 0) {  
            HuffmanNode* node = (HuffmanNode*)malloc(sizeof(HuffmanNode));  
            if (!node) {  
                free(nodes);  
                return NULL;  
            }  
            node->byte = items[i].byte;  
            node->frequency = items[i].frequency;  
            node->left = node->right = NULL;  
            nodes[nodeCount++] = node;  
        }  
    }  
    if (nodeCount == 0) {  
        free(nodes);  
        return NULL;  
    }  
    // 合并节点构造 Huffman树  
    while (nodeCount > 1) {  
        qsort(nodes, nodeCount, sizeof(HuffmanNode*), cmp_huffnode_asc);  
        HuffmanNode* left = nodes[0];  
        HuffmanNode* right = nodes[1];  

        HuffmanNode* parent = (HuffmanNode*)malloc(sizeof(HuffmanNode));  
        if (!parent) {  
            free_huffman_tree(left);  
            free_huffman_tree(right);  
            free(nodes);  
            return NULL;  
        }  
        parent->byte = 0; // 非叶子节点无实际字节  
        parent->frequency = left->frequency + right->frequency;  
        parent->left = left;  
        parent->right = right;  
        nodes[0] = parent;  
        nodes[1] = nodes[nodeCount - 1]; // 将最后一个节点移到索引1  
        nodeCount--;  
    }  
    HuffmanNode* root = nodes[0];  
    free(nodes);  
    return root;  
}  

// 递归计算 Huffman树的加权路径长度 (WPL)  
static size_t calc_wpl(HuffmanNode* root, size_t depth) {  
    if (!root)  
        return 0;  
    // 如果是叶子节点  
    if (!root->left && !root->right)  
        return root->frequency * depth;  
    return calc_wpl(root->left, depth + 1) + calc_wpl(root->right, depth + 1);  
}  

// 统计指定文件中各字节的出现频率，结果存入 items 数组（下标即字节值），返回文件总字节数  
static size_t count_file_frequencies(const char* filename, FreqItem items[256]) {  
    memset(items, 0, 256 * sizeof(FreqItem));  
    for (int i = 0; i < 256; i++) {  
        items[i].byte = (unsigned char)i;  
        items[i].frequency = 0;  
    }  
    
    FILE* fp = fopen(filename, "rb");  
    if (!fp) return 0;  
    size_t total = 0;  
    int ch;  
    while ((ch = fgetc(fp)) != EOF) {  
        items[(unsigned char)ch].frequency++;  
        total++;  
    }  
    fclose(fp);  
    return total;  
}  

// 显示词频统计列表：仅显示出现频率不为0的项，并按频率降序排列  
static void display_frequency_list(const FreqItem items[256]) {  
    FreqItem temp[256];  
    memcpy(temp, items, 256 * sizeof(FreqItem));  
    qsort(temp, 256, sizeof(FreqItem), cmp_freq_desc);  
    printf("字节值(HEX)   次数\n");  
    for (int i = 0; i < 256; i++) {  
        if (temp[i].frequency > 0)  
            printf("0x%02X\t\t%zu\n", temp[i].byte, temp[i].frequency);  
    }  
}  

// -------------------- 主函数 ----------------------  
int main() {  
    // 设置控制台UTF-8编码，并初始化UI  
    set_console_utf8();  
    ui_init();  

    // 建立持久化连接：循环处理多次请求，直到用户选择退出  
    while (1) {  
        ui_display_message("主菜单：压缩(c)、解压(d) 或 退出(q)");  
        char choice = ui_show_main_menu();  

        // 用户选择退出  
        if (choice == 'q' || choice == 'Q') {  
            ui_display_message("退出程序...");  
            break;  
        }  

        if (choice == 'c' || choice == 'C') {  
            // ---------- 压缩操作 ----------  
            char filename[256];  
            char sender_info[128];  
            char receiver_info[128];  

            ui_prompt_filename("\n请输入待压缩的文件名: ", filename, sizeof(filename));  
            ui_prompt_info("请输入发送人信息（格式：10位学号,姓名）：", sender_info, sizeof(sender_info));  
            ui_prompt_info("请输入接收人信息（格式：10位学号,姓名）：", receiver_info, sizeof(receiver_info));  

            int method_choice = ui_prompt_encryption_method();  
            EncryptMethod method = (method_choice == 1) ? ENCRYPT_METHOD_XOR : ENCRYPT_METHOD_OFFSET;  

            // 调用压缩接口（内部会按原有逻辑输出提示信息）  
            if (compress_file(filename, sender_info, receiver_info, method) != 0) {  
                ui_display_message("文件压缩失败！");  
                ui_pause();  
                continue;  
            }  
            ui_display_message("文件压缩成功！");  

            // ---------- 1. 输出原文件HASH值 ----------  
            {  
                FILE *fp = fopen(filename, "rb");  
                if (!fp) {  
                    fprintf(stderr, "无法打开原文件: %s\n", filename);  
                } else {  
                    fseek(fp, 0, SEEK_END);  
                    long fileSize = ftell(fp);  
                    rewind(fp);  
                    unsigned char *buffer = (unsigned char*)malloc(fileSize);  
                    if (!buffer) {  
                        fprintf(stderr, "内存分配失败\n");  
                    } else {  
                        if (fread(buffer, 1, fileSize, fp) != (size_t)fileSize) {  
                            fprintf(stderr, "读取原文件数据失败\n");  
                        } else {  
                            uint64_t origHash = fnv1a_64(buffer, fileSize);  
                            printf("\n原文件 HASH 值: 0x%016llx\n", origHash);  
                        }  
                        free(buffer);  
                    }  
                    fclose(fp);  
                }  
            }  

            // 构造生成文件名  
            char compressed_filename[256];  
            char code_filename[256];  
            // 假设压缩后文件扩展名为 ".hfm"，编码表扩展名为 ".hfm.code"  
            snprintf(compressed_filename, sizeof(compressed_filename), "%s.hfm", filename);  
            snprintf(code_filename, sizeof(code_filename), "%s.hfm.code", filename);  

            // ---------- 2. 显示编码表内容 ----------  
            printf("\n=== 编码表内容 (HEX 显示) ===\n");  
            FILE* fpCode = fopen(code_filename, "r");  
            if (fpCode) {  
                char line[1024];  
                while (fgets(line, sizeof(line), fpCode)) {  
                    printf("%s", line);  
                }  
                fclose(fpCode);  
            } else {  
                printf("无法打开编码表文件: %s\n", code_filename);  
            }  

            // ---------- 3. 显示压缩文件最后16个字节及压缩数据HASH值 ----------  
            printf("\n=== 压缩文件最后16个字节 (HEX 显示) ===\n");  
            FILE* fpComp = fopen(compressed_filename, "rb");  
            if (fpComp) {  
                fseek(fpComp, 0, SEEK_END);  
                long compSize = ftell(fpComp);  
                if (compSize < 0) {  
                    printf("获取压缩文件大小失败\n");  
                } else {  
                    unsigned char* compData = (unsigned char*)malloc(compSize);  
                    if (!compData) {  
                        printf("内存分配失败\n");  
                    } else {  
                        fseek(fpComp, 0, SEEK_SET);  
                        if (fread(compData, 1, compSize, fpComp) != (size_t)compSize) {  
                            printf("读取压缩文件数据失败\n");  
                        } else {  
                            size_t start = (compSize >= 16) ? compSize - 16 : 0;  
                            printf("总字节数: %ld\n", compSize);  
                            printf("最后16个字节: ");  
                            for (size_t i = start; i < (size_t)compSize; i++) {  
                                printf("%02X ", compData[i]);  
                            }  
                            printf("\n");  
                            uint64_t compHash = fnv1a_64(compData, compSize);  
                            printf("压缩数据 HASH 值: 0x%016llx\n", compHash);  
                        }  
                        free(compData);  
                    }  
                }  
                fclose(fpComp);  
            } else {  
                printf("无法打开压缩文件: %s\n", compressed_filename);  
            }  

            // ---------- 4. 显示词频统计列表 ----------  
            printf("\n=== 词频统计列表 ===\n");  
            FreqItem freqItems[256];  
            size_t totalBytes = count_file_frequencies(filename, freqItems);  
            if (totalBytes == 0) {  
                printf("读取原始文件失败或文件为空。\n");  
            } else {  
                printf("原始文件总字节数: %zu\n", totalBytes);  
                display_frequency_list(freqItems);  
            }  

            // ---------- 5. 计算并显示 Huffman 树的 WPL 值 ----------  
            printf("\n=== Huffman 树的加权路径长度 (WPL) ===\n");  
            HuffmanNode* root = build_huffman_tree(freqItems, 256);  
            if (root) {  
                size_t wpl = calc_wpl(root, 0);  
                printf("Huffman 树 WPL 值: %zu\n", wpl);  
                free_huffman_tree(root);  
            } else {  
                printf("构造 Huffman 树失败。\n");  
            }  
        } else if (choice == 'd' || choice == 'D') {  
            // ---------- 解压操作 ----------  
            char filename[256];  
            char sender_info[128];  
            char receiver_info[128];  

            ui_prompt_filename("\n请输入待解压的压缩文件名 (*.hfm): ", filename, sizeof(filename));  
            ui_prompt_info("请输入发送人信息（格式：10位学号,姓名）：", sender_info, sizeof(sender_info));  
            ui_prompt_info("请输入接收人信息（格式：10位学号,姓名）：", receiver_info, sizeof(receiver_info));  

            if (decompress_file(filename, sender_info, receiver_info) == 0) {  
                ui_display_message("文件解压成功！");  
            } else {  
                ui_display_message("文件解压失败！");  
            }  
        } else {  
            ui_display_message("无效的选择，请输入 'c'、'd' 或 'q'。");  
        }  

        ui_pause();  // 操作结束后等待用户按键  
    }  

    ui_cleanup();  
    return EXIT_SUCCESS;  
}
