// Dynamic partition discovery. This file owns MBR/GPT parsing and keeps structural identity out of filesystem providers.

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <common/boot_volume_selector.h>
#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddk/mutex.h>
#include <kernel/io_partition.h>
#include <kernel/ke_debug.h>
#include <kernel/mm_pool.h>

#define IOP_SECTOR_SIZE 512u
#define IOP_GPT_SIGNATURE "EFI PART"
#define IOP_GPT_REVISION_1_0 0x00010000u
#define IOP_GPT_MIN_ENTRY_SIZE 128u

#pragma pack(push, 1)
typedef struct IOP_MBR_PARTITION_ENTRY {
    uint8_t BootIndicator;
    uint8_t StartingChs[3];
    uint8_t Type;
    uint8_t EndingChs[3];
    uint32_t StartingLba;
    uint32_t SectorCount;
} IOP_MBR_PARTITION_ENTRY;

typedef struct IOP_GPT_HEADER {
    char Signature[8];
    uint32_t Revision;
    uint32_t HeaderSize;
    uint32_t HeaderCrc32;
    uint32_t Reserved;
    uint64_t MyLba;
    uint64_t AlternateLba;
    uint64_t FirstUsableLba;
    uint64_t LastUsableLba;
    BOOT_VOLUME_GUID DiskGuid;
    uint64_t PartitionEntryLba;
    uint32_t NumberOfPartitionEntries;
    uint32_t SizeOfPartitionEntry;
    uint32_t PartitionEntryArrayCrc32;
} IOP_GPT_HEADER;

typedef struct IOP_GPT_ENTRY_PREFIX {
    BOOT_VOLUME_GUID PartitionTypeGuid;
    BOOT_VOLUME_GUID UniquePartitionGuid;
    uint64_t StartingLba;
    uint64_t EndingLba;
    uint64_t Attributes;
} IOP_GPT_ENTRY_PREFIX;
#pragma pack(pop)

_Static_assert(sizeof(IOP_MBR_PARTITION_ENTRY) == 16u, "MBR partition entry layout must match disk format");
_Static_assert(sizeof(IOP_GPT_HEADER) == 92u, "GPT header prefix layout must match disk format");
_Static_assert(sizeof(IOP_GPT_ENTRY_PREFIX) == 56u, "GPT entry prefix layout must match disk format");

typedef struct IOP_PARTITION_RECORD {
    struct IOP_PARTITION_RECORD* Next;
    IO_PARTITION_INFO Info;
    BOOT_VOLUME_SELECTOR Identity;
} IOP_PARTITION_RECORD;

typedef struct IOP_PARTITION_CALLBACK_ENTRY {
    struct IOP_PARTITION_CALLBACK_ENTRY* Next;
    IO_PARTITION_CHANGE_CALLBACK Callback;
    void* Context;
} IOP_PARTITION_CALLBACK_ENTRY;

typedef struct IOP_GPT_VALIDATED_HEADER {
    IOP_GPT_HEADER Header;
    uint64_t EntryBytes;
    uint64_t EntrySectors;
} IOP_GPT_VALIDATED_HEADER;

static MUTEX* gPartitionMutex;
static IOP_PARTITION_RECORD* gPartitions;
static IOP_PARTITION_CALLBACK_ENTRY* gPartitionCallbacks;
static BOOT_VOLUME_SELECTOR gSystemVolumeSelector;
static bool gPartitionManagerInitialized;

// Returns whether two shared GUID values are identical.
static bool IopGuidEquals(const BOOT_VOLUME_GUID* left, const BOOT_VOLUME_GUID* right) {
    return left && right && memcmp(left, right, sizeof(*left)) == 0;
}

// Returns whether one shared GUID value is all zeroes.
static bool IopGuidIsZero(const BOOT_VOLUME_GUID* guid) {
    static const BOOT_VOLUME_GUID zero = {0};
    return guid && IopGuidEquals(guid, &zero);
}

// Returns whether one internal partition identity equals another.
static bool IopIdentityEquals(const BOOT_VOLUME_SELECTOR* left, const BOOT_VOLUME_SELECTOR* right) {
    if (!left || !right || left->Kind != right->Kind)
        return false;

    if (left->Kind == BootVolumeSelectorMbr) {
        return left->Identity.Mbr.DiskSignature == right->Identity.Mbr.DiskSignature &&
               left->Identity.Mbr.PartitionNumber == right->Identity.Mbr.PartitionNumber;
    }
    if (left->Kind == BootVolumeSelectorGpt) {
        return IopGuidEquals(&left->Identity.Gpt.PartitionGuid,
                             &right->Identity.Gpt.PartitionGuid);
    }
    return false;
}

// Acquires the partition-manager mutex for registry and callback serialization.
static bool IopAcquirePartitionMutex(void) {
    if (!gPartitionMutex)
        return false;
    return KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(gPartitionMutex, KE_WAIT_INFINITE));
}

// Frees one dynamically discovered partition list.
static void IopFreePartitionList(IOP_PARTITION_RECORD* head) {
    while (head) {
        IOP_PARTITION_RECORD* next = head->Next;
        MmFreePool(head);
        head = next;
    }
}

// Updates one CRC32 accumulator using the GPT/IEEE polynomial.
static uint32_t IopCrc32Update(uint32_t crc, const uint8_t* data, size_t length) {
    for (size_t i = 0; i < length; ++i) {
        crc ^= data[i];
        for (uint32_t bit = 0; bit < 8; ++bit)
            crc = (crc >> 1) ^ (0xedb88320u & (uint32_t)-(int32_t)(crc & 1u));
    }
    return crc;
}

// Computes a complete GPT CRC32 value.
static uint32_t IopCrc32(const void* data, size_t length) {
    uint32_t crc = IopCrc32Update(0xffffffffu, (const uint8_t*)data, length);
    return crc ^ 0xffffffffu;
}

// Reads an arbitrary byte span through the kernel's 512-byte block-device API.
static KSTATUS IopReadDiskBytes(uint8_t diskIndex,
                                uint64_t diskSectors,
                                uint64_t byteOffset,
                                void* buffer,
                                size_t length) {
    if (!buffer || length == 0)
        return NTLSTATUS_INVALID_PARAMETER;
    if (diskSectors > UINT64_MAX / IOP_SECTOR_SIZE)
        return NTLSTATUS_INVALID_PARAMETER;

    uint64_t diskBytes = diskSectors * IOP_SECTOR_SIZE;
    if (byteOffset > diskBytes || (uint64_t)length > diskBytes - byteOffset)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t sector[IOP_SECTOR_SIZE];
    uint8_t* out = (uint8_t*)buffer;
    size_t remaining = length;
    while (remaining) {
        uint64_t lba = byteOffset / IOP_SECTOR_SIZE;
        uint32_t inSector = (uint32_t)(byteOffset % IOP_SECTOR_SIZE);
        size_t chunk = IOP_SECTOR_SIZE - inSector;
        if (chunk > remaining)
            chunk = remaining;

        KSTATUS status = IoReadDiskSectors(diskIndex, lba, 1, sector);
        if (!KSUCCESS(status))
            return status;
        memcpy(out, sector + inSector, chunk);
        out += chunk;
        byteOffset += chunk;
        remaining -= chunk;
    }
    return NTLSTATUS_SUCCESS;
}

// Returns whether one MBR type is an extended-partition container we intentionally do not enumerate.
static bool IopMbrTypeIsExtended(uint8_t type) {
    return type == 0x05u || type == 0x0fu || type == 0x85u;
}

// Allocates one immutable partition record.
static IOP_PARTITION_RECORD* IopAllocatePartition(uint8_t diskIndex,
                                                   uint64_t startLba,
                                                   uint64_t sectorCount,
                                                   const BOOT_VOLUME_SELECTOR* identity) {
    if (!sectorCount || !identity)
        return NULL;

    IOP_PARTITION_RECORD* record = MmAllocatePool(sizeof(*record));
    if (!record)
        return NULL;
    memset(record, 0, sizeof(*record));
    record->Info.DiskIndex = diskIndex;
    record->Info.StartLba = startLba;
    record->Info.SectorCount = sectorCount;
    record->Identity = *identity;
    return record;
}

// Inserts one discovered partition in ascending LBA order and rejects overlap or duplicate identity.
static KSTATUS IopInsertDiscoveredPartition(IOP_PARTITION_RECORD** head,
                                             IOP_PARTITION_RECORD* record) {
    if (!head || !record || record->Info.SectorCount == 0)
        return NTLSTATUS_INVALID_PARAMETER;

    if (record->Info.StartLba > UINT64_MAX - record->Info.SectorCount)
        return NTLSTATUS_INVALID_PARAMETER;
    uint64_t recordEnd = record->Info.StartLba + record->Info.SectorCount;

    IOP_PARTITION_RECORD* previous = NULL;
    IOP_PARTITION_RECORD* current = *head;
    while (current && current->Info.StartLba < record->Info.StartLba) {
        if (IopIdentityEquals(&current->Identity, &record->Identity))
            return NTLSTATUS_OBJECT_NAME_COLLISION;
        previous = current;
        current = current->Next;
    }
    for (IOP_PARTITION_RECORD* scan = current; scan; scan = scan->Next) {
        if (IopIdentityEquals(&scan->Identity, &record->Identity))
            return NTLSTATUS_OBJECT_NAME_COLLISION;
    }

    if (previous) {
        uint64_t previousEnd = previous->Info.StartLba + previous->Info.SectorCount;
        if (previousEnd > record->Info.StartLba)
            return NTLSTATUS_INVALID_PARAMETER;
    }
    if (current && recordEnd > current->Info.StartLba)
        return NTLSTATUS_INVALID_PARAMETER;

    record->Next = current;
    if (previous)
        previous->Next = record;
    else
        *head = record;
    return NTLSTATUS_SUCCESS;
}

// Validates the exact CRC32 of one GPT partition-entry array without imposing an entry-count cap.
static KSTATUS IopValidateGptEntryArrayCrc(uint8_t diskIndex,
                                           const IOP_GPT_VALIDATED_HEADER* validated) {
    uint8_t sector[IOP_SECTOR_SIZE];
    uint64_t remaining = validated->EntryBytes;
    uint64_t lba = validated->Header.PartitionEntryLba;
    uint32_t crc = 0xffffffffu;
    while (remaining) {
        KSTATUS status = IoReadDiskSectors(diskIndex, lba++, 1, sector);
        if (!KSUCCESS(status))
            return status;
        size_t chunk = remaining < IOP_SECTOR_SIZE ? (size_t)remaining : IOP_SECTOR_SIZE;
        crc = IopCrc32Update(crc, sector, chunk);
        remaining -= chunk;
    }
    crc ^= 0xffffffffu;
    return crc == validated->Header.PartitionEntryArrayCrc32
         ? NTLSTATUS_SUCCESS
         : NTLSTATUS_CHECKSUM_MISMATCH;
}

// Validates one primary or backup GPT header and its partition-entry-array CRC.
static KSTATUS IopValidateGptHeader(uint8_t diskIndex,
                                    uint64_t diskSectors,
                                    uint64_t headerLba,
                                    uint64_t expectedAlternateLba,
                                    IOP_GPT_VALIDATED_HEADER* validatedOut) {
    if (!validatedOut || diskSectors < 3 || headerLba >= diskSectors ||
        expectedAlternateLba >= diskSectors || headerLba == expectedAlternateLba) {
        return NTLSTATUS_INVALID_PARAMETER;
    }

    uint8_t sector[IOP_SECTOR_SIZE];
    KSTATUS status = IoReadDiskSectors(diskIndex, headerLba, 1, sector);
    if (!KSUCCESS(status))
        return status;

    IOP_GPT_HEADER header;
    memcpy(&header, sector, sizeof(header));
    if (memcmp(header.Signature, IOP_GPT_SIGNATURE, 8) != 0 ||
        header.Revision < IOP_GPT_REVISION_1_0 ||
        header.HeaderSize < sizeof(IOP_GPT_HEADER) ||
        header.HeaderSize > IOP_SECTOR_SIZE ||
        header.Reserved != 0 ||
        header.MyLba != headerLba ||
        header.AlternateLba != expectedAlternateLba ||
        header.FirstUsableLba > header.LastUsableLba ||
        header.LastUsableLba >= diskSectors ||
        IopGuidIsZero(&header.DiskGuid) ||
        header.NumberOfPartitionEntries == 0 ||
        header.SizeOfPartitionEntry < IOP_GPT_MIN_ENTRY_SIZE ||
        (header.SizeOfPartitionEntry & 7u) != 0) {
        return NTLSTATUS_INVALID_PARAMETER;
    }

    uint32_t expectedHeaderCrc = header.HeaderCrc32;
    memset(sector + offsetof(IOP_GPT_HEADER, HeaderCrc32), 0, sizeof(header.HeaderCrc32));
    if (IopCrc32(sector, header.HeaderSize) != expectedHeaderCrc)
        return NTLSTATUS_CHECKSUM_MISMATCH;

    if ((uint64_t)header.NumberOfPartitionEntries > UINT64_MAX / header.SizeOfPartitionEntry)
        return NTLSTATUS_INVALID_PARAMETER;
    uint64_t entryBytes = (uint64_t)header.NumberOfPartitionEntries * header.SizeOfPartitionEntry;
    if (entryBytes > UINT64_MAX - (IOP_SECTOR_SIZE - 1u))
        return NTLSTATUS_INVALID_PARAMETER;
    uint64_t entrySectors = (entryBytes + IOP_SECTOR_SIZE - 1u) / IOP_SECTOR_SIZE;
    if (entrySectors == 0 || header.PartitionEntryLba >= diskSectors ||
        entrySectors > diskSectors - header.PartitionEntryLba) {
        return NTLSTATUS_INVALID_PARAMETER;
    }

    uint64_t entryLastLba = header.PartitionEntryLba + entrySectors - 1u;
    if (headerLba == 1u) {
        if (header.PartitionEntryLba <= headerLba || entryLastLba >= header.FirstUsableLba)
            return NTLSTATUS_INVALID_PARAMETER;
    } else {
        if (header.PartitionEntryLba <= header.LastUsableLba || entryLastLba >= headerLba)
            return NTLSTATUS_INVALID_PARAMETER;
    }

    IOP_GPT_VALIDATED_HEADER validated = {0};
    validated.Header = header;
    validated.EntryBytes = entryBytes;
    validated.EntrySectors = entrySectors;
    status = IopValidateGptEntryArrayCrc(diskIndex, &validated);
    if (!KSUCCESS(status))
        return status;

    *validatedOut = validated;
    return NTLSTATUS_SUCCESS;
}

// Returns whether independently valid primary and backup GPT headers describe the same table.
static bool IopGptHeadersAgree(const IOP_GPT_VALIDATED_HEADER* primary,
                               const IOP_GPT_VALIDATED_HEADER* backup) {
    if (!primary || !backup)
        return false;
    return IopGuidEquals(&primary->Header.DiskGuid, &backup->Header.DiskGuid) &&
           primary->Header.FirstUsableLba == backup->Header.FirstUsableLba &&
           primary->Header.LastUsableLba == backup->Header.LastUsableLba &&
           primary->Header.NumberOfPartitionEntries == backup->Header.NumberOfPartitionEntries &&
           primary->Header.SizeOfPartitionEntry == backup->Header.SizeOfPartitionEntry &&
           primary->Header.PartitionEntryArrayCrc32 == backup->Header.PartitionEntryArrayCrc32;
}

// Parses all populated GPT entries from one already validated table.
static KSTATUS IopParseGptEntries(uint8_t diskIndex,
                                  uint64_t diskSectors,
                                  const IOP_GPT_VALIDATED_HEADER* validated,
                                  IOP_PARTITION_RECORD** partitionsOut) {
    if (!validated || !partitionsOut)
        return NTLSTATUS_INVALID_PARAMETER;

    IOP_PARTITION_RECORD* partitions = NULL;
    if (validated->Header.PartitionEntryLba > UINT64_MAX / IOP_SECTOR_SIZE)
        return NTLSTATUS_INVALID_PARAMETER;
    uint64_t arrayByteBase = validated->Header.PartitionEntryLba * IOP_SECTOR_SIZE;

    for (uint32_t index = 0; index < validated->Header.NumberOfPartitionEntries; ++index) {
        uint64_t entryOffset = (uint64_t)index * validated->Header.SizeOfPartitionEntry;
        if (arrayByteBase > UINT64_MAX - entryOffset) {
            IopFreePartitionList(partitions);
            return NTLSTATUS_INVALID_PARAMETER;
        }

        IOP_GPT_ENTRY_PREFIX entry;
        KSTATUS status = IopReadDiskBytes(diskIndex,
                                          diskSectors,
                                          arrayByteBase + entryOffset,
                                          &entry,
                                          sizeof(entry));
        if (!KSUCCESS(status)) {
            IopFreePartitionList(partitions);
            return status;
        }
        if (IopGuidIsZero(&entry.PartitionTypeGuid))
            continue;
        if (IopGuidIsZero(&entry.UniquePartitionGuid) ||
            entry.EndingLba < entry.StartingLba ||
            entry.StartingLba < validated->Header.FirstUsableLba ||
            entry.EndingLba > validated->Header.LastUsableLba) {
            IopFreePartitionList(partitions);
            return NTLSTATUS_INVALID_PARAMETER;
        }

        uint64_t sectorCount = entry.EndingLba - entry.StartingLba + 1u;
        BOOT_VOLUME_SELECTOR identity = {0};
        identity.Kind = BootVolumeSelectorGpt;
        identity.Identity.Gpt.PartitionGuid = entry.UniquePartitionGuid;
        IOP_PARTITION_RECORD* record = IopAllocatePartition(diskIndex,
                                                             entry.StartingLba,
                                                             sectorCount,
                                                             &identity);
        if (!record) {
            IopFreePartitionList(partitions);
            return NTLSTATUS_NO_MEMORY;
        }
        status = IopInsertDiscoveredPartition(&partitions, record);
        if (!KSUCCESS(status)) {
            MmFreePool(record);
            IopFreePartitionList(partitions);
            return status;
        }
    }

    *partitionsOut = partitions;
    return NTLSTATUS_SUCCESS;
}

// Discovers one GPT disk from primary/backup metadata and rejects inconsistent valid copies.
static KSTATUS IopDiscoverGptDisk(uint8_t diskIndex,
                                  uint64_t diskSectors,
                                  IOP_PARTITION_RECORD** partitionsOut) {
    IOP_GPT_VALIDATED_HEADER primary = {0};
    IOP_GPT_VALIDATED_HEADER backup = {0};
    KSTATUS primaryStatus = IopValidateGptHeader(diskIndex, diskSectors, 1u,
                                                 diskSectors - 1u, &primary);
    KSTATUS backupStatus = IopValidateGptHeader(diskIndex, diskSectors, diskSectors - 1u,
                                                1u, &backup);
    bool primaryValid = KSUCCESS(primaryStatus);
    bool backupValid = KSUCCESS(backupStatus);
    if (!primaryValid && !backupValid) {
        DbgPrint("partition: disk %u GPT primary=0x%x backup=0x%x\n",
                 diskIndex, primaryStatus, backupStatus);
        return primaryStatus == NTLSTATUS_CHECKSUM_MISMATCH || backupStatus == NTLSTATUS_CHECKSUM_MISMATCH
             ? NTLSTATUS_CHECKSUM_MISMATCH
             : NTLSTATUS_INVALID_PARAMETER;
    }
    if (primaryValid && backupValid && !IopGptHeadersAgree(&primary, &backup)) {
        DbgPrint("partition: disk %u GPT primary/backup metadata disagree\n", diskIndex);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    const IOP_GPT_VALIDATED_HEADER* selected = primaryValid ? &primary : &backup;
    if (!primaryValid)
        DbgPrint("partition: disk %u using valid backup GPT\n", diskIndex);
    return IopParseGptEntries(diskIndex, diskSectors, selected, partitionsOut);
}

// Discovers either an MBR primary-partition map or a protective-MBR GPT map.
static KSTATUS IopDiscoverDiskPartitions(uint8_t diskIndex,
                                         DEVICE_OBJECT* device,
                                         IOP_PARTITION_RECORD** partitionsOut) {
    if (!device || !partitionsOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *partitionsOut = NULL;

    if (device->u.Disk.SectorSize != IOP_SECTOR_SIZE || device->u.Disk.SectorCount == 0) {
        DbgPrint("partition: disk %u unsupported sector geometry size=%u sectors=%llu\n",
                 diskIndex, device->u.Disk.SectorSize,
                 (unsigned long long)device->u.Disk.SectorCount);
        return NTLSTATUS_NOT_SUPPORTED;
    }
    uint64_t diskSectors = device->u.Disk.SectorCount;

    uint8_t sector[IOP_SECTOR_SIZE];
    KSTATUS status = IoReadDiskSectors(diskIndex, 0, 1, sector);
    if (!KSUCCESS(status))
        return status;
    if (sector[510] != 0x55u || sector[511] != 0xaau)
        return NTLSTATUS_NOT_FOUND;

    uint32_t diskSignature = 0;
    memcpy(&diskSignature, sector + 440u, sizeof(diskSignature));
    IOP_MBR_PARTITION_ENTRY entries[4];
    memcpy(entries, sector + 446u, sizeof(entries));

    bool protective = false;
    for (uint32_t i = 0; i < 4; ++i) {
        if (entries[i].Type == 0xeeu && entries[i].SectorCount != 0)
            protective = true;
    }
    if (protective)
        return IopDiscoverGptDisk(diskIndex, diskSectors, partitionsOut);

    // Validate the complete primary table before emitting supported partitions. Extended containers remain
    // unsupported, but malformed container extents must not make a partially trusted MBR map acceptable.
    for (uint32_t i = 0; i < 4; ++i) {
        const IOP_MBR_PARTITION_ENTRY* entry = &entries[i];
        if (entry->BootIndicator != 0x00u && entry->BootIndicator != 0x80u)
            return NTLSTATUS_INVALID_PARAMETER;
        if (entry->Type == 0 && entry->SectorCount == 0)
            continue;
        if (entry->Type == 0 || entry->SectorCount == 0 ||
            entry->StartingLba == 0 || entry->StartingLba >= diskSectors ||
            (uint64_t)entry->SectorCount > diskSectors - entry->StartingLba) {
            return NTLSTATUS_INVALID_PARAMETER;
        }

        uint64_t entryEnd = (uint64_t)entry->StartingLba + entry->SectorCount;
        for (uint32_t previous = 0; previous < i; ++previous) {
            const IOP_MBR_PARTITION_ENTRY* prior = &entries[previous];
            if (prior->Type == 0 && prior->SectorCount == 0)
                continue;
            uint64_t priorEnd = (uint64_t)prior->StartingLba + prior->SectorCount;
            if ((uint64_t)entry->StartingLba < priorEnd &&
                (uint64_t)prior->StartingLba < entryEnd) {
                return NTLSTATUS_INVALID_PARAMETER;
            }
        }
    }

    IOP_PARTITION_RECORD* partitions = NULL;
    for (uint32_t i = 0; i < 4; ++i) {
        const IOP_MBR_PARTITION_ENTRY* entry = &entries[i];
        if (entry->Type == 0 || IopMbrTypeIsExtended(entry->Type))
            continue;

        BOOT_VOLUME_SELECTOR identity = {0};
        identity.Kind = BootVolumeSelectorMbr;
        identity.Identity.Mbr.DiskSignature = diskSignature;
        identity.Identity.Mbr.PartitionNumber = i + 1u;
        IOP_PARTITION_RECORD* record = IopAllocatePartition(diskIndex,
                                                             entry->StartingLba,
                                                             entry->SectorCount,
                                                             &identity);
        if (!record) {
            IopFreePartitionList(partitions);
            return NTLSTATUS_NO_MEMORY;
        }
        status = IopInsertDiscoveredPartition(&partitions, record);
        if (!KSUCCESS(status)) {
            MmFreePool(record);
            IopFreePartitionList(partitions);
            return status;
        }
    }

    *partitionsOut = partitions;
    return NTLSTATUS_SUCCESS;
}

// Returns whether the registry already contains at least one partition for this disk index.
static bool IopDiskAlreadyDiscovered(uint8_t diskIndex) {
    bool found = false;
    if (!IopAcquirePartitionMutex())
        return false;
    for (IOP_PARTITION_RECORD* record = gPartitions; record; record = record->Next) {
        if (record->Info.DiskIndex == diskIndex) {
            found = true;
            break;
        }
    }
    KeReleaseMutex(gPartitionMutex);
    return found;
}

// Appends one discovered list to the live registry and synchronously notifies every subscriber.
static KSTATUS IopPublishPartitions(uint8_t diskIndex, IOP_PARTITION_RECORD* discovered) {
    if (!IopAcquirePartitionMutex()) {
        IopFreePartitionList(discovered);
        return NTLSTATUS_FAILURE;
    }

    for (IOP_PARTITION_RECORD* record = gPartitions; record; record = record->Next) {
        if (record->Info.DiskIndex == diskIndex) {
            KeReleaseMutex(gPartitionMutex);
            IopFreePartitionList(discovered);
            return NTLSTATUS_SUCCESS;
        }
    }

    IOP_PARTITION_RECORD** tail = &gPartitions;
    while (*tail)
        tail = &(*tail)->Next;
    *tail = discovered;

    for (IOP_PARTITION_RECORD* record = discovered; record; record = record->Next) {
        if (record->Identity.Kind == BootVolumeSelectorMbr) {
            DbgPrint("partition: disk=%u mbr/%08X/%u start=%llu sectors=%llu\n",
                     record->Info.DiskIndex,
                     record->Identity.Identity.Mbr.DiskSignature,
                     record->Identity.Identity.Mbr.PartitionNumber,
                     (unsigned long long)record->Info.StartLba,
                     (unsigned long long)record->Info.SectorCount);
        } else {
            const BOOT_VOLUME_GUID* guid = &record->Identity.Identity.Gpt.PartitionGuid;
            DbgPrint("partition: disk=%u gpt/%08X-%04X-%04X-%02X%02X-%02X%02X%02X%02X%02X%02X start=%llu sectors=%llu\n",
                     record->Info.DiskIndex,
                     guid->Data1, guid->Data2, guid->Data3,
                     guid->Data4[0], guid->Data4[1], guid->Data4[2], guid->Data4[3],
                     guid->Data4[4], guid->Data4[5], guid->Data4[6], guid->Data4[7],
                     (unsigned long long)record->Info.StartLba,
                     (unsigned long long)record->Info.SectorCount);
        }
        if (IopIdentityEquals(&record->Identity, &gSystemVolumeSelector)) {
            DbgPrint("partition: configured system_volume matches disk=%u start=%llu sectors=%llu\n",
                     record->Info.DiskIndex,
                     (unsigned long long)record->Info.StartLba,
                     (unsigned long long)record->Info.SectorCount);
        }
        for (IOP_PARTITION_CALLBACK_ENTRY* entry = gPartitionCallbacks; entry; entry = entry->Next)
            entry->Callback(&record->Info, IoPartitionArrived, entry->Context);
    }

    KeReleaseMutex(gPartitionMutex);
    return NTLSTATUS_SUCCESS;
}

// Discovers and publishes all partitions for one newly registered disk.
static void IopDiskArrived(uint8_t diskIndex, DEVICE_OBJECT* device) {
    if (!device || IopDiskAlreadyDiscovered(diskIndex))
        return;

    IOP_PARTITION_RECORD* discovered = NULL;
    KSTATUS status = IopDiscoverDiskPartitions(diskIndex, device, &discovered);
    if (!KSUCCESS(status)) {
        if (status != NTLSTATUS_NOT_FOUND) {
            DbgPrint("partition: disk %u discovery failed status=0x%x\n", diskIndex, status);
        } else {
            DbgPrint("partition: disk %u has no partition table\n", diskIndex);
        }
        return;
    }
    if (!discovered) {
        DbgPrint("partition: disk %u contains no supported partitions\n", diskIndex);
        return;
    }

    status = IopPublishPartitions(diskIndex, discovered);
    if (!KSUCCESS(status))
        DbgPrint("partition: disk %u publish failed status=0x%x\n", diskIndex, status);
}

// Removes all partitions belonging to one disk and notifies consumers before freeing their snapshots.
static void IopDiskRemoved(uint8_t diskIndex, IO_PARTITION_CHANGE change) {
    if (!IopAcquirePartitionMutex())
        return;

    IOP_PARTITION_RECORD** link = &gPartitions;
    while (*link) {
        IOP_PARTITION_RECORD* record = *link;
        if (record->Info.DiskIndex != diskIndex) {
            link = &record->Next;
            continue;
        }

        *link = record->Next;
        for (IOP_PARTITION_CALLBACK_ENTRY* entry = gPartitionCallbacks; entry; entry = entry->Next)
            entry->Callback(&record->Info, change, entry->Context);
        MmFreePool(record);
    }

    KeReleaseMutex(gPartitionMutex);
}

// Receives the underlying disk namespace and translates it into partition changes.
static void IopDiskChange(uint8_t diskIndex,
                          DEVICE_OBJECT* device,
                          IO_DEVICE_CHANGE change,
                          void* context) {
    (void)context;
    switch (change) {
        case IoDeviceArrived:
            IopDiskArrived(diskIndex, device);
            break;
        case IoDeviceSurpriseRemoved:
        case IoDeviceRemoved:
            IopDiskRemoved(diskIndex,
                           change == IoDeviceSurpriseRemoved
                               ? IoPartitionSurpriseRemoved
                               : IoPartitionRemoved);
            break;
        case IoDeviceQueryRemove:
        default:
            break;
    }
}

// Replays one disk that predates partition-manager initialization.
static void IopReplayDisk(uint8_t diskIndex, DEVICE_OBJECT* device, void* context) {
    (void)context;
    IopDiskArrived(diskIndex, device);
}

// Initializes dynamic partition discovery and the mandatory structural system-volume selector.
KSTATUS IoInitializePartitionManager(const char* systemVolumeSelector) {
    if (!systemVolumeSelector || !*systemVolumeSelector)
        return NTLSTATUS_INVALID_PARAMETER;
    if (gPartitionManagerInitialized)
        return NTLSTATUS_ALREADY_EXISTS;

    BOOT_VOLUME_SELECTOR selector = {0};
    if (!BootVolumeSelectorParse(systemVolumeSelector, strlen(systemVolumeSelector), &selector))
        return NTLSTATUS_INVALID_PARAMETER;

    MUTEX* mutex = NULL;
    KSTATUS status = KeMutexCreate(&mutex);
    if (!KSUCCESS(status))
        return status;

    gPartitionMutex = mutex;
    gSystemVolumeSelector = selector;

    status = IoRegisterDiskChangeCallback(IopDiskChange, NULL);
    if (!KSUCCESS(status)) {
        gPartitionMutex = NULL;
        memset(&gSystemVolumeSelector, 0, sizeof(gSystemVolumeSelector));
        ObDereferenceObject(&mutex->Header);
        return status;
    }

    // The disk-change subscription is permanent, so the manager becomes live before replaying existing disks.
    gPartitionManagerInitialized = true;
    status = IoEnumerateDiskDevices(IopReplayDisk, NULL);
    if (!KSUCCESS(status))
        return status;

    DbgPrint("partition: dynamic discovery initialized for system_volume=%s\n", systemVolumeSelector);
    return NTLSTATUS_SUCCESS;
}

// Registers one partition-change consumer and replays every currently discovered partition before returning.
__declspec(dllexport) KSTATUS IoRegisterPartitionChangeCallback(IO_PARTITION_CHANGE_CALLBACK callback,
                                                                 void* context) {
    if (!callback || !gPartitionManagerInitialized)
        return NTLSTATUS_INVALID_PARAMETER;

    IOP_PARTITION_CALLBACK_ENTRY* entry = MmAllocatePool(sizeof(*entry));
    if (!entry)
        return NTLSTATUS_NO_MEMORY;
    entry->Next = NULL;
    entry->Callback = callback;
    entry->Context = context;

    if (!IopAcquirePartitionMutex()) {
        MmFreePool(entry);
        return NTLSTATUS_FAILURE;
    }

    IOP_PARTITION_CALLBACK_ENTRY** tail = &gPartitionCallbacks;
    while (*tail)
        tail = &(*tail)->Next;
    *tail = entry;

    for (IOP_PARTITION_RECORD* record = gPartitions; record; record = record->Next)
        callback(&record->Info, IoPartitionArrived, context);

    KeReleaseMutex(gPartitionMutex);
    return NTLSTATUS_SUCCESS;
}

// Returns the uniquely selected system partition, or NOT_FOUND/COLLISION while discovery is incomplete or ambiguous.
KSTATUS IoQuerySystemPartition(IO_PARTITION_INFO* partitionOut) {
    if (!partitionOut || !gPartitionManagerInitialized)
        return NTLSTATUS_INVALID_PARAMETER;

    if (!IopAcquirePartitionMutex())
        return NTLSTATUS_FAILURE;

    const IOP_PARTITION_RECORD* match = NULL;
    uint32_t matchCount = 0;
    for (IOP_PARTITION_RECORD* record = gPartitions; record; record = record->Next) {
        if (!IopIdentityEquals(&record->Identity, &gSystemVolumeSelector))
            continue;
        match = record;
        matchCount++;
    }

    KSTATUS status;
    if (matchCount == 0) {
        status = NTLSTATUS_NOT_FOUND;
    } else if (matchCount != 1) {
        status = NTLSTATUS_OBJECT_NAME_COLLISION;
    } else {
        *partitionOut = match->Info;
        status = NTLSTATUS_SUCCESS;
    }

    KeReleaseMutex(gPartitionMutex);
    return status;
}