#pragma once
// io_driver.h - kernel driver/device I/O model
//
// Device lifetime is owned by the I/O manager. Bus drivers create PDOs,
// function/filter drivers attach FDOs, and PnP/remove IRPs tear stacks down.
// Disk indexes are only a compatibility namespace above DEVICE_OBJECTs.

#include <stdint.h>
#include <stdbool.h>
#include <stddef.h>
#include <ntlite/status.h>
#include <ntlite/string_view.h>
#include <kernel/ob_object.h>
#include <kernel/ke_list.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irp.h>

typedef struct HAL_API HAL_API;
typedef KSTATUS (*DRIVER_ENTRY_FN)(const HAL_API* hal);
typedef void    (*DRIVER_UNLOAD_FN)(void);

typedef struct DRIVER_OBJECT DRIVER_OBJECT;
typedef struct DEVICE_OBJECT DEVICE_OBJECT;
typedef struct IRP IRP;

typedef KSTATUS (*PDRIVER_DISPATCH)(DEVICE_OBJECT* device, IRP* irp);
typedef KSTATUS (*PDRIVER_ADD_DEVICE)(DRIVER_OBJECT* driver,
                                      DEVICE_OBJECT* physicalDeviceObject);

// -------------------------------------------------------------------------
// Bus / identity
// -------------------------------------------------------------------------

typedef enum IO_BUS_TYPE {
    IoBusUnknown = 0,
    IoBusPci     = 1,
    IoBusUsb     = 2,
} IO_BUS_TYPE;

typedef struct IO_DEVICE_ID {
    IO_BUS_TYPE     BusType;
    SYS_STRING_VIEW Id;
} IO_DEVICE_ID;

typedef struct IO_DRIVER_ID_TABLE {
    const IO_DEVICE_ID* Ids;
    size_t              Count;
} IO_DRIVER_ID_TABLE;

#define IO_MAX_DEVICE_IDS 8u
#define IO_MAX_DEVICE_ID_CHARS 96u
#define IO_MAX_DEVICE_NAME_CHARS 64u

// -------------------------------------------------------------------------
// Driver/device objects
// -------------------------------------------------------------------------

typedef enum DEVICE_OBJECT_TYPE {
    DeviceObjectPhysical = 1,
    DeviceObjectFunction = 2,
    DeviceObjectFilter   = 3,
} DEVICE_OBJECT_TYPE;

typedef enum DEVICE_PNP_STATE {
    DevicePnpNotStarted = 0,
    DevicePnpStarted,
    DevicePnpStopPending,
    DevicePnpStopped,
    DevicePnpRemovePending,
    DevicePnpSurpriseRemoved,
    DevicePnpDeleted,
} DEVICE_PNP_STATE;

struct DRIVER_OBJECT {
    OBJECT              Header;
    SYS_STRING_VIEW     Name;
    PDRIVER_DISPATCH    MajorFunction[IRP_MJ_MAXIMUM_FUNCTION];
    PDRIVER_ADD_DEVICE  DriverAddDevice;
    DRIVER_UNLOAD_FN    DriverUnload;
    const IO_DRIVER_ID_TABLE* IdTable;
};

struct DEVICE_OBJECT {
    OBJECT              Header;

    DRIVER_OBJECT*      DriverObject;
    DEVICE_OBJECT_TYPE  Type;
    DEVICE_PNP_STATE    PnpState;

    DEVICE_OBJECT*      Parent;
    DEVICE_OBJECT*      FirstChild;
    DEVICE_OBJECT*      NextSibling;

    DEVICE_OBJECT*      AttachedDevice; // device directly above this object
    DEVICE_OBJECT*      LowerDevice;    // device directly below this FDO/filter

    void*               DeviceExtension;

    uint32_t            Flags;
    uint32_t            Characteristics;
    uint32_t            AlignmentRequirement;

    SPINLOCK            Lock;
    volatile long       OutstandingIrps;
    bool                DeletePending;

    SYS_STRING_VIEW     Name;
    char                NameStorage[IO_MAX_DEVICE_NAME_CHARS];

    union {
        struct {
            IO_BUS_TYPE     BusType;
            void*           BusContext;
            void*           DeviceContext;
            IO_DEVICE_ID    Ids[IO_MAX_DEVICE_IDS];
            char            IdStorage[IO_MAX_DEVICE_IDS][IO_MAX_DEVICE_ID_CHARS];
            uint8_t         IdCount;
        } Pdo;

        struct {
            uint8_t         DiskIndex;
            uint64_t        SectorCount;
            uint32_t        SectorSize;
        } Disk;
    } u;
};

#define DEVICE_OBJECT_DISK       0x00000001u
#define DEVICE_OBJECT_REMOVABLE  0x00000002u

// -------------------------------------------------------------------------
// Object/device creation and stack management
// -------------------------------------------------------------------------

__declspec(dllexport) void    IoInitializeDriverObject(DRIVER_OBJECT* driver,
                                                       SYS_STRING_VIEW name);
__declspec(dllexport) KSTATUS IoRegisterDriverObject(DRIVER_OBJECT* driver);

__declspec(dllexport) KSTATUS IoCreateDevice(DRIVER_OBJECT* driver,
                                             size_t extensionSize,
                                             DEVICE_OBJECT** outDevice);

__declspec(dllexport) KSTATUS IoCreatePhysicalDevice(DRIVER_OBJECT* busDriver,
                                                     DEVICE_OBJECT* parent,
                                                     const IO_DEVICE_ID* ids,
                                                     size_t idCount,
                                                     size_t extensionSize,
                                                     DEVICE_OBJECT** outPdo);
__declspec(dllexport) KSTATUS IoRegisterPhysicalDevice(DEVICE_OBJECT* pdo);

__declspec(dllexport) KSTATUS IoSetDeviceName(DEVICE_OBJECT* device,
                                                 SYS_STRING_VIEW name);
__declspec(dllexport) void    IoReferenceDeviceObject(DEVICE_OBJECT* device);
__declspec(dllexport) void    IoDereferenceDeviceObject(DEVICE_OBJECT* device);
__declspec(dllexport) void    IoDeleteDevice(DEVICE_OBJECT* device);

__declspec(dllexport) DEVICE_OBJECT* IoAttachDeviceToDeviceStack(DEVICE_OBJECT* source,
                                                                 DEVICE_OBJECT* target);
__declspec(dllexport) void    IoDetachDevice(DEVICE_OBJECT* source);
__declspec(dllexport) DEVICE_OBJECT* IoGetAttachedDeviceReference(DEVICE_OBJECT* device);

__declspec(dllexport) KSTATUS IoCallDriver(DEVICE_OBJECT* device, IRP* irp);
__declspec(dllexport) KSTATUS IoForwardIrpToLower(DEVICE_OBJECT* device, IRP* irp);
__declspec(dllexport) uint8_t IoGetDeviceStackSize(DEVICE_OBJECT* device);
__declspec(dllexport) KSTATUS IoSendPnpIrpSync(DEVICE_OBJECT* device,
                                                IO_PNP_MINOR_FUNCTION minor);

// -------------------------------------------------------------------------
// Legacy USB publish/matcher path. This remains only until USB is migrated to
// PDO/interface-PDO enumeration in the later USB stages.
// -------------------------------------------------------------------------

#define IO_PUBLISHED_USB_MAX_ENDPOINTS 8u

typedef struct IO_PUBLISHED_USB_ENDPOINT {
    uint8_t  Address;
    uint8_t  Attributes;
    uint16_t MaxPacketSize;
    uint8_t  Interval;
} IO_PUBLISHED_USB_ENDPOINT;

typedef struct IO_PUBLISHED_USB_DEVICE {
    uint16_t VendorId;
    uint16_t ProductId;
    uint16_t BcdUsb;
    uint16_t BcdDevice;

    uint8_t DeviceClass;
    uint8_t DeviceSubclass;
    uint8_t DeviceProtocol;

    uint8_t SlotId;
    uint8_t Port;
    uint8_t Speed;
    uint8_t ConfigurationValue;

    uint8_t InterfaceNumber;
    uint8_t AlternateSetting;
    uint8_t InterfaceClass;
    uint8_t InterfaceSubclass;
    uint8_t InterfaceProtocol;

    uint8_t EndpointCount;
    IO_PUBLISHED_USB_ENDPOINT Endpoints[IO_PUBLISHED_USB_MAX_ENDPOINTS];
} IO_PUBLISHED_USB_DEVICE;

typedef struct IO_PUBLISHED_DEVICE {
    IO_BUS_TYPE BusType;
    const char* Name;
    void*       BusContext;
    void*       DeviceContext;
    uint32_t    Flags;
    bool        Bound;
    const char* BoundDriverName;

    union {
        IO_PUBLISHED_USB_DEVICE Usb;
    } u;
} IO_PUBLISHED_DEVICE;

typedef bool (*IO_DRIVER_MATCH_CALLBACK)(const IO_PUBLISHED_DEVICE* device,
                                         void* context);
typedef KSTATUS (*IO_DRIVER_BIND_CALLBACK)(IO_PUBLISHED_DEVICE* device,
                                           void* context);

__declspec(dllexport) KSTATUS IoRegisterDriverMatcher(const char* driverName,
                                                      IO_BUS_TYPE busType,
                                                      IO_DRIVER_MATCH_CALLBACK match,
                                                      IO_DRIVER_BIND_CALLBACK bind,
                                                      void* context);
__declspec(dllexport) KSTATUS IoPublishDevice(const IO_PUBLISHED_DEVICE* device);
__declspec(dllexport) uint32_t IoPublishedDeviceCount(void);

// -------------------------------------------------------------------------
// Disk namespace above storage DEVICE_OBJECTs
// -------------------------------------------------------------------------

typedef enum IO_DEVICE_CHANGE {
    IoDeviceArrived = 0,
    IoDeviceQueryRemove,
    IoDeviceSurpriseRemoved,
    IoDeviceRemoved,
} IO_DEVICE_CHANGE;

typedef void (*IO_DISK_CHANGE_CALLBACK)(uint8_t driveIndex,
                                        DEVICE_OBJECT* device,
                                        IO_DEVICE_CHANGE change,
                                        void* context);

typedef void (*IO_DISK_ARRIVAL_CALLBACK)(uint8_t driveIndex,
                                          DEVICE_OBJECT* device,
                                          void* context);

__declspec(dllexport) KSTATUS IoRegisterDiskDevice(uint8_t driveIndex,
                                                   DEVICE_OBJECT* device);
__declspec(dllexport) KSTATUS IoNotifyDiskSurpriseRemoved(uint8_t driveIndex);
__declspec(dllexport) KSTATUS IoUnregisterDiskDevice(uint8_t driveIndex);
__declspec(dllexport) DEVICE_OBJECT* IoGetDiskDevice(uint8_t driveIndex);
__declspec(dllexport) DEVICE_OBJECT* IoReferenceDiskDevice(uint8_t driveIndex);
__declspec(dllexport) bool    IoDiskDevicePresent(uint8_t driveIndex);
__declspec(dllexport) uint64_t IoDiskSectorCount(uint8_t driveIndex);

__declspec(dllexport) KSTATUS IoRegisterDiskArrivalCallback(IO_DISK_ARRIVAL_CALLBACK callback,
                                                            void* context);
__declspec(dllexport) KSTATUS IoRegisterDiskChangeCallback(IO_DISK_CHANGE_CALLBACK callback,
                                                           void* context);
__declspec(dllexport) KSTATUS IoEnumerateDiskDevices(IO_DISK_ARRIVAL_CALLBACK callback,
                                                     void* context);

__declspec(dllexport) KSTATUS IoReadDiskSectors(uint8_t driveIndex,
                                                uint64_t lba,
                                                uint32_t count,
                                                void* buffer);
__declspec(dllexport) KSTATUS IoWriteDiskSectors(uint8_t driveIndex,
                                                 uint64_t lba,
                                                 uint32_t count,
                                                 const void* buffer);