// DriverEntry.cpp - KMDF 驱动入口
// 职责：初始化 WDF 驱动对象、注册 EvtDeviceAdd（PnP 过滤驱动）
//      创建控制设备 \\.\CommModifyKit 接收 DLL 的 IOCTL
// 注意：本驱动作为串口设备类的 UpperFilters 加载
//      当 PnP Manager 创建串口 PDO 时，调用 EvtDeviceAdd attach 过滤设备
#include <ntddk.h>
#include <wdf.h>
#include <initguid.h>  // 必须在包含定义 GUID 的头文件之前
#include <ntddser.h>   // GUID_DEVINTERFACE_COMPORT
#include "DeviceContext.h"
#include "QueueCallback.h"
#include "ClientConnection.h"
#include "SerialFilter.h"   // g_sequence_counter
#include "EventRingBuffer.h" // EventRingBuffer
#include "../common/CommKitEvents.h"  // COMMKIT_OP_OPEN / OP_CLOSE

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) {
    DbgPrint("[CommModifyKit] DriverEntry ENTERED\n");
    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)) {
        DbgPrint("[CommModifyKit] CreateControlDevice FAILED 0x%08X\n", status);
        return status;
    }

    DbgPrint("[CommModifyKit] DriverEntry OK, EvtDeviceAdd registered\n");
    return STATUS_SUCCESS;
}

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

    InterlockedIncrement(&g_diag_evt_device_add_count);
    DbgPrint("[CommModifyKit] >>> EvtDeviceAdd ENTER (count=%ld)\n", g_diag_evt_device_add_count);

    NTSTATUS status;

    // 1. 标记为过滤设备
    WdfFdoInitSetFilter(init);
    DbgPrint("[CommModifyKit] EvtDeviceAdd: WdfFdoInitSetFilter OK\n");

    // 2. 注册 PnP 电源回调
    WDF_PNPPOWER_EVENT_CALLBACKS pnp_callbacks;
    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnp_callbacks);
    pnp_callbacks.EvtDevicePrepareHardware = EvtDevicePrepareHardware;
    WdfDeviceInitSetPnpPowerEventCallbacks(init, &pnp_callbacks);
    DbgPrint("[CommModifyKit] EvtDeviceAdd: PnpPower callbacks registered\n");

    // 2.5 不使用 WdfDeviceInitSetFileObjectConfig
    // -----------------------------------------------
    // 原因：WdfDeviceInitSetFileObjectConfig 会让 WDF 为每个 CREATE 创建
    // WDFFILEOBJECT，并在 WdfDeviceWdmDispatchPreprocessedIrp 中强制查找它。
    // 当 WDFFILEOBJECT 查找失败时（异步 SEND_AND_FORGET 转发时序问题、
    // 或 FileObject 不匹配），WDF 返回 STATUS_FILE_FORCED_CLOSED (0xC0000182)，
    // 导致其他工具的 READ/WRITE IRP 全部失败。
    //
    // 改用 IRP_MJ_CREATE/CLOSE 预处理回调手动追踪文件打开/关闭和捕获 FileObject。
    // 没有 WdfDeviceInitSetFileObjectConfig，WdfDeviceWdmDispatchPreprocessedIrp
    // 不再查找 WDFFILEOBJECT，直接转发 IRP 到下层设备。

    // 3. 注册 WDM IRP 预处理回调
    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchCreate,
        IRP_MJ_CREATE, nullptr, 0);
    if (!NT_SUCCESS(status)) { InterlockedExchange(&g_diag_last_failure_status, status); return status; }
    DbgPrint("[CommModifyKit] EvtDeviceAdd: IRP_MJ_CREATE preprocess registered\n");

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchClose,
        IRP_MJ_CLOSE, nullptr, 0);
    if (!NT_SUCCESS(status)) { InterlockedExchange(&g_diag_last_failure_status, status); return status; }
    DbgPrint("[CommModifyKit] EvtDeviceAdd: IRP_MJ_CLOSE preprocess registered\n");

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchRead,
        IRP_MJ_READ, nullptr, 0);
    if (!NT_SUCCESS(status)) { InterlockedExchange(&g_diag_last_failure_status, status); return status; }
    DbgPrint("[CommModifyKit] EvtDeviceAdd: IRP_MJ_READ preprocess registered\n");

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchWrite,
        IRP_MJ_WRITE, nullptr, 0);
    if (!NT_SUCCESS(status)) { InterlockedExchange(&g_diag_last_failure_status, status); return status; }
    DbgPrint("[CommModifyKit] EvtDeviceAdd: IRP_MJ_WRITE preprocess registered\n");

    // IRP_MJ_CLEANUP / IRP_MJ_DEVICE_CONTROL 不注册预处理回调。
    // CLEANUP 由 WDF 自动转发（过滤驱动模式）。
    // DEVICE_CONTROL 由 WDF 自动转发给下层设备，串口 IOCTL（波特率/DCB 等）
    // 由虚拟串口驱动直接处理。
    InterlockedIncrement(&g_diag_irp_preprocess_ok);

    // 4. SDDL: WDF filter device 不支持 WdfDeviceInitAssignSDDLString
    //    它会返回 STATUS_INVALID_SECURITY_DESCR (0xC0000079)
    //    Filter device 会继承下层设备 (serial.sys) 的安全描述符
    //    不需要显式设置 SDDL

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

    DbgPrint("[CommModifyKit] EvtDeviceAdd: calling WdfDeviceCreate...\n");
    status = WdfDeviceCreate(&init, &attrs, &device);
    if (!NT_SUCCESS(status)) {
        DbgPrint("[CommModifyKit] EvtDeviceAdd: WdfDeviceCreate FAILED 0x%08X\n", status);
        InterlockedExchange(&g_diag_last_failure_status, status);
        return status;
    }
    DbgPrint("[CommModifyKit] EvtDeviceAdd: WdfDeviceCreate OK device=%p\n", device);

    InterlockedIncrement(&g_diag_device_create_ok);

    // 6. 初始化设备上下文
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    ctx->ComNumber = 0;
    ctx->MonitoringEnabled = FALSE;
    ctx->RingBuffer = nullptr;
    ctx->WdfDevice = device;
    ctx->LowerDevice = WdfDeviceWdmGetAttachedDevice(device);
    ctx->SavedFileObject = nullptr;
    KeInitializeSpinLock(&ctx->FileObjectLock);

    DbgPrint("[CommModifyKit] EvtDeviceAdd: ctx initialized, LowerDevice=%p\n", ctx->LowerDevice);

    // 注意：不再设置过滤设备 / 下层设备的 SecurityDescriptor。
    // 之前直接覆盖 ctx->LowerDevice->SecurityDescriptor 是严重错误：
    //   1. 下层设备不归本驱动管理（由虚拟串口驱动创建）
    //   2. 原始 SD 内存被泄漏
    //   3. 时机错误：EvtDeviceAdd 中下层设备可能未完成初始化
    //   4. 破坏下层驱动的内部状态，导致 CreateFile 异常失败
    // 过滤设备是未命名的，其 SD 不参与 CreateFile 的安全检查
    // （I/O Manager 查的是命名设备即下层 FDO 的 SD）。
    // 若需放宽非 elevated 进程的访问权限，应通过 INF 文件在设备类
    // 注册表项 HKLM\...\Control\Class\{Ports GUID}\Security 中设置。

    // 7. 注册到全局端口表
    commkit_driver::GetClientConnection().RegisterFilterDevice(device);
    DbgPrint("[CommModifyKit] EvtDeviceAdd: RegisterFilterDevice done, EvtDeviceAdd COMPLETE\n");

    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);

    InterlockedIncrement(&g_diag_prepare_hardware_count);
    DbgPrint("[CommModifyKit] EvtDevicePrepareHardware ENTER\n");

    // 从 PnP 属性解析 COM 编号
    // 方式：获取 DevicePropertyFriendlyName 或 DevicePropertyDeviceDescription，
    // 匹配 "COM%d" 模式
    WDFMEMORY memory;
    status = WdfDeviceAllocAndQueryProperty(
        device, DevicePropertyFriendlyName, NonPagedPoolNx, 0, &memory);
    if (!NT_SUCCESS(status)) {
        DbgPrint("[CommModifyKit] PrepareHardware: FriendlyName query failed 0x%08X, trying DeviceDescription\n", status);
        // 尝试 DeviceDescription 作为后备
        status = WdfDeviceAllocAndQueryProperty(
            device, DevicePropertyDeviceDescription, NonPagedPoolNx, 0, &memory);
    }

    if (NT_SUCCESS(status)) {
        PWCHAR name = (PWCHAR)WdfMemoryGetBuffer(memory, nullptr);
        if (name) {
            // 打印 FriendlyName/DeviceDescription 前 40 字符用于诊断
            DbgPrint("[CommModifyKit] PrepareHardware: name='%ws'\n", 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;
                            DbgPrint("[CommModifyKit] PrepareHardware: parsed COM%u from name\n", num);
                            break;
                        }
                    }
                }
            }
        }
        WdfObjectDelete(memory);
    }

    // 后备方案：FriendlyName/DeviceDescription 均未解析出 COM 编号时，
    // 从 PDO 名称查询 HKLM\HARDWARE\DEVICEMAP\SERIALCOMM 获取 "COMx" 映射
    if (ctx->ComNumber == 0) {
        DbgPrint("[CommModifyKit] PrepareHardware: trying SERIALCOMM registry fallback\n");
        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') {
                DbgPrint("[CommModifyKit] PrepareHardware: PDO name='%ws'\n", pdo_name);
                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);
                            DbgPrint("[CommModifyKit] PrepareHardware: SERIALCOMM value='%ws'\n", com_name);
                            // 解析 "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;
                                    DbgPrint("[CommModifyKit] PrepareHardware: parsed COM%u from SERIALCOMM\n", num);
                                }
                            }
                        }
                    } else {
                        DbgPrint("[CommModifyKit] PrepareHardware: SERIALCOMM query failed 0x%08X\n", status);
                    }
                    ZwClose(hkey);
                } else {
                    DbgPrint("[CommModifyKit] PrepareHardware: ZwOpenKey(SERIALCOMM) failed 0x%08X\n", status);
                }
            }
            WdfObjectDelete(pdo_memory);
        }
    }

    // 后备方案 2：查询本设备注册的 GUID_DEVINTERFACE_COMPORT 接口
    // 虚拟串口可能不在 SERIALCOMM 中，但会注册 COMPORT 接口
    // 注意：必须传入本设备 PDO，不能传 nullptr 枚举所有设备后用 IoGetDeviceObjectPointer 打开，
    //       因为那会发送 IRP_MJ_CREATE 到正在初始化的设备栈，导致死锁
    if (ctx->ComNumber == 0) {
        DbgPrint("[CommModifyKit] PrepareHardware: trying GUID_DEVINTERFACE_COMPORT fallback\n");
        PDEVICE_OBJECT currentPdo = WdfDeviceWdmGetPhysicalDevice(device);
        if (currentPdo) {
            PWSTR interfaceList = nullptr;
            status = IoGetDeviceInterfaces(
                &GUID_DEVINTERFACE_COMPORT,
                currentPdo,  // 仅查询本设备的接口
                DEVICE_INTERFACE_INCLUDE_NONACTIVE,
                &interfaceList);
            if (NT_SUCCESS(status) && interfaceList) {
                PWCHAR p = interfaceList;
                while (*p && ctx->ComNumber == 0) {
                    ULONG len = (ULONG)wcslen(p);
                    DbgPrint("[CommModifyKit] PrepareHardware: interface='%ws'\n", p);
                    // 从符号链接路径解析 COM 编号（如 "\\?\COM12"）
                    for (ULONG i = 0; i + 3 <= len; i++) {
                        if ((p[i] == L'C' || p[i] == L'c') &&
                            (p[i+1] == L'O' || p[i+1] == L'o') &&
                            (p[i+2] == L'M' || p[i+2] == L'm') &&
                            p[i+3] >= L'0' && p[i+3] <= L'9') {
                            ULONG num = 0;
                            ULONG j = i + 3;
                            while (j < len && p[j] >= L'0' && p[j] <= L'9') {
                                num = num * 10 + (p[j] - L'0');
                                j++;
                            }
                            if (num > 0 && num <= 256) {
                                ctx->ComNumber = num;
                                DbgPrint("[CommModifyKit] PrepareHardware: parsed COM%u from interface\n", num);
                                break;
                            }
                        }
                    }
                    p += len + 1;
                }
                ExFreePool(interfaceList);
            } else {
                DbgPrint("[CommModifyKit] PrepareHardware: IoGetDeviceInterfaces failed 0x%08X or empty\n", status);
            }
        }
    }

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

    DbgPrint("[CommModifyKit] EvtDevicePrepareHardware COMPLETE, ComNumber=%u\n", ctx->ComNumber);
    return STATUS_SUCCESS;
}

// 设备上下文清理回调：设备移除时注销端口表条目
VOID EvtDeviceContextCleanup(WDFOBJECT object) {
    WDFDEVICE device = (WDFDEVICE)object;
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    if (!ctx) return;

    // 释放保存的 FileObject
    KIRQL old_irql;
    KeAcquireSpinLock(&ctx->FileObjectLock, &old_irql);
    PFILE_OBJECT saved = ctx->SavedFileObject;
    ctx->SavedFileObject = nullptr;
    KeReleaseSpinLock(&ctx->FileObjectLock, old_irql);
    if (saved) {
        ObDereferenceObject(saved);
    }

    if (ctx->ComNumber != 0) {
        commkit_driver::GetClientConnection().UnregisterFilterDevice(ctx->ComNumber);
    }
}

// 文件打开/关闭/数据捕获现在通过 IRP_MJ_CREATE/CLOSE/READ/WRITE 预处理回调处理
// （见 SerialFilter.cpp 中的 DispatchCreate/DispatchClose/DispatchRead/DispatchWrite）。
// 不再使用 WdfDeviceInitSetFileObjectConfig + EvtDeviceFileCreate/EvtFileClose/EvtFileCleanup，
// 因为 WdfDeviceInitSetFileObjectConfig 会导致 WdfDeviceWdmDispatchPreprocessedIrp
// 强制查找 WDFFILEOBJECT，查找失败时返回 STATUS_FILE_FORCED_CLOSED (0xC0000182)。

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