// SerialFilter.cpp - 串口 IRP 拦截实现
// KMDF 过滤驱动模式：通过 WdfDeviceInitAssignWdmIrpPreprocessCallback 注册
// 在预处理回调中注册完成例程，捕获 READ/WRITE 数据
#include <ntddk.h>
#include <wdf.h>
#include "SerialFilter.h"
#include "EventRingBuffer.h"
#include "ClientConnection.h"
#include "DeviceContext.h"
#include "../common/CommKitEvents.h"

namespace commkit_driver {

// 全局原子序列号（所有事件类型共享，保证 Sequence 字段全局唯一递增）
LONG g_sequence_counter = 0;

// 获取当前时间戳（100ns 单位）
static UINT64 GetCurrentTimeStamp() {
    return KeQueryInterruptTime();
}

void SerialFilter::CaptureIrpData(PDEVICE_CONTEXT ctx, PIRP irp, ULONG event_type) {
    if (!ctx || !ctx->MonitoringEnabled || !ctx->RingBuffer) {
        return;
    }

    PIO_STACK_LOCATION sl = IoGetCurrentIrpStackLocation(irp);
    if (!sl) return;

    // 仅在 IRP 成功完成时捕获
    if (!NT_SUCCESS(irp->IoStatus.Status)) {
        return;
    }

    ULONG data_size = 0;
    PVOID data_ptr = nullptr;

    // 根据 IRP 类型提取数据
    if (event_type == COMMKIT_OP_READ) {
        // 实际读取的字节数（IRP 完成后 Information = 读取字节数）
        data_size = (ULONG)irp->IoStatus.Information;
        if (data_size == 0) return;
        // 数据可能在 SystemBuffer / MdlAddress / UserBuffer
        if (irp->MdlAddress) {
            data_ptr = MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority);
        } else if (irp->AssociatedIrp.SystemBuffer) {
            data_ptr = irp->AssociatedIrp.SystemBuffer;
        }
    } else if (event_type == COMMKIT_OP_WRITE) {
        data_size = (ULONG)irp->IoStatus.Information;
        if (data_size == 0) {
            data_size = sl->Parameters.Write.Length;
        }
        if (data_size == 0) return;
        if (irp->MdlAddress) {
            data_ptr = MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority);
        } else if (irp->AssociatedIrp.SystemBuffer) {
            data_ptr = irp->AssociatedIrp.SystemBuffer;
        }
    }

    if (!data_ptr || data_size == 0) {
        return;
    }

    // 推入环形缓冲
    LONG seq = InterlockedIncrement(&g_sequence_counter);
    auto* ring = (EventRingBuffer*)ctx->RingBuffer;
    ring->Push((ULONG)seq, GetCurrentTimeStamp(),
                ctx->ComNumber, event_type,
                data_size, data_ptr);
}

NTSTATUS SerialFilter::OnReadComplete(PDEVICE_OBJECT dev, PIRP irp, PVOID ctx) {
    UNREFERENCED_PARAMETER(dev);
    PDEVICE_CONTEXT device_ctx = (PDEVICE_CONTEXT)ctx;
    if (device_ctx) {
        CaptureIrpData(device_ctx, irp, COMMKIT_OP_READ);
    }
    // 下层驱动返回 STATUS_PENDING 时需标记 IRP 为 pending
    if (irp->PendingReturned) {
        IoMarkIrpPending(irp);
    }
    return STATUS_SUCCESS;
}

NTSTATUS SerialFilter::OnWriteComplete(PDEVICE_OBJECT dev, PIRP irp, PVOID ctx) {
    UNREFERENCED_PARAMETER(dev);
    PDEVICE_CONTEXT device_ctx = (PDEVICE_CONTEXT)ctx;
    if (device_ctx) {
        CaptureIrpData(device_ctx, irp, COMMKIT_OP_WRITE);
    }
    if (irp->PendingReturned) {
        IoMarkIrpPending(irp);
    }
    return STATUS_SUCCESS;
}

NTSTATUS SerialFilter::DispatchCreate(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);
    ULONG com = ctx ? ctx->ComNumber : 0;

    // 从 IRP 栈位置获取 WDM FileObject，用于 HandleWritePort/HandleReadPort
    PIO_STACK_LOCATION sl = IoGetCurrentIrpStackLocation(irp);
    if (ctx && sl && sl->FileObject) {
        PFILE_OBJECT wdm_fo = sl->FileObject;
        ObReferenceObject(wdm_fo);
        KIRQL old_irql;
        KeAcquireSpinLock(&ctx->FileObjectLock, &old_irql);
        PFILE_OBJECT old_fo = ctx->SavedFileObject;
        ctx->SavedFileObject = wdm_fo;
        KeReleaseSpinLock(&ctx->FileObjectLock, old_irql);
        if (old_fo) {
            ObDereferenceObject(old_fo);
        }
        DbgPrint("[CommModifyKit] DispatchCreate COM%u: SavedFileObject=%p\n", com, wdm_fo);
    } else {
        DbgPrint("[CommModifyKit] DispatchCreate COM%u: no FileObject in IRP\n", com);
    }

    // OP_OPEN 事件
    if (ctx && ctx->MonitoringEnabled && ctx->RingBuffer) {
        EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
        LONG seq = InterlockedIncrement(&g_sequence_counter);
        UINT64 timestamp = KeQueryInterruptTime();
        ring->Push((ULONG)seq, timestamp, com, COMMKIT_OP_OPEN, 0, nullptr);
    }

    // 转发到下层设备（无 WdfDeviceInitSetFileObjectConfig，WDF 不查找 WDFFILEOBJECT）
    NTSTATUS fwd_status = WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
    DbgPrint("[CommModifyKit] DispatchCreate COM%u: forward status=0x%08X LowerDev=%p\n",
             com, fwd_status, ctx ? ctx->LowerDevice : nullptr);
    return fwd_status;
}

NTSTATUS SerialFilter::DispatchClose(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);
    ULONG com = ctx ? ctx->ComNumber : 0;

    // 释放捕获的 FileObject
    PIO_STACK_LOCATION sl = IoGetCurrentIrpStackLocation(irp);
    if (ctx && sl && sl->FileObject) {
        PFILE_OBJECT wdm_fo = sl->FileObject;
        KIRQL old_irql;
        KeAcquireSpinLock(&ctx->FileObjectLock, &old_irql);
        PFILE_OBJECT saved = ctx->SavedFileObject;
        if (saved == wdm_fo) {
            ctx->SavedFileObject = nullptr;
        } else {
            saved = nullptr;
        }
        KeReleaseSpinLock(&ctx->FileObjectLock, old_irql);
        if (saved) {
            ObDereferenceObject(saved);
            DbgPrint("[CommModifyKit] DispatchClose COM%u: SavedFileObject released\n", com);
        }
    }

    // OP_CLOSE 事件
    if (ctx && ctx->MonitoringEnabled && ctx->RingBuffer) {
        EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
        LONG seq = InterlockedIncrement(&g_sequence_counter);
        UINT64 timestamp = KeQueryInterruptTime();
        ring->Push((ULONG)seq, timestamp, com, COMMKIT_OP_CLOSE, 0, nullptr);
    }

    NTSTATUS fwd_status = WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
    DbgPrint("[CommModifyKit] DispatchClose COM%u: forward status=0x%08X LowerDev=%p\n",
             com, fwd_status, ctx ? ctx->LowerDevice : nullptr);
    return fwd_status;
}

NTSTATUS SerialFilter::DispatchRead(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);
    ULONG com = ctx ? ctx->ComNumber : 0;

    // 在转发前读取栈位置（转发后可能失效）
    PIO_STACK_LOCATION sl_before = IoGetCurrentIrpStackLocation(irp);
    ULONG req_len = sl_before ? sl_before->Parameters.Read.Length : 0;
    PFILE_OBJECT irp_fo = sl_before ? sl_before->FileObject : nullptr;
    PDEVICE_OBJECT lower_dev = ctx ? ctx->LowerDevice : nullptr;
    BOOLEAN monitored = (ctx && ctx->MonitoringEnabled && ctx->RingBuffer);

    NTSTATUS status;
    const char* fwd_method = "";

    if (monitored) {
        // 监控端口：设置完成例程后手动转发（捕获数据需要完成例程）
        IoSetCompletionRoutine(irp, OnReadComplete, ctx, TRUE, TRUE, TRUE);
        IoSkipCurrentIrpStackLocation(irp);
        status = IoCallDriver(lower_dev, irp);
        fwd_method = "IoCallDriver";
    } else {
        // 非监控端口：使用 WdfDeviceWdmDispatchPreprocessedIrp 转发
        // 原因：CREATE/CLOSE 使用 WdfDeviceWdmDispatchPreprocessedIrp，
        // READ/WRITE 若用 IoCallDriver 会绕过 WDF 内部 IRP 跟踪，
        // 导致 WDF 状态不一致，可能引发后续 IRP 异常。
        // 统一使用 WdfDeviceWdmDispatchPreprocessedIrp 确保 WDF 一致性。
        status = WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
        fwd_method = "WdfDispatch";
    }

    // 诊断：每个端口前 5 次 IRP 记录详细信息
    if (com > 0 && com <= 256) {
        static volatile LONG g_read_log_count[257] = {0};
        LONG count = InterlockedIncrement(&g_read_log_count[com]);
        if (count <= 5) {
            DbgPrint("[CommModifyKit] IRP_MJ_READ COM%u status=0x%08X reqLen=%u mon=%d FO=%p Low=%p [%s] (#%ld)\n",
                     com, status, req_len, (int)monitored, irp_fo, lower_dev, fwd_method, count);
        }
    }
    return status;
}

NTSTATUS SerialFilter::DispatchWrite(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);
    ULONG com = ctx ? ctx->ComNumber : 0;

    PIO_STACK_LOCATION sl_before = IoGetCurrentIrpStackLocation(irp);
    ULONG req_len = sl_before ? sl_before->Parameters.Write.Length : 0;
    PFILE_OBJECT irp_fo = sl_before ? sl_before->FileObject : nullptr;
    PDEVICE_OBJECT lower_dev = ctx ? ctx->LowerDevice : nullptr;
    BOOLEAN monitored = (ctx && ctx->MonitoringEnabled && ctx->RingBuffer);

    NTSTATUS status;
    const char* fwd_method = "";

    if (monitored) {
        // 监控端口：设置完成例程后手动转发
        IoSetCompletionRoutine(irp, OnWriteComplete, ctx, TRUE, TRUE, TRUE);
        IoSkipCurrentIrpStackLocation(irp);
        status = IoCallDriver(lower_dev, irp);
        fwd_method = "IoCallDriver";
    } else {
        // 非监控端口：使用 WdfDeviceWdmDispatchPreprocessedIrp 转发（同 DispatchRead）
        status = WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
        fwd_method = "WdfDispatch";
    }

    if (com > 0 && com <= 256) {
        static volatile LONG g_write_log_count[257] = {0};
        LONG count = InterlockedIncrement(&g_write_log_count[com]);
        if (count <= 5) {
            DbgPrint("[CommModifyKit] IRP_MJ_WRITE COM%u status=0x%08X reqLen=%u mon=%d FO=%p Low=%p [%s] (#%ld)\n",
                     com, status, req_len, (int)monitored, irp_fo, lower_dev, fwd_method, count);
        }
    }
    return status;
}

} // namespace commkit_driver