#include <ntddk.h>
#include <wdf.h>
#include "Ioctl.hpp"
#include "kmdfrc.hpp"
#include "staticclassessupport.h"
#include "trace.hpp"
#include "kmdfrc.tmh"

EXTERN_C_START

DRIVER_INITIALIZE DriverEntry;

NTSTATUS
DriverEntry(
    IN OUT PDRIVER_OBJECT   DriverObject,
    IN PUNICODE_STRING      RegistryPath
);

EXTERN_C_END

typedef struct _DEVICE_EXTENSION
{
    int a;
}DEVICE_CONTEXT, *PDEVICE_EXTENSION;
WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(DEVICE_CONTEXT, DeviceGetContext)

VOID
KmdfrcUnload(
    IN WDFDRIVER Driver
)
{
    PAGED_CODE();
    DrvLogInfo("Unload Driver=0x%p", Driver);
    auto driverObject = WdfDriverWdmGetDriverObject(Driver);
    WPP_CLEANUP(driverObject);
}

_Function_class_(EVT_WDF_OBJECT_CONTEXT_CLEANUP)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcContextCleanup(
    _In_ WDFOBJECT Object
)
{
    DrvLogInfo("KmdfrcContextCleanup Object=0x%p", Object);
}

_Function_class_(EVT_WDF_OBJECT_CONTEXT_DESTROY)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcContextDestroy(
    _In_
    WDFOBJECT Object
)
{
    DrvLogInfo("KmdfrcContextDestroy Object=0x%p", Object);
}

EVT_WDF_DEVICE_FILE_CREATE KmdfrcEvtWdfDeviceFileCreate;
EVT_WDF_FILE_CLOSE KmdfrcEvtWdfFileClose;
EVT_WDF_FILE_CLEANUP KmdfrcEvtWdfFileCleanup;

void
KmdfrcEvtWdfDeviceFileCreate(
    _In_ WDFDEVICE Device,
    _In_ WDFREQUEST Request,
    _In_ WDFFILEOBJECT FileObject
)
{
    DrvLogInfo("KmdfrcEvtWdfDeviceFileCreate Device=0x%p, Request=0x%p, FileObject=0x%p", Device, Request, FileObject);
    WdfRequestComplete(Request, STATUS_SUCCESS);

}

void
KmdfrcEvtWdfFileCleanup(
    _In_ WDFFILEOBJECT FileObject
)
{
    DrvLogInfo("KmdfrcEvtWdfFileCleanup FileObject=0x%p", FileObject);
}

void
KmdfrcEvtWdfFileClose(
    _In_ WDFFILEOBJECT FileObject
)
{
    DrvLogInfo("KmdfrcEvtWdfFileClose FileObject=0x%p", FileObject);
}

_Function_class_(EVT_WDF_IO_QUEUE_IO_READ)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcIoQueueRead(
    _In_ WDFQUEUE Queue,
    _In_ WDFREQUEST Request,
    _In_ size_t Length
)
{
    DrvLogInfo("KmdfrcIoQueueRead Queue=0x%p, Request=0x%p, Length=0x%Iu", Queue, Request, Length);
    WdfRequestComplete(Request, STATUS_ACCESS_DENIED);
}

_Function_class_(EVT_WDF_IO_QUEUE_IO_WRITE)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcIoQueueWrite(
    _In_ WDFQUEUE Queue,
    _In_ WDFREQUEST Request,
    _In_ size_t Length
)
{
    DrvLogInfo("KmdfrcIoQueueWrite Queue=0x%p, Request=0x%p, Length=0x%Iu", Queue, Request, Length);
    WdfRequestComplete(Request, STATUS_ACCESS_DENIED);
}

_Function_class_(EVT_WDF_IO_QUEUE_IO_DEVICE_CONTROL)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcIoQueueDeviceControl(
    _In_ WDFQUEUE Queue,
    _In_ WDFREQUEST Request,
    _In_ size_t OutputBufferLength,
    _In_ size_t InputBufferLength,
    _In_ ULONG IoControlCode
)
{
    DrvLogInfo("KmdfrcIoQueueDeviceControl Queue=0x%p, Request=0x%p, OutputBufferLength=0x%Iu, InputBufferLength=0x%Iu, IoControlCode=0x%x", Queue, Request, OutputBufferLength, InputBufferLength, IoControlCode);
    WdfRequestComplete(Request, STATUS_ACCESS_DENIED);
}

_Function_class_(EVT_WDF_IO_IN_CALLER_CONTEXT)
_IRQL_requires_same_
_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
KmdfrcEvtIoInCallerContext(
    _In_ WDFDEVICE Device,
    _In_ WDFREQUEST Request
)
{
    DrvLogInfo("KmdfrcEvtIoInCallerContext Device=0x%p, Request=0x%p", Device, Request);
}

NTSTATUS
KmdfrcInitDevice(
    _In_ WDFDRIVER Driver,
    _In_ PWDFDEVICE_INIT DeviceInit
)
{
    UNREFERENCED_PARAMETER(Driver);
    NTSTATUS status = STATUS_UNSUCCESSFUL;;
    WDF_OBJECT_ATTRIBUTES attributes = { 0 };
    WDF_IO_QUEUE_CONFIG ioQueueConfig = { 0 };
    WDF_FILEOBJECT_CONFIG fileConfig = { 0 };
    WDFQUEUE queue = nullptr;
    WDFDEVICE controlDevice = nullptr;
    DECLARE_CONST_UNICODE_STRING(ntDeviceName, KMDFRC_DEVICE_NAME_STRING);
    DECLARE_CONST_UNICODE_STRING(symbolicLinkName, KMDFRC_SYMBOLIC_NAME_STRING);

    WdfDeviceInitSetExclusive(DeviceInit, TRUE);
    WdfDeviceInitSetIoType(DeviceInit, WdfDeviceIoBuffered);

    status = WdfDeviceInitAssignName(DeviceInit, &ntDeviceName);
    if (!VERIFY_SUCCESS(status))
    {
        DrvLogError("WdfDeviceInitAssignName failed with status: 0x%x", status);
        goto exit;
    }

    // WdfControlDeviceInitSetShutdownNotification

    WDF_FILEOBJECT_CONFIG_INIT(
        &fileConfig,
        KmdfrcEvtWdfDeviceFileCreate,
        KmdfrcEvtWdfFileClose,
        KmdfrcEvtWdfFileCleanup // not interested in Cleanup
    );

    WdfDeviceInitSetFileObjectConfig(
        DeviceInit,
        &fileConfig,
        WDF_NO_OBJECT_ATTRIBUTES
    );

    //WdfDeviceInitSetIoInCallerContextCallback(
    //    DeviceInit,
    //    KmdfrcEvtIoInCallerContext
    //);

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    WDF_OBJECT_ATTRIBUTES_SET_CONTEXT_TYPE(&attributes, DEVICE_CONTEXT);

    status = WdfDeviceCreate(
        &DeviceInit,
        &attributes,
        &controlDevice
    );
    if (!VERIFY_SUCCESS(status))
    {
        DrvLogError("WdfDeviceCreate failed with status: 0x%x", status);
        goto exit;
    }

    status = WdfDeviceCreateSymbolicLink(
        controlDevice,
        &symbolicLinkName
    );
    if (!VERIFY_SUCCESS(status))
    {
        DrvLogError("WdfDeviceCreateSymbolicLink failed with status: 0x%x", status);
        goto exit;
    }

    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(
        &ioQueueConfig,
        WdfIoQueueDispatchSequential
    );

    ioQueueConfig.EvtIoRead = KmdfrcIoQueueRead;
    ioQueueConfig.EvtIoWrite = KmdfrcIoQueueWrite;
    ioQueueConfig.EvtIoDeviceControl = KmdfrcIoQueueDeviceControl;

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    status = WdfIoQueueCreate(
        controlDevice,
        &ioQueueConfig,
        &attributes,
        &queue // pointer to default queue
    );
    if (!VERIFY_SUCCESS(status))
    {
        DrvLogError("WdfIoQueueCreate failed with status: 0x%x", status);
        goto exit;
    }
    WdfControlFinishInitializing(controlDevice);
exit:
    if (DeviceInit != NULL)
    {
        WdfDeviceInitFree(DeviceInit);
    }

    return status;
}

NTSTATUS
DriverEntry(
    _Inout_ PDRIVER_OBJECT DriverObject,
    _In_ PUNICODE_STRING RegistryPath
)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    WDF_DRIVER_CONFIG config = { 0 };
    WDFDRIVER driverObject = nullptr;
    PWDFDEVICE_INIT deviceInit = nullptr;
    WDF_OBJECT_ATTRIBUTES attributes = { 0 };

    WPP_INIT_TRACING(DriverObject, RegistryPath);

    CppInitStaticDestructorList();
    CppCallStaticConstructors(__crtXca, __crtXcz);
    if (!NT_VERIFY(::CppWereAllDestructorsRegistered()))
    {
        DrvLogCritical("Could not register all static destructors.");
        status = STATUS_INSUFFICIENT_RESOURCES;
        goto exit;
    }

    WDF_DRIVER_CONFIG_INIT(
        &config,
        WDF_NO_EVENT_CALLBACK
    );

    config.DriverInitFlags |= WdfDriverInitNonPnpDriver; // non-pnp driver
    config.EvtDriverUnload = KmdfrcUnload;

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.EvtCleanupCallback = KmdfrcContextCleanup;
    attributes.EvtDestroyCallback = KmdfrcContextDestroy;
    status = WdfDriverCreate(
        DriverObject,
        RegistryPath,
        &attributes,
        &config,
        &driverObject
    );
    if (!NT_SUCCESS(status))
    {
        DrvLogCritical("NonPnp: WdfDriverCreate failed with status 0x%x", status);
        goto exit;
    }

    // manually allocate device init since we
    // don't have a device add callback
    deviceInit = WdfControlDeviceInitAllocate(
        driverObject,
        &SDDL_DEVOBJ_SYS_ALL_ADM_RWX_WORLD_RW_RES_R
    );

    status = KmdfrcInitDevice(driverObject, deviceInit);

exit:
    if (!VERIFY_SUCCESS(status))
    {
        CppExit();
        WPP_CLEANUP(DriverObject);
    }
    return status;
}