// DriverEntry.cpp - KMDF 驱动入口
// 职责：初始化 WDF 驱动对象、注册 EvtDeviceAdd（PnP 过滤驱动）
//      创建控制设备 \\.\CommModifyKit 接收 DLL 的 IOCTL
// 注意：本驱动作为串口设备类的 UpperFilters 加载
//      当 PnP Manager 创建串口 PDO 时，调用 EvtDeviceAdd attach 过滤设备
#include <ntddk.h>
#include <wdf.h>
#include "DeviceContext.h"
#include "QueueCallback.h"
#include "ClientConnection.h"
#include "SerialFilter.h"

extern "C" {

DRIVER_INITIALIZE DriverEntry;
EVT_WDF_DRIVER_DEVICE_ADD EvtDeviceAdd;
EVT_WDF_DRIVER_UNLOAD EvtDriverUnload;
EVT_WDF_OBJECT_CONTEXT_CLEANUP EvtDeviceContextCleanup;

} // extern "C"

// PnP 回调使用 C++ 链接（WDF 宏内部处理）
EVT_WDF_DEVICE_PREPARE_HARDWARE EvtDevicePrepareHardware;

namespace commkit_driver {

// 全局驱动对象（仅用于内部日志/调试）
static WDFDRIVER g_wdf_driver = nullptr;

} // namespace commkit_driver

extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT driver_object,
                                  PUNICODE_STRING registry_path) {
    NTSTATUS status;
    WDF_DRIVER_CONFIG cfg;

    // 初始化全局客户端连接管理器
    commkit_driver::GetClientConnection().Initialize();

    // 初始化 WDF 配置
    WDF_DRIVER_CONFIG_INIT(&cfg, EvtDeviceAdd);
    cfg.EvtDriverUnload = EvtDriverUnload;

    // 创建 WDF 驱动对象
    status = WdfDriverCreate(driver_object, registry_path,
                                WDF_NO_OBJECT_ATTRIBUTES, &cfg,
                                &commkit_driver::g_wdf_driver);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 初始化 ClientConnection（创建控制设备）
    status = commkit_driver::CreateControlDevice(commkit_driver::g_wdf_driver);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    return STATUS_SUCCESS;
}

extern "C" NTSTATUS EvtDeviceAdd(WDFDRIVER driver, PWDFDEVICE_INIT init) {
    UNREFERENCED_PARAMETER(driver);

    NTSTATUS status;

    // 1. 标记为过滤设备
    WdfFdoInitSetFilter(init);

    // 2. 注册 PnP 电源回调（EvtDevicePrepareHardware 解析 COM 编号）
    WDF_PNPPOWER_EVENT_CALLBACKS pnp_callbacks;
    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnp_callbacks);
    pnp_callbacks.EvtDevicePrepareHardware = EvtDevicePrepareHardware;
    WdfDeviceInitSetPnpPowerEventCallbacks(init, &pnp_callbacks);

    // 3. 注册 WDM IRP 预处理回调（拦截 READ/WRITE/DEVICE_CONTROL）
    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchRead,
        IRP_MJ_READ, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchWrite,
        IRP_MJ_WRITE, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchDeviceControl,
        IRP_MJ_DEVICE_CONTROL, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 4. 创建 WDF 设备
    WDFDEVICE device;
    WDF_OBJECT_ATTRIBUTES attrs;
    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attrs, DEVICE_CONTEXT);
    attrs.EvtCleanupCallback = EvtDeviceContextCleanup;

    status = WdfDeviceCreate(&init, &attrs, &device);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 5. 初始化设备上下文
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    ctx->ComNumber = 0;          // 延迟到 EvtDevicePrepareHardware 解析
    ctx->MonitoringEnabled = FALSE;
    ctx->RingBuffer = nullptr;
    ctx->WdfDevice = device;
    ctx->LowerDevice = WdfDeviceWdmGetAttachedDevice(device);

    // 6. 注册到全局端口表（com_number 暂为 0，PrepareHardware 时更新）
    commkit_driver::GetClientConnection().RegisterFilterDevice(device);

    return STATUS_SUCCESS;
}

NTSTATUS EvtDevicePrepareHardware(WDFDEVICE device,
                                     WDFCMRESLIST resources,
                                     WDFCMRESLIST resources_translated) {
    UNREFERENCED_PARAMETER(resources);
    UNREFERENCED_PARAMETER(resources_translated);

    NTSTATUS status;
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);

    // 从 PnP 属性解析 COM 编号
    // 方式：获取 DevicePropertyFriendlyName 或 DevicePropertyDeviceDescription，
    // 匹配 "COM%d" 模式
    WDFMEMORY memory;
    status = WdfDeviceAllocAndQueryProperty(
        device, DevicePropertyFriendlyName, NonPagedPoolNx, 0, &memory);
    if (!NT_SUCCESS(status)) {
        // 尝试 DeviceDescription 作为后备
        status = WdfDeviceAllocAndQueryProperty(
            device, DevicePropertyDeviceDescription, NonPagedPoolNx, 0, &memory);
    }

    if (NT_SUCCESS(status)) {
        PWCHAR name = (PWCHAR)WdfMemoryGetBuffer(memory, nullptr);
        if (name) {
            // 解析 "COMx" 模式（可能出现在 "COMx (port)" 或 "通信端口 (COMx)" 等格式中）
            ULONG name_len = (ULONG)wcslen(name);
            for (ULONG i = 0; i + 3 <= name_len; i++) {
                if ((name[i] == L'C' || name[i] == L'c') &&
                    (name[i+1] == L'O' || name[i+1] == L'o') &&
                    (name[i+2] == L'M' || name[i+2] == L'm')) {
                    if (name[i+3] >= L'0' && name[i+3] <= L'9') {
                        ULONG num = 0;
                        ULONG j = i + 3;
                        while (j < name_len && name[j] >= L'0' && name[j] <= L'9') {
                            num = num * 10 + (name[j] - L'0');
                            j++;
                        }
                        if (num > 0 && num <= 256) {
                            ctx->ComNumber = num;
                            break;
                        }
                    }
                }
            }
        }
        WdfObjectDelete(memory);
    }

    // 后备方案：FriendlyName/DeviceDescription 均未解析出 COM 编号时，
    // 从 PDO 名称查询 HKLM\HARDWARE\DEVICEMAP\SERIALCOMM 获取 "COMx" 映射
    if (ctx->ComNumber == 0) {
        WDFMEMORY pdo_memory;
        status = WdfDeviceAllocAndQueryProperty(
            device, DevicePropertyPhysicalDeviceObjectName,
            NonPagedPoolNx, 0, &pdo_memory);
        if (NT_SUCCESS(status)) {
            PWCHAR pdo_name = (PWCHAR)WdfMemoryGetBuffer(pdo_memory, nullptr);
            if (pdo_name && pdo_name[0] != L'\0') {
                UNICODE_STRING key_path;
                RtlInitUnicodeString(&key_path,
                    L"\\Registry\\Machine\\HARDWARE\\DEVICEMAP\\SERIALCOMM");
                OBJECT_ATTRIBUTES obj_attr;
                InitializeObjectAttributes(&obj_attr, &key_path,
                    OBJ_CASE_INSENSITIVE, nullptr, nullptr);

                HANDLE hkey;
                status = ZwOpenKey(&hkey, KEY_READ, &obj_attr);
                if (NT_SUCCESS(status)) {
                    UNICODE_STRING value_name;
                    RtlInitUnicodeString(&value_name, pdo_name);

                    UCHAR result_buf[256] = {0};
                    ULONG result_len = 0;
                    status = ZwQueryValueKey(hkey, &value_name,
                        KeyValuePartialInformation,
                        result_buf, sizeof(result_buf), &result_len);
                    if (NT_SUCCESS(status)) {
                        KEY_VALUE_PARTIAL_INFORMATION* kv_info =
                            (KEY_VALUE_PARTIAL_INFORMATION*)result_buf;
                        if (kv_info->Type == REG_SZ && kv_info->DataLength >= sizeof(WCHAR)) {
                            PWCHAR com_name = (PWCHAR)kv_info->Data;
                            ULONG com_len = kv_info->DataLength / sizeof(WCHAR);
                            // 解析 "COMx"
                            if (com_len >= 4 &&
                                (com_name[0] == L'C' || com_name[0] == L'c') &&
                                (com_name[1] == L'O' || com_name[1] == L'o') &&
                                (com_name[2] == L'M' || com_name[2] == L'm') &&
                                com_name[3] >= L'0' && com_name[3] <= L'9') {
                                ULONG num = 0;
                                ULONG j = 3;
                                while (j < com_len &&
                                       com_name[j] >= L'0' && com_name[j] <= L'9') {
                                    num = num * 10 + (com_name[j] - L'0');
                                    j++;
                                }
                                if (num > 0 && num <= 256) {
                                    ctx->ComNumber = num;
                                }
                            }
                        }
                    }
                    ZwClose(hkey);
                }
            }
            WdfObjectDelete(pdo_memory);
        }
    }

    // 更新端口表中的 COM 编号
    commkit_driver::GetClientConnection().UpdatePortComNumber(device, ctx->ComNumber);

    return STATUS_SUCCESS;
}

// 设备上下文清理回调：设备移除时注销端口表条目
VOID EvtDeviceContextCleanup(WDFOBJECT object) {
    WDFDEVICE device = (WDFDEVICE)object;
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    if (ctx && ctx->ComNumber != 0) {
        commkit_driver::GetClientConnection().UnregisterFilterDevice(ctx->ComNumber);
    }
}

extern "C" VOID EvtDriverUnload(WDFDRIVER driver) {
    UNREFERENCED_PARAMETER(driver);
    commkit_driver::GetClientConnection().Cleanup();
    commkit_driver::g_wdf_driver = nullptr;
}