#include "vhd_mount.h"
#include "utils/strings.h"
#include <initguid.h>
#include <virtdisk.h>
#include <string>
#include <vector>
#include <format>

namespace vhd_mount {
    // ===== 内部 helper =====

    // 打开VHD句柄，失败返回INVALID_HANDLE_VALUE并填充error
    static HANDLE open_vhd(const std::wstring &vhd_path, VIRTUAL_DISK_ACCESS_MASK access_mask, DWORD rw_depth,
                           std::string *error) {
        VIRTUAL_STORAGE_TYPE vst = {};
        vst.DeviceId = VIRTUAL_STORAGE_TYPE_DEVICE_VHD;
        vst.VendorId = VIRTUAL_STORAGE_TYPE_VENDOR_MICROSOFT;

        OPEN_VIRTUAL_DISK_PARAMETERS open_params = {};
        open_params.Version = OPEN_VIRTUAL_DISK_VERSION_1;
        open_params.Version1.RWDepth = rw_depth;

        HANDLE hVhd = nullptr;
        DWORD err = OpenVirtualDisk(&vst, vhd_path.c_str(),
                                    access_mask,
                                    OPEN_VIRTUAL_DISK_FLAG_NONE,
                                    &open_params, &hVhd);
        if (err != ERROR_SUCCESS) {
            if (error)
                *error = std::format("OpenVirtualDisk 失败：{} path={}", err, utils::wstr_to_str(vhd_path));
            return INVALID_HANDLE_VALUE;
        }
        return hVhd;
    }

    // 通过任意卷设备路径（\\.\X: 或 \\?\Volume{GUID}）获取卷所在磁盘号
    static bool device_get_disk_number(const std::wstring &device_path, DWORD &disk_number) {
        HANDLE hVol = CreateFileW(device_path.c_str(),
                                  GENERIC_READ,
                                  FILE_SHARE_READ | FILE_SHARE_WRITE,
                                  nullptr,
                                  OPEN_EXISTING,
                                  FILE_ATTRIBUTE_NORMAL,
                                  nullptr);
        if (hVol == INVALID_HANDLE_VALUE) return false;

        struct VolExtentsBuf {
            DWORD NumberOfDiskExtents;
            DISK_EXTENT Extents[256];
        } buf{};
        std::memset(&buf, 0, sizeof(buf));
        DWORD bytes_returned = 0;
        BOOL ok = DeviceIoControl(hVol, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
                                  nullptr, 0,
                                  &buf, sizeof(buf),
                                  &bytes_returned, nullptr);
        CloseHandle(hVol);
        if (!ok || buf.NumberOfDiskExtents == 0) return false;
        disk_number = buf.Extents[0].DiskNumber;
        return true;
    }

    // 获取指定物理磁盘号上的所有卷GUID路径（封装FindFirstVolumeW，不重复扫描）
    // 返回的路径形如 \\?\Volume{GUID}\（带尾部反斜杠）
    static bool get_volumes_on_disk(DWORD disk_number,
                                    std::vector<std::wstring> &out_volumes) {
        out_volumes.clear();
        wchar_t vol_name[MAX_PATH];
        HANDLE hFind = FindFirstVolumeW(vol_name, _countof(vol_name));
        if (hFind == INVALID_HANDLE_VALUE) {
            return false;
        }
        do {
            std::wstring guid_path(vol_name);
            // FindFirstVolumeW返回\\?\前缀，CreateFile需要\\.\前缀才能打开卷设备
            std::wstring dev_path = guid_path;
            if (dev_path.size() >= 4 && dev_path.substr(0, 4) == L"\\\\?\\") {
                dev_path = std::format(L"\\\\.\\{}", dev_path.substr(4));
            }
            if (!dev_path.empty() && dev_path.back() == L'\\') dev_path.pop_back();

            DWORD disk_num = 0;
            if (!device_get_disk_number(dev_path, disk_num)) continue;
            if (disk_num != disk_number) continue;
            out_volumes.push_back(guid_path);
        } while (FindNextVolumeW(hFind, vol_name, _countof(vol_name)));
        FindVolumeClose(hFind);
        return true;
    }

    // ===== 公开接口 =====

    // 可写挂载VHD，返回磁盘号、物理路径、卷GUID（不分配盘符）
    // 流程：Open(ATTACH_RW|GET_INFO) → Attach(NO_DRIVE_LETTER) → GetVirtualDiskPhysicalPath → 查找卷GUID
    bool attach_vhd(const std::wstring &vhd_path, MountedVhdInfo &info, std::string *error) {
        // 可写方式打开overlay层（父盘仍是只读的，overlay本身本地可写）
        HANDLE hVhd = open_vhd(vhd_path,
                               VIRTUAL_DISK_ACCESS_ATTACH_RW | VIRTUAL_DISK_ACCESS_GET_INFO,
                               1, // 顶层1层可写
                               error);
        if (hVhd == INVALID_HANDLE_VALUE) return false;

        // 挂载VHD（可写，不自动分配盘符）
        ATTACH_VIRTUAL_DISK_PARAMETERS attach_params = {};
        attach_params.Version = ATTACH_VIRTUAL_DISK_VERSION_1;
        attach_params.Version1.Reserved = 0;
        DWORD err = AttachVirtualDisk(hVhd, nullptr,
                                      ATTACH_VIRTUAL_DISK_FLAG_PERMANENT_LIFETIME |
                                      ATTACH_VIRTUAL_DISK_FLAG_NO_DRIVE_LETTER,
                                      0, &attach_params, nullptr);
        if (err != ERROR_SUCCESS) {
            if (error) *error = std::format("AttachVirtualDisk（RW）失败：{}", err);
            CloseHandle(hVhd);
            return false;
        }

        // 通过VHD句柄直接获取物理磁盘路径（\\.\PhysicalDriveX），重试等待系统定位
        // 等待：20次 × 100ms = 最长2秒
        wchar_t phys_path[MAX_PATH] = {};
        ULONG path_size = sizeof(phys_path);
        err = ERROR_NOT_READY;
        for (int retry = 0; retry < 20; retry++) {
            path_size = sizeof(phys_path);
            err = GetVirtualDiskPhysicalPath(hVhd, &path_size, phys_path);
            if (err == ERROR_SUCCESS) break;
            Sleep(100);
        }
        CloseHandle(hVhd);

        if (err != ERROR_SUCCESS) {
            if (error) *error = std::format("GetVirtualDiskPhysicalPath 失败：{}", err);
            return false;
        }

        info.physical_path = phys_path;

        // 解析 "\\.\PhysicalDriveX" 中的磁盘号
        DWORD disk_num = 0;
        if (swscanf_s(phys_path, L"\\\\.\\PhysicalDrive%u", &disk_num) != 1) {
            if (error)
                *error = std::format("解析物理磁盘路径失败：{}", utils::wstr_to_str(phys_path));
            return false;
        }
        info.disk_number = disk_num;

        // 在该磁盘上查找卷GUID路径（重试等待卷枚举就绪）
        // 等待：20次 × 100ms = 最长2秒
        for (int retry = 0; retry < 20; retry++) {
            std::vector<std::wstring> vols;
            get_volumes_on_disk(disk_num, vols);
            if (!vols.empty()) {
                info.volume_guid = vols.front();
                return true;
            }
            Sleep(100);
        }

        if (error) *error = std::format("磁盘 {} 上未找到卷", disk_num);
        return false;
    }

    bool detach_vhd(const std::wstring &vhd_path, std::string *error) {
        HANDLE hVhd = open_vhd(vhd_path, VIRTUAL_DISK_ACCESS_DETACH | VIRTUAL_DISK_ACCESS_GET_INFO, 0, error);
        if (hVhd == INVALID_HANDLE_VALUE) return false;

        DWORD err = DetachVirtualDisk(hVhd, DETACH_VIRTUAL_DISK_FLAG_NONE, 0);
        CloseHandle(hVhd);

        if (err != ERROR_SUCCESS) {
            if (error) *error = std::format("DetachVirtualDisk 失败：{}", err);
            return false;
        }
        return true;
    }

    bool create_diff_vhd(const std::wstring &vhd_path, const std::wstring &parent_vhd_path, std::string *error) {
        // 如果上次运行遗留了已挂载的overlay.vhd，先尝试卸载再删除
        std::string detach_err;
        (void) detach_vhd(vhd_path, &detach_err);
        DeleteFileW(vhd_path.c_str());

        VIRTUAL_STORAGE_TYPE vst = {};
        vst.DeviceId = VIRTUAL_STORAGE_TYPE_DEVICE_VHD;
        vst.VendorId = VIRTUAL_STORAGE_TYPE_VENDOR_MICROSOFT;

        CREATE_VIRTUAL_DISK_PARAMETERS params = {};
        params.Version = CREATE_VIRTUAL_DISK_VERSION_2;
        params.Version2.UniqueId = GUID_NULL;
        params.Version2.MaximumSize = 0; // 差分盘跟随父盘大小
        params.Version2.BlockSizeInBytes = 0;
        params.Version2.SectorSizeInBytes = 512; // VHD1必须为512
        params.Version2.ParentPath = parent_vhd_path.c_str();
        params.Version2.SourcePath = nullptr;
        params.Version2.OpenFlags = OPEN_VIRTUAL_DISK_FLAG_NONE;
        params.Version2.ResiliencyGuid = GUID_NULL;

        // Version2强制要求VirtualDiskAccessMask必须为VIRTUAL_DISK_ACCESS_NONE
        HANDLE hVhd = nullptr;
        DWORD err = CreateVirtualDisk(&vst,
                                      vhd_path.c_str(),
                                      VIRTUAL_DISK_ACCESS_NONE,
                                      nullptr,
                                      CREATE_VIRTUAL_DISK_FLAG_NONE,
                                      0,
                                      &params,
                                      nullptr,
                                      &hVhd);
        if (err != ERROR_SUCCESS) {
            if (error)
                *error = std::format("CreateVirtualDisk 失败：{} path={} parent={}", err,
                                     utils::wstr_to_str(vhd_path), utils::wstr_to_str(parent_vhd_path));
            return false;
        }
        CloseHandle(hVhd);
        return true;
    }
} // namespace vhd_mount