// Network device registry. This file owns net-device registration and RX handoff to the active stack provider.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <ntlite/ddi/net_device.h>
#include <ntlite/ddi/net_stack.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_spinlock.h>

#define NET_MAX_DEVICES 8u

static NET_DEVICE* gNetDevices[NET_MAX_DEVICES];
static uint32_t    gNetDeviceCount = 0;
static SPINLOCK    gNetLock = SPINLOCK_INIT;

// Implements the net Make Default Name operation.
static void NetMakeDefaultName(NET_DEVICE* device, uint32_t index) {
    if (!device || device->Name[0])
        return;

    device->Name[0] = 'e';
    device->Name[1] = 't';
    device->Name[2] = 'h';
    device->Name[3] = (char)('0' + (index % 10u));
    device->Name[4] = '\0';
}

// Implements the net Is Device Registered Locked operation.
static bool NetIsDeviceRegisteredLocked(const NET_DEVICE* device) {
    for (uint32_t i = 0; i < gNetDeviceCount; ++i) {
        if (gNetDevices[i] == device)
            return true;
    }
    return false;
}

// Implements the net Device Has Ethernet Address operation.
static bool NetDeviceHasEthernetAddress(const NET_DEVICE* device) {
    return device && device->HardwareAddressLength == NET_ETHERNET_ADDRESS_LENGTH;
}

// Implements the net Log Hardware Address operation.
static void NetLogHardwareAddress(const NET_DEVICE* device) {
    if (NetDeviceHasEthernetAddress(device)) {
        DbgPrint("%02x:%02x:%02x:%02x:%02x:%02x",
                 device->HardwareAddress[0], device->HardwareAddress[1],
                 device->HardwareAddress[2], device->HardwareAddress[3],
                 device->HardwareAddress[4], device->HardwareAddress[5]);
    } else {
        DbgPrint("hwlen=%u", (unsigned)device->HardwareAddressLength);
    }
}

// Implements the net Register Device operation.
__declspec(dllexport) KSTATUS NetRegisterDevice(NET_DEVICE* device) {
    if (!device || !device->SendFrame)
        return NTLSTATUS_INVALID_PARAMETER;
    if (device->HardwareAddressLength > NET_HWADDR_MAX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    if (NetIsDeviceRegisteredLocked(device)) {
        KeReleaseSpinlock(&gNetLock, flags);
        return NTLSTATUS_ALREADY_EXISTS;
    }
    if (gNetDeviceCount >= NET_MAX_DEVICES) {
        KeReleaseSpinlock(&gNetLock, flags);
        return NTLSTATUS_QUEUE_FULL;
    }

    uint32_t index = gNetDeviceCount;
    device->Index = index;
    if (device->Mtu == 0)
        device->Mtu = NET_MTU_ETHERNET;
    memset(&device->Stats, 0, sizeof(device->Stats));
    NetMakeDefaultName(device, index);
    gNetDevices[index] = device;
    gNetDeviceCount++;
    KeReleaseSpinlock(&gNetLock, flags);

    DbgPrint("NET: registered %s index=%u mac=", device->Name, device->Index);
    NetLogHardwareAddress(device);
    DbgPrint(" mtu=%u\n", device->Mtu);

    const NET_STACK_PROVIDER* provider = NetGetStackProvider();
    if (provider && provider->AttachDevice) {
        KSTATUS attach = provider->AttachDevice(device);
        if (!KSUCCESS(attach)) {
            DbgPrint("NET: provider %s attach for %s failed: 0x%x\n",
                     provider->Name, device->Name, attach);
        }
    }

    return NTLSTATUS_SUCCESS;
}

// Implements the net Get Device Count operation.
__declspec(dllexport) uint32_t NetGetDeviceCount(void) {
    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    uint32_t count = gNetDeviceCount;
    KeReleaseSpinlock(&gNetLock, flags);
    return count;
}

// Implements the net Get Device operation.
__declspec(dllexport) NET_DEVICE* NetGetDevice(uint32_t index) {
    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    NET_DEVICE* device = (index < gNetDeviceCount) ? gNetDevices[index] : (NET_DEVICE*)0;
    KeReleaseSpinlock(&gNetLock, flags);
    return device;
}

__declspec(dllexport) KSTATUS NetGetDeviceStats(uint32_t index,
                                                NET_DEVICE_STATS* stats) {
    if (!stats)
        return NTLSTATUS_INVALID_PARAMETER;

    NET_DEVICE* device = NetGetDevice(index);
    if (!device)
        return NTLSTATUS_NOT_FOUND;

    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    memcpy(stats, &device->Stats, sizeof(*stats));
    KeReleaseSpinlock(&gNetLock, flags);
    return NTLSTATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS NetSendFrame(uint32_t index,
                                           const void* frame,
                                           size_t length) {
    NET_DEVICE* device = NetGetDevice(index);
    if (!device || !device->SendFrame)
        return NTLSTATUS_NOT_FOUND;
    if (!frame || length < NET_ETHERNET_HEADER_LENGTH)
        return NTLSTATUS_INVALID_PARAMETER;
    if (length > (size_t)device->Mtu + NET_ETHERNET_HEADER_LENGTH)
        return NTLSTATUS_BUFFER_TOO_SMALL;

    KSTATUS st = device->SendFrame(device, frame, length);
    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    if (KSUCCESS(st)) {
        device->Stats.TxPackets++;
        device->Stats.TxBytes += length;
    } else {
        device->Stats.TxErrors++;
    }
    KeReleaseSpinlock(&gNetLock, flags);
    return st;
}

__declspec(dllexport) void NetReceiveFrame(NET_DEVICE* device,
                                           const void* frame,
                                           size_t length) {
    if (!device || !frame || length < NET_ETHERNET_HEADER_LENGTH) {
        if (device) {
            uintptr_t dropFlags = KeAcquireSpinlock(&gNetLock);
            device->Stats.RxDrops++;
            KeReleaseSpinlock(&gNetLock, dropFlags);
        }
        return;
    }

    uintptr_t flags = KeAcquireSpinlock(&gNetLock);
    device->Stats.RxPackets++;
    device->Stats.RxBytes += length;
    KeReleaseSpinlock(&gNetLock, flags);

    const NET_STACK_PROVIDER* provider = NetGetStackProvider();
    KSTATUS rxStatus = NTLSTATUS_NOT_SUPPORTED;
    if (provider && provider->ReceiveFrame)
        rxStatus = provider->ReceiveFrame(device, frame, length);

    if (!KSUCCESS(rxStatus)) {
        uintptr_t dropFlags = KeAcquireSpinlock(&gNetLock);
        device->Stats.RxDrops++;
        KeReleaseSpinlock(&gNetLock, dropFlags);
    }
}