#include "battery.h"

#include "common/io.h"
#include "common/windows/nt.h"
#include "common/windows/unicode.h"
#include "common/mallocHelper.h"
#include "common/smbios.h"

#undef WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <batclass.h>
#include <devguid.h>
#include <cfgmgr32.h>

static const char* detectWithCmApi(FFBatteryOptions* options, FFlist* results) {
    // https://learn.microsoft.com/en-us/windows-hardware/drivers/install/using-device-interfaces
    ULONG cchDeviceInterfaces = 0;
    if (CM_Get_Device_Interface_List_SizeW(
            &cchDeviceInterfaces,
            (LPGUID) &GUID_DEVCLASS_BATTERY,
            NULL,
            CM_GET_DEVICE_INTERFACE_LIST_PRESENT) != CR_SUCCESS) {
        return "CM_Get_Device_Interface_List_SizeW() failed";
    }

    if (cchDeviceInterfaces <= 1) {
        return NULL; // Not found
    }

    wchar_t* FF_AUTO_FREE mszDeviceInterfaces = (wchar_t*) malloc(cchDeviceInterfaces * sizeof(wchar_t));
    if (CM_Get_Device_Interface_ListW(
            (LPGUID) &GUID_DEVCLASS_BATTERY,
            NULL,
            mszDeviceInterfaces,
            cchDeviceInterfaces,
            CM_GET_DEVICE_INTERFACE_LIST_PRESENT) != CR_SUCCESS) {
        return "CM_Get_Device_Interface_ListW() failed";
    }

    for (const wchar_t* p = mszDeviceInterfaces; *p; p += wcslen(p) + 1) {
        HANDLE FF_AUTO_CLOSE_FD hBattery =
            CreateFileW(p, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

        if (hBattery == INVALID_HANDLE_VALUE) {
            continue;
        }

        BATTERY_QUERY_INFORMATION bqi = { .InformationLevel = BatteryInformation };

        DWORD dwWait = 0;
        DWORD dwOut;

        if (!DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_TAG, &dwWait, sizeof(dwWait), &bqi.BatteryTag, sizeof(bqi.BatteryTag), &dwOut, NULL) && bqi.BatteryTag) {
            continue;
        }

        BATTERY_INFORMATION bi = { 0 };
        if (!DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), &bi, sizeof(bi), &dwOut, NULL)) {
            continue;
        }

        if (!(bi.Capabilities & BATTERY_SYSTEM_BATTERY)) {
            continue;
        }

        FFBatteryResult* battery = FF_LIST_ADD(FFBatteryResult, *results);

        if (memcmp(bi.Chemistry, "PbAc", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Lead Acid");
        } else if (memcmp(bi.Chemistry, "LION", 4) == 0 || memcmp(bi.Chemistry, "Li-I", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Lithium Ion");
        } else if (memcmp(bi.Chemistry, "NiCd", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Nickel Cadmium");
        } else if (memcmp(bi.Chemistry, "NiMH", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Nickel Metal Hydride");
        } else if (memcmp(bi.Chemistry, "NiZn", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Nickel Zinc");
        } else if (memcmp(bi.Chemistry, "RAM\0", 4) == 0) {
            ffStrbufInitStatic(&battery->technology, "Rechargeable Alkaline-Manganese");
        } else {
            ffStrbufInitStatic(&battery->technology, "Unknown");
        }

        {
            ffStrbufInit(&battery->modelName);
            bqi.InformationLevel = BatteryDeviceName;
            wchar_t name[64];
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), name, sizeof(name), &dwOut, NULL)) {
                ffStrbufSetWS(&battery->modelName, name);
            }
        }

        {
            ffStrbufInit(&battery->manufacturer);
            bqi.InformationLevel = BatteryManufactureName;
            wchar_t name[64];
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), name, sizeof(name), &dwOut, NULL)) {
                ffStrbufSetWS(&battery->manufacturer, name);
            }
        }

        {
            ffStrbufInit(&battery->manufactureDate);
            bqi.InformationLevel = BatteryManufactureDate;
            BATTERY_MANUFACTURE_DATE date;
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), &date, sizeof(date), &dwOut, NULL)) {
                ffStrbufSetF(&battery->manufactureDate, "%.4d-%.2d-%.2d", date.Year < 1000 ? date.Year + 1900 : date.Year, date.Month, date.Day);
            }
        }

        {
            ffStrbufInit(&battery->serial);
            bqi.InformationLevel = BatterySerialNumber;
            wchar_t name[64];
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), name, sizeof(name), &dwOut, NULL)) {
                ffStrbufSetWS(&battery->serial, name);
            }
        }

        battery->cycleCount = bi.CycleCount;

        battery->temperature = FF_BATTERY_TEMP_UNSET;
        if (options->temp) {
            bqi.InformationLevel = BatteryTemperature;
            ULONG temp;
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), &temp, sizeof(temp), &dwOut, NULL)) {
                battery->temperature = temp / 10.0 - 273.15;
            }
        }

        {
            bqi.InformationLevel = BatteryEstimatedTime;
            ULONG time;
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_INFORMATION, &bqi, sizeof(bqi), &time, sizeof(time), &dwOut, NULL)) {
                battery->timeRemaining = time == BATTERY_UNKNOWN_TIME ? -1 : (int32_t) time;
            }
        }

        {
            BATTERY_STATUS bs;
            BATTERY_WAIT_STATUS bws = { .BatteryTag = bqi.BatteryTag };
            if (DeviceIoControl(hBattery, IOCTL_BATTERY_QUERY_STATUS, &bws, sizeof(bws), &bs, sizeof(bs), &dwOut, NULL) && bs.Capacity != BATTERY_UNKNOWN_CAPACITY && bi.FullChargedCapacity != 0) {
                battery->capacity = bs.Capacity * 100.0 / bi.FullChargedCapacity;

                battery->status = FF_BATTERY_STATUS_NONE;
                if (bs.PowerState & BATTERY_POWER_ON_LINE) {
                    battery->status |= FF_BATTERY_STATUS_AC_CONNECTED;
                }
                if (bs.PowerState & BATTERY_DISCHARGING) {
                    battery->status |= FF_BATTERY_STATUS_DISCHARGING;
                }
                if (bs.PowerState & BATTERY_CHARGING) {
                    battery->status |= FF_BATTERY_STATUS_CHARGING;
                }
                if (bs.PowerState & BATTERY_CRITICAL) {
                    battery->status |= FF_BATTERY_STATUS_CRITICAL;
                }
            } else {
                battery->status = FF_BATTERY_STATUS_UNKNOWN;
                battery->capacity = 0;
            }
        }
    }
    return NULL;
}

typedef struct FFSmbiosPortableBattery {
    FFSmbiosHeader Header;

    // 2.1+
    uint8_t Location;                  // string
    uint8_t Manufacturer;              // string
    uint8_t ManufactureDate;           // string
    uint8_t SerialNumber;              // string
    uint8_t DeviceName;                // string
    uint8_t DeviceChemistry;           // enum
    uint16_t DesignCapacity;           // varies
    uint16_t DesignVoltage;            // varies
    uint8_t SbdsVersionNumber;         // string
    uint8_t MaximumErrorInBatteryData; // varies

    // 2.2+
    uint16_t SbdsSerialNumber;        // varies
    uint16_t SbdsManufactureDate;     // varies
    uint8_t SbdsDeviceChemistry;      // string
    uint8_t DesignCapacityMultiplier; // varies
    uint16_t OEMSpecific;             // varies
} FF_A_PACKED FFSmbiosPortableBattery;

static_assert(offsetof(FFSmbiosPortableBattery, OEMSpecific) == 0x16,
    "FFSmbiosPortableBattery: Wrong struct alignment");

static const char* detectBySmbios(FFBatteryResult* battery) {
    const FFSmbiosHeaderTable* smbiosTable = ffGetSmbiosHeaderTable();
    if (!smbiosTable) {
        return "Failed to get SMBIOS data";
    }

    const FFSmbiosPortableBattery* data = (const FFSmbiosPortableBattery*) (*smbiosTable)[FF_SMBIOS_TYPE_PORTABLE_BATTERY];
    if (!data) {
        return "Portable battery section is not found in SMBIOS data";
    }

    const char* strings = (const char*) data + data->Header.Length;

    ffStrbufSetStatic(&battery->modelName, ffSmbiosLocateString(strings, data->DeviceName));
    ffCleanUpSmbiosValue(&battery->modelName);
    ffStrbufSetStatic(&battery->manufacturer, ffSmbiosLocateString(strings, data->Manufacturer));
    ffCleanUpSmbiosValue(&battery->manufacturer);

    if (data->ManufactureDate) {
        ffStrbufSetStatic(&battery->manufactureDate, ffSmbiosLocateString(strings, data->ManufactureDate));
        ffCleanUpSmbiosValue(&battery->manufactureDate);
    } else if (data->Header.Length > offsetof(FFSmbiosPortableBattery, SbdsManufactureDate)) {
        int day = data->SbdsManufactureDate & 0b11111;
        int month = (data->SbdsManufactureDate >> 5) & 0b1111;
        int year = (data->SbdsManufactureDate >> 9) + 1800;
        ffStrbufSetF(&battery->manufactureDate, "%.4d-%.2d-%.2d", year, month, day);
    }

    switch (data->DeviceChemistry) {
        case 0x01:
            ffStrbufSetStatic(&battery->technology, "Other");
            break;
        case 0x02:
            ffStrbufSetStatic(&battery->technology, "Unknown");
            break;
        case 0x03:
            ffStrbufSetStatic(&battery->technology, "Lead Acid");
            break;
        case 0x04:
            ffStrbufSetStatic(&battery->technology, "Nickel Cadmium");
            break;
        case 0x05:
            ffStrbufSetStatic(&battery->technology, "Nickel metal hydride");
            break;
        case 0x06:
            ffStrbufSetStatic(&battery->technology, "Lithium-ion");
            break;
        case 0x07:
            ffStrbufSetStatic(&battery->technology, "Zinc air");
            break;
        case 0x08:
            ffStrbufSetStatic(&battery->technology, "Lithium Polymer");
            break;
    }

    if (data->SerialNumber) {
        ffStrbufSetStatic(&battery->serial, ffSmbiosLocateString(strings, data->SerialNumber));
        ffCleanUpSmbiosValue(&battery->serial);
    } else if (data->Header.Length > offsetof(FFSmbiosPortableBattery, SbdsSerialNumber)) {
        ffStrbufSetF(&battery->serial, "%4X", data->SbdsSerialNumber);
    }

    return NULL;
}

static const char* detectWithNtApi(FF_A_UNUSED FFBatteryOptions* options, FFlist* results) {
    SYSTEM_BATTERY_STATE info;
    if (NT_SUCCESS(NtPowerInformation(SystemBatteryState, NULL, 0, &info, sizeof(info))) &&
        info.BatteryPresent) {
        FFBatteryResult* battery = FF_LIST_ADD(FFBatteryResult, *results);
        ffStrbufInit(&battery->modelName);
        ffStrbufInit(&battery->manufacturer);
        ffStrbufInit(&battery->manufactureDate);
        ffStrbufInit(&battery->technology);
        ffStrbufInit(&battery->serial);
        battery->temperature = FF_BATTERY_TEMP_UNSET;
        battery->cycleCount = 0;
        battery->timeRemaining = info.EstimatedTime == BATTERY_UNKNOWN_TIME ? -1 : (int32_t) info.EstimatedTime;
        battery->status = FF_BATTERY_STATUS_NONE;

        battery->capacity = info.RemainingCapacity * 100.0 / info.MaxCapacity;
        if (info.AcOnLine) {
            battery->status |= FF_BATTERY_STATUS_AC_CONNECTED;
        }
        if (info.Charging) {
            battery->status |= FF_BATTERY_STATUS_CHARGING;
        }
        if (info.Discharging) {
            battery->status |= FF_BATTERY_STATUS_DISCHARGING;
        }
        if (info.DefaultAlert1 > 0 && info.RemainingCapacity <= info.DefaultAlert1) {
            battery->status |= FF_BATTERY_STATUS_CRITICAL;
        }

        detectBySmbios(battery);

        return NULL;
    }
    return "NtPowerInformation(SystemBatteryState) failed";
}

const char* ffDetectBattery(FFBatteryOptions* options, FFlist* results) {
    return options->useSetupApi
        ? detectWithCmApi(options, results)
        : detectWithNtApi(options, results);
}
