#include "diskio.h"
#include "common/io.h"
#include "common/windows/unicode.h"

#include <windows.h>
#include <winioctl.h>

static bool detectPhysicalDisk(const wchar_t* szDevice, FFlist* result, FFDiskIOOptions* options) {
    FF_AUTO_CLOSE_FD HANDLE hDevice = CreateFileW(szDevice, FILE_READ_ATTRIBUTES, FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
    if (hDevice == INVALID_HANDLE_VALUE) {
        return false;
    }

    DWORD retSize;
    char sddBuffer[4096];
    if (!DeviceIoControl(
            hDevice,
            IOCTL_STORAGE_QUERY_PROPERTY,
            &(STORAGE_PROPERTY_QUERY) {
                .PropertyId = StorageDeviceProperty,
                .QueryType = PropertyStandardQuery,
            },
            sizeof(STORAGE_PROPERTY_QUERY),
            &sddBuffer,
            ARRAY_SIZE(sddBuffer),
            &retSize,
            nullptr) ||
        retSize == 0) {
        return true;
    }

    FFDiskIOResult* device = FF_LIST_ADD(FFDiskIOResult, *result);
    STORAGE_DEVICE_DESCRIPTOR* sdd = (STORAGE_DEVICE_DESCRIPTOR*) sddBuffer;

    ffStrbufInit(&device->name);
    if (sdd->VendorIdOffset != 0) {
        ffStrbufSetS(&device->name, (const char*) sddBuffer + sdd->VendorIdOffset);
        ffStrbufTrim(&device->name, ' ');
    }
    if (sdd->ProductIdOffset != 0) {
        if (device->name.length) {
            ffStrbufAppendC(&device->name, ' ');
        }

        ffStrbufAppendS(&device->name, (const char*) sddBuffer + sdd->ProductIdOffset);
        ffStrbufTrimRight(&device->name, ' ');
    }

    if (!device->name.length) {
        ffStrbufSetWS(&device->name, szDevice);
    }

    if (options->namePrefix.length && !ffStrbufStartsWith(&device->name, &options->namePrefix)) {
        ffStrbufDestroy(&device->name);
        result->length--;
        return true;
    }

    DISK_PERFORMANCE dp = {};
    if (DeviceIoControl(hDevice, IOCTL_DISK_PERFORMANCE, nullptr, 0, &dp, sizeof(dp), &retSize, nullptr)) {
        device->bytesRead = (uint64_t) dp.BytesRead.QuadPart;
        device->readCount = (uint64_t) dp.ReadCount;
        device->bytesWritten = (uint64_t) dp.BytesWritten.QuadPart;
        device->writeCount = (uint64_t) dp.WriteCount;
    } else {
        ffStrbufDestroy(&device->name);
        result->length--;
    }

    ffStrbufInitWS(&device->devPath, szDevice);

    return true;
}

const char* ffDiskIOGetIoCounters(FFlist* result, FFDiskIOOptions* options) {
    {
        wchar_t szPhysicalDrive[32] = L"\\\\.\\PhysicalDrive";
        wchar_t* pNum = szPhysicalDrive + strlen("\\\\.\\PhysicalDrive");
        for (uint32_t idev = 0;; ++idev) {
            _ultow(idev, pNum, 10);

            if (!detectPhysicalDisk(szPhysicalDrive, result, options)) {
                break;
            }
        }
    }

    {
        wchar_t szCdrom[32] = L"\\\\.\\CDROM";
        wchar_t* pNum = szCdrom + strlen("\\\\.\\CDROM");
        for (uint32_t idev = 0;; ++idev) {
            _ultow(idev, pNum, 10);

            if (!detectPhysicalDisk(szCdrom, result, options)) {
                break;
            }
        }
    }

    return nullptr;
}
