#define LOG_MODULE "ICF"
#include "icf_mount.h"
#include "icf_parser.h"
#include "fscrypt.h"
#include "vhd_mount.h"
#include "utils/files.h"
#include "utils/strings.h"
#include "utils/log.h"
#include <map>
#include <format>
#include <ranges>
#include <windows.h>

namespace icf_mount {
    // 前向声明
    static std::vector<const icf::DecodedEntry *> get_active_apps(const icf::IcfFile &icf_file);

    static std::vector<const icf::DecodedEntry *> get_active_opts(const icf::IcfFile &icf_file);

    // 解析ICF文件并挂载所有Active状态的APP和OPT容器
    // 流程：挂载APP → 创建/挂载overlay → 挂载OPT → 创建LINK
    IcfMountResult mount_icf_apps(const Config &cfg) {
        IcfMountResult result{};
        result.success = false;

        // 读取ICF文件
        auto icf_data = utils::read_file(cfg.icf_path);
        if (icf_data.empty()) {
            result.error = "Failed to read ICF file";
            return result;
        }

        // 解密并解析ICF（parse_file内部会处理解密）
        std::string parse_err;
        auto icf_file = icf::parse_file(icf_data, &parse_err);
        if (!icf_file) {
            result.error = std::format("Parse ICF failed: {}", parse_err);
            return result;
        }

        // 筛选Active状态的APP条目，按part_index排序
        auto app_entries = get_active_apps(*icf_file);
        if (app_entries.empty()) {
            result.error = "No active APP entries found in ICF";
            return result;
        }

        // 验证base version依赖关系
        // part 0 没有base version
        // part n (n>0) 的 base_version 应该等于 part n-1 的 version
        for (size_t i = 1; i < app_entries.size(); i++) {
            auto *cur = app_entries[i];
            auto *prev = app_entries[i - 1];
            if (cur->part_index != prev->part_index + 1) {
                result.error = std::format("Part index not consecutive: {} -> {}", prev->part_index, cur->part_index);
                return result;
            }
            if (cur->base_version != prev->version) {
                result.error = std::format("Base version mismatch for part {}: {} != {}", cur->part_index,
                                           cur->base_version, prev->version);
                return result;
            }
        }

        // APP/OPT文件所在目录：优先使用image_dir，否则使用exe所在目录
        std::wstring dir = cfg.image_dir.empty() ? utils::get_exe_dir() : cfg.image_dir;

        // 构造密钥override（config.ini [KEYS] section，空则用内置密钥表）
        fscrypt::KeyOverride key_override;
        key_override.app_key_hex = cfg.app_key;
        key_override.app_iv_hex = cfg.app_iv;
        key_override.opt_key_hex = cfg.opt_key;
        key_override.opt_iv_hex = cfg.opt_iv;

        // 依次挂载每个APP容器
        bool all_ok = true;
        for (auto entry: app_entries) {
            MountEntry mount_entry{};
            mount_entry.filename = entry->filename;
            mount_entry.version = entry->version;
            mount_entry.part_index = entry->part_index;
            mount_entry.tag = std::format("APP_{}", entry->part_index);
            mount_entry.vhd_success = false;

            std::wstring app_path = std::format(L"{}\\{}", dir, utils::to_wstring(entry->filename));
            std::wstring tag = std::format(L"APP_{}", entry->part_index);

            LOG_INFO("Mount part {}: {}", entry->part_index, entry->filename);

            auto mount_res = fscrypt::mount_container(app_path, tag, &key_override);
            mount_entry.success = mount_res.success;
            mount_entry.error = mount_res.error;

            if (!mount_res.success) {
                LOG_ERROR("Mount part {} failed: {}", entry->part_index, mount_res.error);
                all_ok = false;
            }

            result.apps.push_back(mount_entry);
        }

        // 所有APP容器挂载完成后，处理overlay写层：
        // overlay.vhd不存在 → 创建差分盘 + 拷贝overlay目录文件
        // overlay.vhd已存在 → 直接挂载（数据已在里面，无需重新拷贝）
        if (all_ok && !app_entries.empty()) {
            auto *last = app_entries.back();
            std::wstring exe_dir = utils::get_exe_dir();
            std::wstring overlay_vhd = std::format(L"{}\\overlay.vhd", exe_dir);
            std::wstring last_tag = std::format(L"APP_{}", last->part_index);
            std::wstring parent_for_create = std::format(L"\\\\?\\FscryptDisk_{}\\internal_{}.vhd", last_tag,
                                                         last->part_index);

            bool vhd_exists = GetFileAttributesW(overlay_vhd.c_str()) != INVALID_FILE_ATTRIBUTES;

            if (!vhd_exists) {
                // 创建overlay.vhd差分盘
                LOG_INFO("Creating overlay.vhd (parent: internal_{}.vhd)", last->part_index);
                std::string ov_err;
                if (!vhd_mount::create_diff_vhd(overlay_vhd, parent_for_create, &ov_err)) {
                    LOG_ERROR("Overlay create failed: {}", ov_err);
                    all_ok = false;
                }
            } else {
                LOG_INFO("Reusing existing overlay.vhd");
                // 先尝试detach，处理上次未正常卸载的情况
                vhd_mount::detach_vhd(overlay_vhd, nullptr);
            }

            if (all_ok) {
                // 挂载overlay.vhd
                vhd_mount::MountedVhdInfo vhd_info{};
                std::string attach_err;
                if (!vhd_mount::attach_vhd(overlay_vhd, vhd_info, &attach_err)) {
                    LOG_ERROR("Overlay attach failed: {}", attach_err);
                    all_ok = false;
                } else {
                    result.apps.back().vhd_success = true;

                    // 分配盘符到overlay卷
                    char drive = static_cast<char>(cfg.mount_letter);
                    std::wstring mount_point = std::format(L"{}:\\", cfg.mount_letter);
                    DeleteVolumeMountPointW(mount_point.c_str()); // 清理可能残留的旧挂载点
                    if (!SetVolumeMountPointW(mount_point.c_str(), vhd_info.volume_guid.c_str())) {
                        LOG_ERROR("Assign drive {} failed: Error = {}", drive, GetLastError());
                        all_ok = false;
                    } else {
                        LOG_INFO("Overlay mounted to {}:", drive);

                        // 仅新建overlay.vhd时才拷贝overlay目录文件
                        // OVERLAY_DIR为空或不存在时跳过拷贝
                        if (!vhd_exists) {
                            if (!cfg.overlay_dir.empty()) {
                                DWORD attr = GetFileAttributesW(cfg.overlay_dir.c_str());
                                bool dir_exists = attr != INVALID_FILE_ATTRIBUTES &&
                                                  attr & FILE_ATTRIBUTE_DIRECTORY;
                                if (dir_exists) {
                                    std::string cp_err;
                                    size_t copied = 0;
                                    if (!utils::copy_dir_overwrite(cfg.overlay_dir, mount_point,
                                                                   &cp_err, &copied)) {
                                        LOG_ERROR("Overlay copy failed: {}", cp_err);
                                        all_ok = false;
                                    } else {
                                        LOG_INFO("Copied {} files", copied);
                                    }
                                } else {
                                    LOG_INFO("OVERLAY_DIR not found, skip copy");
                                }
                            } else {
                                LOG_INFO("OVERLAY_DIR empty, skip copy");
                            }
                        }
                    }
                }
            }
        }

        // 挂载OPT容器（独立于APP/overlay，不受其结果影响）
        auto opt_entries = get_active_opts(*icf_file);
        for (auto entry: opt_entries) {
            MountEntry mount_entry{};
            mount_entry.filename = entry->filename;
            mount_entry.version = entry->version;
            mount_entry.part_index = entry->part_index;
            mount_entry.tag = std::format("OPT_{}", entry->version);
            mount_entry.vhd_success = false;

            std::wstring opt_path = std::format(L"{}\\{}", dir, utils::to_wstring(entry->filename));
            std::wstring opt_tag = std::format(L"OPT_{}", utils::to_wstring(entry->version));

            LOG_INFO("Mount OPT {}: {}", entry->version, entry->filename);

            auto mount_res = fscrypt::mount_container(opt_path, opt_tag, &key_override);
            mount_entry.success = mount_res.success;
            mount_entry.error = mount_res.error;

            if (!mount_res.success) {
                LOG_ERROR("Mount OPT {} failed: {}", entry->version, mount_res.error);
            }

            result.opts.push_back(mount_entry);
        }

        // 创建APP LINK（仅对挂载成功的APP）
        if (!cfg.app_link.empty()) {
            for (const auto &app: result.apps) {
                if (!app.success) continue;
                std::wstring tag = std::format(L"APP_{}", app.part_index);
                std::wstring link_path = std::format(L"{}\\{}", cfg.app_link, app.part_index);
                if (!fscrypt::link_container(tag, link_path)) {
                    LOG_ERROR("Link APP {} failed", app.part_index);
                }
            }
        }

        // 创建OPT LINK（仅对挂载成功的OPT）
        if (!cfg.opt_link.empty()) {
            for (const auto &opt: result.opts) {
                if (!opt.success) continue;
                std::wstring tag = std::format(L"OPT_{}", utils::to_wstring(opt.version));
                std::wstring link_path = std::format(L"{}\\{}", cfg.opt_link, utils::to_wstring(opt.version));
                if (!fscrypt::link_container(tag, link_path)) {
                    LOG_ERROR("Link OPT {} failed", opt.version);
                }
            }
        }

        result.success = all_ok;
        if (!all_ok) {
            result.error = "Some APPs failed to mount";
        }
        return result;
    }

    // 从ICF中获取所有Active状态的APP条目，同一part_index只保留timestamp最新的，按part_index升序
    static std::vector<const icf::DecodedEntry *> get_active_apps(const icf::IcfFile &icf_file) {
        // 按 part_index 分组，每组保留 timestamp 最大的条目（timestamp格式为 YYYYMMDDhhmmss，定长数字串可直接比较）
        std::map<uint8_t, const icf::DecodedEntry *> latest_by_part;
        for (const auto &entry: icf_file.entries) {
            if (entry.type == icf::EntryType::App &&
                static_cast<uint16_t>(entry.status) == 0x0102) {
                auto it = latest_by_part.find(entry.part_index);
                if (it == latest_by_part.end() || entry.timestamp > it->second->timestamp) {
                    latest_by_part[entry.part_index] = &entry;
                }
            }
        }
        // map 按 key 升序，直接收集 values 即可
        std::vector<const icf::DecodedEntry *> app_entries;
        app_entries.reserve(latest_by_part.size());
        for (const auto *entry: latest_by_part | std::views::values) {
            app_entries.push_back(entry);
        }
        return app_entries;
    }

    // 从ICF中获取所有Active状态的OPT条目，同一version只保留timestamp最新的
    static std::vector<const icf::DecodedEntry *> get_active_opts(const icf::IcfFile &icf_file) {
        // OPT的version字段为4字符名称（如A001），按此分组去重
        std::map<std::string, const icf::DecodedEntry *> latest_by_name;
        for (const auto &entry: icf_file.entries) {
            if (entry.type == icf::EntryType::Opt &&
                static_cast<uint16_t>(entry.status) == 0x0102) {
                auto it = latest_by_name.find(entry.version);
                if (it == latest_by_name.end() || entry.timestamp > it->second->timestamp) {
                    latest_by_name[entry.version] = &entry;
                }
            }
        }
        std::vector<const icf::DecodedEntry *> opt_entries;
        opt_entries.reserve(latest_by_name.size());
        for (const auto *entry: latest_by_name | std::views::values) {
            opt_entries.push_back(entry);
        }
        return opt_entries;
    }

    // 解析ICF文件并卸载所有容器
    // 流程：移除LINK → 卸载overlay → 倒序卸载APP → 卸载OPT
    IcfMountResult unmount_icf_apps(const Config &cfg) {
        IcfMountResult result{};
        result.success = false;

        // 读取并解析ICF
        auto icf_data = utils::read_file(cfg.icf_path);
        if (icf_data.empty()) {
            result.error = "Failed to read ICF file";
            return result;
        }

        std::string parse_err;
        auto icf_file = icf::parse_file(icf_data, &parse_err);
        if (!icf_file) {
            result.error = std::format("Parse ICF failed: {}", parse_err);
            return result;
        }

        auto app_entries = get_active_apps(*icf_file);
        auto opt_entries = get_active_opts(*icf_file);

        bool all_ok = true;

        // 1) 移除APP LINK
        if (!cfg.app_link.empty()) {
            for (auto *entry: app_entries) {
                std::wstring link_path = std::format(L"{}\\{}", cfg.app_link, entry->part_index);
                fscrypt::unlink_container(link_path);
            }
        }

        // 2) 移除OPT LINK
        if (!cfg.opt_link.empty()) {
            for (auto *entry: opt_entries) {
                std::wstring link_path = std::format(L"{}\\{}", cfg.opt_link,
                                                     utils::to_wstring(entry->version));
                fscrypt::unlink_container(link_path);
            }
        }

        // 3) overlay.vhd：移除盘符 + detach（保留文件供下次沿用）
        std::wstring exe_dir = utils::get_exe_dir();
        std::wstring overlay_vhd = std::format(L"{}\\overlay.vhd", exe_dir);
        {
            char drive = static_cast<char>(cfg.mount_letter);
            std::wstring mount_point = std::format(L"{}:\\", cfg.mount_letter);
            if (DeleteVolumeMountPointW(mount_point.c_str())) {
                LOG_INFO("Removed drive letter {}", drive);
            }

            WIN32_FILE_ATTRIBUTE_DATA fa;
            if (GetFileAttributesExW(overlay_vhd.c_str(), GetFileExInfoStandard, &fa)) {
                LOG_INFO("Detaching overlay.vhd...");
                std::string vhd_err;
                if (!vhd_mount::detach_vhd(overlay_vhd, &vhd_err)) {
                    LOG_ERROR("Overlay detach failed: {}", vhd_err);
                    all_ok = false;
                } else {
                    LOG_INFO("Detached overlay.vhd");
                }
            } else {
                LOG_INFO("overlay.vhd not found, skip detach");
            }
        }

        // 4) 倒序卸载APP容器
        for (auto it = app_entries.rbegin(); it != app_entries.rend(); ++it) {
            auto *entry = *it;
            MountEntry mount_entry{};
            mount_entry.filename = entry->filename;
            mount_entry.version = entry->version;
            mount_entry.part_index = entry->part_index;
            mount_entry.tag = std::format("APP_{}", entry->part_index);
            mount_entry.vhd_success = false;

            std::wstring tag = std::format(L"APP_{}", entry->part_index);

            LOG_INFO("Unmount part {}: {}", entry->part_index, entry->filename);

            bool ok = fscrypt::unmount_container(tag);
            mount_entry.success = ok;
            if (!ok) {
                mount_entry.error = "Unmount failed";
                LOG_ERROR("Unmount part {} failed", entry->part_index);
                all_ok = false;
            }

            result.apps.push_back(mount_entry);
        }

        // 5) 卸载OPT容器
        for (auto *entry: opt_entries) {
            MountEntry mount_entry{};
            mount_entry.filename = entry->filename;
            mount_entry.version = entry->version;
            mount_entry.part_index = entry->part_index;
            mount_entry.tag = std::format("OPT_{}", entry->version);
            mount_entry.vhd_success = false;

            std::wstring opt_tag = std::format(L"OPT_{}", utils::to_wstring(entry->version));

            LOG_INFO("Unmount OPT {}: {}", entry->version, entry->filename);

            bool ok = fscrypt::unmount_container(opt_tag);
            mount_entry.success = ok;
            if (!ok) {
                mount_entry.error = "Unmount failed";
                LOG_ERROR("Unmount OPT {} failed", entry->version);
                all_ok = false;
            }

            result.opts.push_back(mount_entry);
        }

        result.success = all_ok;
        if (!all_ok) {
            result.error = "Some containers failed to unmount";
        }
        return result;
    }
}