// 单元测试：挂载/卸载/链接/取消链接 OPT 容器
// 需要管理员权限运行（驱动安装+挂载都需要）
//
// 用法:
//   test_mount_opt mount  <opt_path> <tag>           挂载容器，验证设备可访问
//   test_mount_opt unmount <tag>                     卸载容器，验证设备已消失
//   test_mount_opt link   <tag> <link_dir>           为已挂载的卷创建目录链接(Junction)
//   test_mount_opt unlink <link_dir>                 移除目录链接(Junction)
//
// 示例:
//   test_mount_opt mount  "C:\Users\aiyah\Desktop\SDGA_A011_20260716114938_0.opt" OPT_TEST
//   test_mount_opt link   OPT_TEST "C:\mnt\opt_test"
//   test_mount_opt unlink "C:\mnt\opt_test"
//   test_mount_opt unmount OPT_TEST
#define LOG_MODULE "TEST"
#include "fscrypt.h"
#include "driver_service.h"
#include "utils/strings.h"
#include "utils/log.h"
#include <windows.h>
#include <cstdio>
#include <string>
#include <format>

static void print_usage() {
    puts("Usage:");
    puts("  test_mount_opt mount   <opt_path> <tag>          Mount container");
    puts("  test_mount_opt unmount <tag>                      Unmount container");
    puts("  test_mount_opt link    <tag> <link_dir>           Link device to directory");
    puts("  test_mount_opt unlink  <link_dir>                 Remove directory link");
}

// 确保驱动服务运行
static bool ensure_driver() {
    std::string svc_err;
    if (!driver::ensure_running(&svc_err)) {
        LOG_ERROR("driver::ensure_running failed: {}", svc_err);
        LOG_ERROR("Hint: run this program as Administrator");
        return false;
    }
    LOG_INFO("Driver service ready");
    return true;
}

// 验证设备 \\.\FscryptDisk_<tag> 是否可访问
static bool check_device(const std::wstring &tag, bool should_exist) {
    std::wstring device_path = std::format(L"\\\\.\\FscryptDisk_{}", tag);
    HANDLE hCheck = CreateFileW(device_path.c_str(), 0,
                                FILE_SHARE_READ | FILE_SHARE_WRITE,
                                nullptr, OPEN_EXISTING, 0, nullptr);
    bool exists = (hCheck != INVALID_HANDLE_VALUE);
    if (exists) CloseHandle(hCheck);

    if (should_exist && !exists) {
        LOG_ERROR("Device {} not accessible: {}", utils::wstr_to_str(device_path), GetLastError());
        return false;
    }
    if (!should_exist && exists) {
        LOG_ERROR("Device {} still exists after unmount!", utils::wstr_to_str(device_path));
        return false;
    }
    if (should_exist) {
        LOG_INFO("Device {} accessible", utils::wstr_to_str(device_path));
    } else {
        LOG_INFO("Device {} removed", utils::wstr_to_str(device_path));
    }
    return true;
}

// ---- 子命令实现 ----

static int cmd_mount(int argc, char *argv[]) {
    if (argc < 4) { print_usage(); return 1; }
    std::wstring opt_path = utils::to_wstring(argv[2]);
    std::wstring tag = utils::to_wstring(argv[3]);

    if (!ensure_driver()) return 1;

    LOG_INFO("Mounting '{}' as tag '{}' ...", utils::wstr_to_str(opt_path), utils::wstr_to_str(tag));
    auto result = fscrypt::mount_container(opt_path, tag);
    if (!result.success) {
        LOG_ERROR("Mount failed: {}", result.error);
        return 1;
    }
    LOG_INFO("Mount succeeded");

    if (!check_device(tag, true)) return 1;

    LOG_INFO("=== MOUNT OK ===");
    return 0;
}

static int cmd_unmount(int argc, char *argv[]) {
    if (argc < 3) { print_usage(); return 1; }
    std::wstring tag = utils::to_wstring(argv[2]);

    if (!ensure_driver()) return 1;

    LOG_INFO("Unmounting tag '{}' ...", utils::wstr_to_str(tag));
    if (!fscrypt::unmount_container(tag)) {
        LOG_ERROR("Unmount failed");
        return 1;
    }
    LOG_INFO("Unmount succeeded");

    if (!check_device(tag, false)) return 1;

    LOG_INFO("=== UNMOUNT OK ===");
    return 0;
}

static int cmd_link(int argc, char *argv[]) {
    if (argc < 4) { print_usage(); return 1; }
    std::wstring tag = utils::to_wstring(argv[2]);
    std::wstring link_dir = utils::to_wstring(argv[3]);

    LOG_INFO("Linking tag '{}' -> '{}' ...", utils::wstr_to_str(tag), utils::wstr_to_str(link_dir));
    if (!fscrypt::link_container(tag, link_dir)) {
        LOG_ERROR("Link failed");
        return 1;
    }
    LOG_INFO("=== LINK OK ===");
    return 0;
}

static int cmd_unlink(int argc, char *argv[]) {
    if (argc < 3) { print_usage(); return 1; }
    std::wstring link_dir = utils::to_wstring(argv[2]);

    LOG_INFO("Unlinking '{}' ...", utils::wstr_to_str(link_dir));
    if (!fscrypt::unlink_container(link_dir)) {
        LOG_ERROR("Unlink failed");
        return 1;
    }
    LOG_INFO("=== UNLINK OK ===");
    return 0;
}

int main(int argc, char *argv[]) {
    SetConsoleOutputCP(CP_UTF8);
    SetConsoleCP(CP_UTF8);

    if (argc < 2) {
        print_usage();
        return 1;
    }

    std::string cmd = argv[1];
    if (cmd == "mount")   return cmd_mount(argc, argv);
    if (cmd == "unmount") return cmd_unmount(argc, argv);
    if (cmd == "link")    return cmd_link(argc, argv);
    if (cmd == "unlink")  return cmd_unlink(argc, argv);

    LOG_ERROR("Unknown command: {}", cmd);
    print_usage();
    return 1;
}