#include "disks.h"
#include "util.h"

#include <winioctl.h>

#include <algorithm>
#include <cctype>
#include <cwctype>
#include <utility>

namespace {

HANDLE OpenPhysicalDrive(int disk)
{
    std::wstring path = L"\\\\.\\PhysicalDrive" + std::to_wstring(disk);
    HANDLE handle = CreateFileW(path.c_str(), GENERIC_READ,
                                FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                OPEN_EXISTING, 0, nullptr);
    if (handle == INVALID_HANDLE_VALUE) {
        handle = CreateFileW(path.c_str(), 0,
                             FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                             OPEN_EXISTING, 0, nullptr);
    }
    return handle;
}

HANDLE OpenVolume(const wchar_t letter)
{
    std::wstring path = L"\\\\.\\";
    path.push_back(letter);
    path.push_back(L':');
    HANDLE handle = CreateFileW(path.c_str(), 0,
                                FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                                OPEN_EXISTING, 0, nullptr);
    if (handle == INVALID_HANDLE_VALUE) {
        handle = CreateFileW(path.c_str(), GENERIC_READ,
                             FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
                             OPEN_EXISTING, 0, nullptr);
    }
    return handle;
}

std::wstring TrimAscii(std::string value)
{
    size_t begin = 0;
    while (begin < value.size() && (value[begin] == ' ' || value[begin] == '\t' ||
                                    value[begin] == '\r' || value[begin] == '\n')) {
        ++begin;
    }
    size_t end = value.size();
    while (end > begin && (value[end - 1] == ' ' || value[end - 1] == '\t' ||
                           value[end - 1] == '\r' || value[end - 1] == '\n')) {
        --end;
    }
    return AnsiToWide(value.substr(begin, end - begin), CP_ACP);
}

std::wstring TrimWide(std::wstring value)
{
    size_t begin = 0;
    while (begin < value.size() && iswspace(value[begin])) {
        ++begin;
    }
    size_t end = value.size();
    while (end > begin && iswspace(value[end - 1])) {
        --end;
    }
    return value.substr(begin, end - begin);
}

std::wstring DescriptorString(const std::vector<BYTE>& buffer,
                              const STORAGE_DEVICE_DESCRIPTOR& descriptor,
                              DWORD offset)
{
    if (offset == 0 || offset >= descriptor.Size || offset >= buffer.size()) {
        return std::wstring();
    }
    size_t available = std::min<size_t>(descriptor.Size, buffer.size()) - offset;
    const char* start = reinterpret_cast<const char*>(buffer.data() + offset);
    size_t length = 0;
    while (length < available && start[length] != '\0') {
        ++length;
    }
    return TrimAscii(std::string(start, length));
}

struct StorageDescriptor {
    bool valid = false;
    int busType = -1;
    bool removable = false;
    std::wstring vendor;
    std::wstring product;
};

StorageDescriptor QueryStorageDescriptor(HANDLE handle)
{
    StorageDescriptor result;
    STORAGE_PROPERTY_QUERY query{};
    query.PropertyId = StorageDeviceProperty;
    query.QueryType = PropertyStandardQuery;

    std::vector<BYTE> buffer(4096);
    DWORD returned = 0;
    if (!DeviceIoControl(handle, IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(query),
                         buffer.data(), static_cast<DWORD>(buffer.size()), &returned, nullptr)) {
        return result;
    }
    if (returned < sizeof(STORAGE_DEVICE_DESCRIPTOR)) {
        return result;
    }

    auto* descriptor = reinterpret_cast<STORAGE_DEVICE_DESCRIPTOR*>(buffer.data());
    if (descriptor->Size < sizeof(STORAGE_DEVICE_DESCRIPTOR) || descriptor->Size > buffer.size()) {
        return result;
    }

    result.valid = true;
    result.busType = static_cast<int>(descriptor->BusType);
    result.removable = descriptor->RemovableMedia != 0;
    result.vendor = DescriptorString(buffer, *descriptor, descriptor->VendorIdOffset);
    result.product = DescriptorString(buffer, *descriptor, descriptor->ProductIdOffset);
    return result;
}

bool QueryDiskSize(HANDLE handle, unsigned long long& size)
{
    GET_LENGTH_INFORMATION length{};
    DWORD returned = 0;
    if (DeviceIoControl(handle, IOCTL_DISK_GET_LENGTH_INFO, nullptr, 0,
                        &length, sizeof(length), &returned, nullptr)) {
        size = static_cast<unsigned long long>(length.Length.QuadPart);
        return true;
    }

    std::vector<BYTE> buffer(sizeof(DISK_GEOMETRY_EX) + 256);
    if (DeviceIoControl(handle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, nullptr, 0,
                        buffer.data(), static_cast<DWORD>(buffer.size()), &returned, nullptr)) {
        auto* geometry = reinterpret_cast<DISK_GEOMETRY_EX*>(buffer.data());
        size = static_cast<unsigned long long>(geometry->DiskSize.QuadPart);
        return true;
    }
    return false;
}

std::wstring MakeDiskName(const StorageDescriptor& descriptor)
{
    std::wstring name;
    if (!descriptor.vendor.empty()) {
        name = descriptor.vendor;
    }
    if (!descriptor.product.empty()) {
        if (!name.empty()) {
            name += L" ";
        }
        name += descriptor.product;
    }
    name = TrimWide(name);
    if (name.empty()) {
        name = L"(без названия)";
    }
    return name;
}

bool AddDisk(std::vector<DiskItem>& disks, int number,
             const StorageDescriptor& descriptor, unsigned long long size)
{
    DiskItem item;
    item.number = number;
    item.size = size;
    item.name = MakeDiskName(descriptor);
    item.cap = NominalSize(size);
    disks.push_back(std::move(item));
    return true;
}

bool GetVolumeExtents(wchar_t letter,
                      std::vector<std::pair<int, unsigned long long>>& extents)
{
    HANDLE handle = OpenVolume(letter);
    if (handle == INVALID_HANDLE_VALUE) {
        return false;
    }

    std::vector<BYTE> buffer(sizeof(VOLUME_DISK_EXTENTS) + 32 * sizeof(DISK_EXTENT));
    DWORD returned = 0;
    bool ok = DeviceIoControl(handle, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, nullptr, 0,
                              buffer.data(), static_cast<DWORD>(buffer.size()),
                              &returned, nullptr) != FALSE;
    CloseHandle(handle);
    if (!ok || returned < sizeof(VOLUME_DISK_EXTENTS)) {
        return false;
    }

    auto* info = reinterpret_cast<VOLUME_DISK_EXTENTS*>(buffer.data());
    for (DWORD i = 0; i < info->NumberOfDiskExtents; ++i) {
        extents.emplace_back(static_cast<int>(info->Extents[i].DiskNumber),
                             static_cast<unsigned long long>(info->Extents[i].StartingOffset.QuadPart));
    }
    return !extents.empty();
}

bool FindVolumeForPartition(int disk, unsigned long long offset,
                            wchar_t& letter, std::wstring& label, std::wstring& fileSystem)
{
    for (wchar_t ch = L'A'; ch <= L'Z'; ++ch) {
        std::wstring root;
        root.push_back(ch);
        root += L":\\";
        if (GetDriveTypeW(root.c_str()) == DRIVE_NO_ROOT_DIR) {
            continue;
        }

        std::vector<std::pair<int, unsigned long long>> extents;
        if (!GetVolumeExtents(ch, extents)) {
            continue;
        }

        bool matched = false;
        for (const auto& extent : extents) {
            if (extent.first == disk && extent.second == offset) {
                matched = true;
                break;
            }
        }
        if (!matched) {
            continue;
        }

        letter = ch;
        wchar_t volumeName[256] = {};
        wchar_t fsName[64] = {};
        DWORD serial = 0;
        DWORD maxComponent = 0;
        DWORD flags = 0;
        if (GetVolumeInformationW(root.c_str(), volumeName, 256, &serial,
                                  &maxComponent, &flags, fsName, 64)) {
            label = TrimWide(volumeName);
            fileSystem = fsName;
        }
        return true;
    }
    return false;
}

} // namespace

std::vector<DiskItem> ScanUsbDisks()
{
    std::vector<DiskItem> disks;
    struct Candidate {
        int number;
        bool removable;
        unsigned long long size;
        StorageDescriptor descriptor;
    };
    std::vector<Candidate> candidates;

    for (int disk = 0; disk < 128; ++disk) {
        HANDLE handle = OpenPhysicalDrive(disk);
        if (handle == INVALID_HANDLE_VALUE) {
            continue;
        }

        StorageDescriptor descriptor = QueryStorageDescriptor(handle);
        unsigned long long size = 0;
        QueryDiskSize(handle, size);
        CloseHandle(handle);

        if (descriptor.valid && descriptor.busType == BusTypeUsb) {
            AddDisk(disks, disk, descriptor, size);
        } else {
            candidates.push_back({disk, descriptor.valid && descriptor.removable, size, descriptor});
        }
    }

    if (disks.empty()) {
        for (const auto& candidate : candidates) {
            if (candidate.descriptor.valid &&
                (candidate.descriptor.busType == BusTypeUsb || candidate.removable)) {
                AddDisk(disks, candidate.number, candidate.descriptor, candidate.size);
            }
        }
    }

    return disks;
}

bool GetDiskSizeBytes(int disk, unsigned long long& size)
{
    HANDLE handle = OpenPhysicalDrive(disk);
    if (handle == INVALID_HANDLE_VALUE) {
        return false;
    }
    bool ok = QueryDiskSize(handle, size);
    CloseHandle(handle);
    return ok;
}

std::vector<int> GetSystemDiskIds()
{
    std::vector<int> result;
    wchar_t windowsDir[MAX_PATH] = {};
    if (GetWindowsDirectoryW(windowsDir, MAX_PATH) == 0) {
        return result;
    }
    if (windowsDir[1] != L':') {
        return result;
    }

    wchar_t letter = static_cast<wchar_t>(towupper(windowsDir[0]));
    std::vector<std::pair<int, unsigned long long>> extents;
    if (!GetVolumeExtents(letter, extents)) {
        return result;
    }

    for (const auto& extent : extents) {
        if (std::find(result.begin(), result.end(), extent.first) == result.end()) {
            result.push_back(extent.first);
        }
    }
    return result;
}

double NominalGb(unsigned long long bytes)
{
    double gb = static_cast<double>(bytes) / 1000000000.0;
    static const int list[] = {
        1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1000, 1500, 2000, 3000,
        4000, 6000, 8000, 10000, 12000, 14000, 16000, 18000, 20000
    };
    for (int value : list) {
        if (gb <= static_cast<double>(value) * 1.02) {
            return static_cast<double>(value);
        }
    }
    return gb;
}

std::wstring NominalSize(unsigned long long bytes)
{
    double value = NominalGb(bytes);
    if (value >= 1000.0) {
        return Format(L"%g ТБ", value / 1000.0);
    }
    return Format(L"%d ГБ", static_cast<int>(value));
}

std::vector<PartitionInfo> GetPartitions(int disk, bool& ok)
{
    ok = false;
    std::vector<PartitionInfo> result;

    HANDLE handle = OpenPhysicalDrive(disk);
    if (handle == INVALID_HANDLE_VALUE) {
        return result;
    }

    std::vector<BYTE> buffer(sizeof(DRIVE_LAYOUT_INFORMATION_EX) +
                             256 * sizeof(PARTITION_INFORMATION_EX));
    DWORD returned = 0;
    bool success = DeviceIoControl(handle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, nullptr, 0,
                                   buffer.data(), static_cast<DWORD>(buffer.size()),
                                   &returned, nullptr) != FALSE;
    if (!success && GetLastError() == ERROR_INSUFFICIENT_BUFFER && returned > buffer.size()) {
        buffer.resize(returned);
        success = DeviceIoControl(handle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, nullptr, 0,
                                  buffer.data(), static_cast<DWORD>(buffer.size()),
                                  &returned, nullptr) != FALSE;
    }
    CloseHandle(handle);

    if (!success || returned < sizeof(DRIVE_LAYOUT_INFORMATION_EX)) {
        return result;
    }

    auto* layout = reinterpret_cast<DRIVE_LAYOUT_INFORMATION_EX*>(buffer.data());
    for (DWORD i = 0; i < layout->PartitionCount; ++i) {
        const PARTITION_INFORMATION_EX& entry = layout->PartitionEntry[i];
        if (entry.PartitionLength.QuadPart <= 0) {
            continue;
        }
        PartitionInfo info;
        info.offset = static_cast<unsigned long long>(entry.StartingOffset.QuadPart);
        info.size = static_cast<unsigned long long>(entry.PartitionLength.QuadPart);
        info.number = static_cast<int>(entry.PartitionNumber);
        result.push_back(info);
    }

    std::sort(result.begin(), result.end(), [](const PartitionInfo& a, const PartitionInfo& b) {
        return a.offset < b.offset;
    });
    ok = true;
    return result;
}

std::wstring GetPartitionLetter(int disk, int idx)
{
    bool ok = false;
    std::vector<PartitionInfo> partitions = GetPartitions(disk, ok);
    if (!ok || idx < 0 || idx >= static_cast<int>(partitions.size())) {
        return std::wstring();
    }

    wchar_t letter = 0;
    std::wstring label;
    std::wstring fileSystem;
    if (!FindVolumeForPartition(disk, partitions[idx].offset, letter, label, fileSystem)) {
        return std::wstring();
    }

    std::wstring result;
    result.push_back(letter);
    result += L":\\";
    return result;
}

bool DiskIsInstalled(int disk)
{
    std::wstring first = GetPartitionLetter(disk, 0);
    if (first.empty() || !FileExists(first + L"efi\\Paravis\\Boot\\NOT.DEL")) {
        return false;
    }
    std::wstring third = GetPartitionLetter(disk, 2);
    if (third.empty() || !FileExists(third + L"ISO\\Ventoy\\NOT.DEL")) {
        return false;
    }
    return true;
}

std::wstring VerifyLayout(int disk, bool& ok)
{
    ok = false;
    bool layoutOk = false;
    std::vector<PartitionInfo> partitions = GetPartitions(disk, layoutOk);
    if (!layoutOk) {
        return L"разделов: ? (нет доступа к WMI)";
    }

    const wchar_t* want[] = { L"PARAVIS", L"VTOYEFI", L"DATA" };
    std::wstring report;
    std::wstring labels = L"|";
    unsigned long long sizes[3] = { 0, 0, 0 };
    int count = static_cast<int>(partitions.size());

    for (int i = 0; i < count; ++i) {
        unsigned long long megabytes = partitions[i].size / 1048576ULL;
        if (i < 3) {
            sizes[i] = megabytes;
        }
        std::wstring line = Format(L"  раздел %d: %llu МБ", i + 1, megabytes);

        wchar_t letter = 0;
        std::wstring label;
        std::wstring fileSystem;
        if (FindVolumeForPartition(disk, partitions[i].offset, letter, label, fileSystem)) {
            std::wstring device;
            device.push_back(letter);
            device += L":";
            line += L"   " + device + L"  " + label + L"  " + fileSystem;
            if (!label.empty()) {
                labels += label + L"|";
            }
        }
        report += line + L"\r\n";
    }

    std::wstring missing;
    for (const wchar_t* name : want) {
        std::wstring token = std::wstring(L"|") + name + L"|";
        if (labels.find(token) == std::wstring::npos) {
            missing += L" ";
            missing += name;
        }
    }

    auto absDiff = [](unsigned long long value, unsigned long long expected) -> long long {
        return value >= expected
            ? static_cast<long long>(value - expected)
            : static_cast<long long>(expected - value);
    };
    bool sizeOk = absDiff(sizes[0], 16384) <= 8 && absDiff(sizes[1], 32) <= 8;

    std::wstring head = Format(L"разделов: %d", count);
    if (count != 3) {
        head += L"  (ожидалось 3)";
    }
    if (!missing.empty()) {
        report += L"нет меток:" + missing + L"\r\n";
    }
    if (!sizeOk && count == 3) {
        report += L"не совпали размеры 1-го и 2-го разделов\r\n";
    }

    ok = count == 3 && missing.empty() && sizeOk;
    return head + L"\r\n" + report;
}