// Driver module identity and versioned declarative module exports.

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddk/module.h>
#include <kernel/driver_module.h>
#include <kernel/ldr_image.h>
#include <kernel/mm_virt.h>
#include <kernel/io_pnp.h>
#include <kernel/ke_debug.h>

#define DM_NAME_MAX 64u

typedef enum DM_MODULE_STATE {
    DmModuleStarting,
    DmModuleStarted,
    DmModuleFailed,
    DmModuleStopping,
    DmModuleStopped,
} DM_MODULE_STATE;

struct NTL_MODULE {
    struct NTL_MODULE* Next;
    char Name[DM_NAME_MAX];
    uintptr_t ImageBase;
    DM_MODULE_STATE State;
    NTL_MODULE_PREPARE_SHUTDOWN PrepareShutdown;
    NTL_MODULE_COMMIT_SHUTDOWN CommitShutdown;
};

static NTL_MODULE* gDriverModules;

static bool DmImageContains(uintptr_t base, size_t imageSize, uintptr_t ptr, size_t bytes) {
    return ptr >= base && ptr - base <= imageSize && bytes <= imageSize - (ptr - base);
}

__declspec(dllexport) const char* NtlModuleGetName(const NTL_MODULE* module) {
    return module ? module->Name : (const char*)0;
}

static bool DmValidImageString(uintptr_t base, size_t size, const char* text,
                               size_t maxLength) {
    if (!DmImageContains(base, size, (uintptr_t)text, 1u)) return false;
    for (size_t i = 0; i < maxLength; ++i) {
        if (!DmImageContains(base, size, (uintptr_t)text + i, 1u)) return false;
        if (!text[i])
            return i && SysSvIsUtf8Text(SysSvFromParts(text, i));
    }
    return false;
}

static bool DmValidIoDriver(uintptr_t imageBase, size_t imageSize, const DRIVER_OBJECT* driver) {
    if (!DmImageContains(imageBase, imageSize, (uintptr_t)driver, sizeof(*driver)) ||
        driver->Header.Type != ObjectTypeDriver ||
        driver->Header.LifetimeState != ObjectAlive ||
        !driver->Name.Length ||
        !DmImageContains(imageBase, imageSize, (uintptr_t)driver->Name.Data,
                         driver->Name.Length) ||
        !SysSvIsUtf8Text(driver->Name) ||
        (driver->DriverAddDevice &&
         !DmImageContains(imageBase, imageSize, (uintptr_t)driver->DriverAddDevice, 1u)) ||
        (driver->DriverUnload &&
         !DmImageContains(imageBase, imageSize, (uintptr_t)driver->DriverUnload, 1u)))
        return false;

    for (uint32_t i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; ++i) {
        if (driver->MajorFunction[i] &&
            !DmImageContains(imageBase, imageSize,
                             (uintptr_t)driver->MajorFunction[i], 1u))
            return false;
    }

    if (!driver->IdTable) return true;
    if (!driver->DriverAddDevice ||
        !DmImageContains(imageBase, imageSize, (uintptr_t)driver->IdTable,
                         sizeof(*driver->IdTable)) ||
        !driver->IdTable->Count || driver->IdTable->Count > 32u ||
        !DmImageContains(imageBase, imageSize, (uintptr_t)driver->IdTable->Ids,
                         driver->IdTable->Count * sizeof(IO_DEVICE_ID)))
        return false;

    for (size_t i = 0; i < driver->IdTable->Count; ++i) {
        const IO_DEVICE_ID* id = &driver->IdTable->Ids[i];
        if ((id->BusType != IoBusPci && id->BusType != IoBusUsb &&
             id->BusType != IoBusUnknown) ||
            !id->Id.Length ||
            !DmImageContains(imageBase, imageSize, (uintptr_t)id->Id.Data,
                             id->Id.Length) ||
            !SysSvIsUtf8Text(id->Id))
            return false;
    }
    return true;
}

KSTATUS DmStartLoadedDriver(const char* name, uintptr_t imageBase) {
    if (!name || !*name || !imageBase)
        return NTLSTATUS_INVALID_PARAMETER;

    size_t length = 0;
    while (length < DM_NAME_MAX && name[length])
        length++;
    if (!length || length == DM_NAME_MAX)
        return NTLSTATUS_INVALID_PARAMETER;

    for (NTL_MODULE* it = gDriverModules; it; it = it->Next) {
        if (strcmp(it->Name, name) == 0 || it->ImageBase == imageBase)
            return NTLSTATUS_ALREADY_EXISTS;
    }

    size_t imageSize = LdrGetImageSize(imageBase);
    if (!imageSize)
        return NTLSTATUS_INVALID_IMAGE_FORMAT;

    const NTL_MODULE_EXPORT* manifest =
        (const NTL_MODULE_EXPORT*)LdrGetImageExport(imageBase, NTL_DDK_MODULE_EXPORT_NAME);
    NTL_MODULE_INITIALIZE initialize = (NTL_MODULE_INITIALIZE)0;

    if (!manifest) return NTLSTATUS_INVALID_IMAGE_FORMAT;
    {
        if (!DmImageContains(imageBase, imageSize, (uintptr_t)manifest, sizeof(*manifest)) ||
            manifest->Signature != NTL_DDK_MODULE_SIGNATURE ||
            manifest->Version != NTL_DDK_MODULE_ABI_VERSION ||
            manifest->Size != sizeof(*manifest) || manifest->Flags != 0u ||
            manifest->Reserved != 0u || manifest->Reserved2 != 0u ||
            (manifest->Initialize && !DmImageContains(imageBase, imageSize,
                                                      (uintptr_t)manifest->Initialize, 1u)) ||
            (manifest->PrepareShutdown && !DmImageContains(imageBase, imageSize,
                                                      (uintptr_t)manifest->PrepareShutdown, 1u)) ||
            (manifest->CommitShutdown && !DmImageContains(imageBase, imageSize,
                                                      (uintptr_t)manifest->CommitShutdown, 1u)))
            return NTLSTATUS_INVALID_IMAGE_FORMAT;
        if ((manifest->DeviceCount && !manifest->Devices) ||
            manifest->DeviceCount > 64u ||
            (manifest->DeviceCount &&
             !DmImageContains(imageBase, imageSize, (uintptr_t)manifest->Devices,
                              (size_t)manifest->DeviceCount * sizeof(NTL_DEVICE_DESC))))
            return NTLSTATUS_INVALID_IMAGE_FORMAT;
        for (uint32_t i = 0; i < manifest->DeviceCount; ++i) {
            const NTL_DEVICE_DESC* d = &manifest->Devices[i];
            if (d->Size != sizeof(*d) || d->Flags || !d->Name || !d->Start ||
                !DmValidImageString(imageBase, imageSize, d->Name, 64u) ||
                !DmImageContains(imageBase, imageSize, (uintptr_t)d->Start, 1u) ||
                (d->Stop && !DmImageContains(imageBase, imageSize, (uintptr_t)d->Stop, 1u)) ||
                (d->CancelStart && !DmImageContains(imageBase, imageSize,
                                                    (uintptr_t)d->CancelStart, 1u)) ||
                !d->MatchCount || d->MatchCount > 32u ||
                !DmImageContains(imageBase, imageSize, (uintptr_t)d->Matches,
                                 (size_t)d->MatchCount * sizeof(NTL_DEVICE_MATCH)))
                return NTLSTATUS_INVALID_IMAGE_FORMAT;
            for (uint32_t j = 0; j < d->MatchCount; ++j) {
                const NTL_DEVICE_MATCH* m = &d->Matches[j];
                if ((m->Bus != NtlBusPci && m->Bus != NtlBusUsb) ||
                    !DmValidImageString(imageBase, imageSize, m->Id, 96u))
                    return NTLSTATUS_INVALID_IMAGE_FORMAT;
            }
        }
        if ((manifest->IoDriverCount && !manifest->IoDrivers) ||
            manifest->IoDriverCount > 64u ||
            (manifest->IoDriverCount &&
             !DmImageContains(imageBase, imageSize, (uintptr_t)manifest->IoDrivers,
                              (size_t)manifest->IoDriverCount * sizeof(*manifest->IoDrivers))))
            return NTLSTATUS_INVALID_IMAGE_FORMAT;
        for (uint32_t i = 0; i < manifest->IoDriverCount; ++i) {
            if (!DmImageContains(imageBase, imageSize,
                                 (uintptr_t)manifest->IoDrivers[i], sizeof(DRIVER_OBJECT)))
                return NTLSTATUS_INVALID_IMAGE_FORMAT;
        }
        initialize = manifest->Initialize;
    }

    NTL_MODULE* module = (NTL_MODULE*)MmAllocatePool(sizeof(*module));
    if (!module)
        return NTLSTATUS_NO_MEMORY;

    memset(module, 0, sizeof(*module));
    memcpy(module->Name, name, length);
    module->ImageBase = imageBase;
    module->State = DmModuleStarting;
    module->PrepareShutdown = manifest->PrepareShutdown;
    module->CommitShutdown = manifest->CommitShutdown;
    module->Next = gDriverModules;
    gDriverModules = module;

    KSTATUS status = initialize ? initialize(module) : NTLSTATUS_SUCCESS;
    if (!KSUCCESS(status) && !KFAILED(status))
        status = NTLSTATUS_INVALID_PARAMETER;

    if (KSUCCESS(status) && manifest) {
        for (uint32_t i = 0; i < manifest->IoDriverCount; ++i) {
            if (!DmValidIoDriver(imageBase, imageSize, manifest->IoDrivers[i])) {
                status = NTLSTATUS_INVALID_IMAGE_FORMAT;
                break;
            }
        }
    }
    if (KSUCCESS(status))
        status = IoPnpPublishModule(module, manifest->Devices, manifest->DeviceCount,
                                    manifest->IoDrivers, manifest->IoDriverCount);
    module->State = KSUCCESS(status) ? DmModuleStarted : DmModuleFailed;
    return status;
}

KSTATUS DmPrepareShutdown(void) {
    KSTATUS firstError = NTLSTATUS_SUCCESS;
    for (NTL_MODULE* module = gDriverModules; module; module = module->Next) {
        if (module->State != DmModuleStarted) continue;
        module->State = DmModuleStopping;
        if (!module->PrepareShutdown) continue;
        KSTATUS st = module->PrepareShutdown(module);
        if (!KSUCCESS(st)) {
            DbgPrint("shutdown: prepare %s failed 0x%x\n", module->Name, st);
            if (KSUCCESS(firstError)) firstError = st;
        }
    }
    return firstError;
}

void DmCommitShutdown(void) {
    for (NTL_MODULE* module = gDriverModules; module; module = module->Next) {
        if (module->State != DmModuleStopping) continue;
        if (module->CommitShutdown) module->CommitShutdown(module);
        module->State = DmModuleStopped;
    }
}