#include <ntlite/ddk/module.h>
// Intel High Definition Audio PCI attachment and controller lifetime management.

#include "hda.h"

#include <string.h>

static DRIVER_OBJECT gHdaDriver;
static HDA_CONTROLLER* gHdaControllers[HDA_MAX_CONTROLLERS];
static uint32_t gHdaControllerCount;
static SPINLOCK gHdaRegistryLock = SPINLOCK_INIT;

static const IO_DEVICE_ID gHdaIds[] = {
    { IoBusPci, { "PCI\\VEN_8086&DEV_2668", sizeof("PCI\\VEN_8086&DEV_2668") - 1u } },
    { IoBusPci, { "PCI\\VEN_15AD&DEV_1977", sizeof("PCI\\VEN_15AD&DEV_1977") - 1u } },
    { IoBusPci, { "PCI\\CC_040300", sizeof("PCI\\CC_040300") - 1u } },
};

static const IO_DRIVER_ID_TABLE gHdaIdTable = {
    gHdaIds,
    sizeof(gHdaIds) / sizeof(gHdaIds[0]),
};

// Registers a controller in the shared interrupt scan set.
static KSTATUS HdaRegisterController(HDA_CONTROLLER* controller) {
    if (!controller)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gHdaRegistryLock);
    for (uint32_t index = 0; index < gHdaControllerCount; index++) {
        if (gHdaControllers[index] == controller) {
            KeReleaseSpinlock(&gHdaRegistryLock, flags);
            return NTLSTATUS_SUCCESS;
        }
    }
    if (gHdaControllerCount == HDA_MAX_CONTROLLERS) {
        KeReleaseSpinlock(&gHdaRegistryLock, flags);
        return NTLSTATUS_NO_MEMORY;
    }
    gHdaControllers[gHdaControllerCount++] = controller;
    KeReleaseSpinlock(&gHdaRegistryLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Removes a controller from the shared interrupt scan set.
static void HdaUnregisterController(HDA_CONTROLLER* controller) {
    if (!controller)
        return;

    uintptr_t flags = KeAcquireSpinlock(&gHdaRegistryLock);
    for (uint32_t index = 0; index < gHdaControllerCount; index++) {
        if (gHdaControllers[index] != controller)
            continue;
        gHdaControllers[index] = gHdaControllers[gHdaControllerCount - 1u];
        gHdaControllers[gHdaControllerCount - 1u] = (HDA_CONTROLLER*)0;
        gHdaControllerCount--;
        break;
    }
    KeReleaseSpinlock(&gHdaRegistryLock, flags);
}

// Dispatches a hardware interrupt to every live HDA controller.
static void HdaIsr(void) {
    DEVICE_OBJECT* devices[HDA_MAX_CONTROLLERS];
    uint32_t count = 0;

    uintptr_t flags = KeAcquireSpinlock(&gHdaRegistryLock);
    count = gHdaControllerCount;
    if (count > HDA_MAX_CONTROLLERS)
        count = HDA_MAX_CONTROLLERS;
    for (uint32_t index = 0; index < count; index++) {
        HDA_CONTROLLER* controller = gHdaControllers[index];
        devices[index] = controller ? controller->Fdo : (DEVICE_OBJECT*)0;
        if (devices[index])
            IoReferenceDeviceObject(devices[index]);
    }
    KeReleaseSpinlock(&gHdaRegistryLock, flags);

    for (uint32_t index = 0; index < count; index++) {
        DEVICE_OBJECT* device = devices[index];
        if (!device)
            continue;
        HDA_CONTROLLER* controller = (HDA_CONTROLLER*)device->DeviceExtension;
        if (controller && controller->Magic == HDA_CONTROLLER_MAGIC)
            HdaHandleInterrupt(controller);
        IoDereferenceDeviceObject(device);
    }
}

// Registers the global ISR once for each kernel IRQ slot used by HDA.
static KSTATUS HdaRegisterIrq(uint8_t irq) {
    if (irq >= 24u)
        return NTLSTATUS_NOT_SUPPORTED;
    KeIrqRegister(irq, HdaIsr);
    return NTLSTATUS_SUCCESS;
}

// Maps BAR0 and validates the stream-descriptor aperture described by GCAP.
static KSTATUS HdaMapRegisters(HDA_CONTROLLER* controller) {
    PCI_BAR bar;
    KSTATUS status = PciGetBar(controller->Pdo, 0, &bar);
    if (!KSUCCESS(status))
        return status;
    if (bar.Type != PciBarMemory32 && bar.Type != PciBarMemory64)
        return NTLSTATUS_NOT_SUPPORTED;

    status = PciMapBar(controller->Pdo,
                       0,
                       PAGE_PRESENT | PAGE_WRITE | PAGE_CACHE_UC,
                       (void**)&controller->Registers,
                       &controller->RegisterBytes);
    if (!KSUCCESS(status))
        return status;
    if (controller->RegisterBytes < HDA_REG_STREAM_BASE)
        return NTLSTATUS_BUFFER_TOO_SMALL;

    controller->Gcap = HdaRead16(controller, HDA_REG_GCAP);
    controller->InputStreamCount =
        (uint8_t)((controller->Gcap & HDA_GCAP_ISS_MASK) >> HDA_GCAP_ISS_SHIFT);
    controller->OutputStreamCount =
        (uint8_t)((controller->Gcap & HDA_GCAP_OSS_MASK) >> HDA_GCAP_OSS_SHIFT);
    controller->BidirectionalStreamCount =
        (uint8_t)((controller->Gcap & HDA_GCAP_BSS_MASK) >> HDA_GCAP_BSS_SHIFT);
    controller->TotalStreamCount =
        (uint8_t)(controller->InputStreamCount +
                  controller->OutputStreamCount +
                  controller->BidirectionalStreamCount);
    if ((controller->OutputStreamCount == 0 &&
         controller->BidirectionalStreamCount == 0) ||
        controller->TotalStreamCount == 0 ||
        controller->TotalStreamCount > 30u)
        return NTLSTATUS_NOT_SUPPORTED;

    uint64_t requiredBytes = HDA_REG_STREAM_BASE +
                             (uint64_t)controller->TotalStreamCount *
                             HDA_REG_STREAM_STRIDE;
    if (controller->RegisterBytes < requiredBytes)
        return NTLSTATUS_BUFFER_TOO_SMALL;

    // Dedicated output descriptors precede bidirectional descriptors. Use the
    // first bidirectional descriptor when no output-only descriptor exists.
    uint8_t outputDescriptorBase = controller->InputStreamCount;
    uint8_t bidirectionalDescriptorBase =
        (uint8_t)(controller->InputStreamCount + controller->OutputStreamCount);
    controller->Output.DescriptorIndex = controller->OutputStreamCount
                                       ? outputDescriptorBase
                                       : bidirectionalDescriptorBase;
    controller->Output.StreamTag = HDA_OUTPUT_STREAM_TAG;
    return NTLSTATUS_SUCCESS;
}

// Configures MSI first and enables INTx only after MSI setup reports failure.
static KSTATUS HdaConfigureInterrupt(HDA_CONTROLLER* controller) {
    KSTATUS status = PciEnableMsiAuto(controller->Pdo, &controller->Irq);
    if (KSUCCESS(status)) {
        controller->UsingMsi = true;
        return NTLSTATUS_SUCCESS;
    }

    DbgPrint(
        "HDA: MSI setup failed status=0x%x; attempting explicit INTx fallback\n",
        status);
    controller->UsingMsi = false;
    status = PciGetInterrupt(controller->Pdo, &controller->Irq);
    if (!KSUCCESS(status))
        return status;
    if (controller->Irq == 0xFFu || controller->Irq >= 24u)
        return NTLSTATUS_NOT_SUPPORTED;

    status = PciEnableDevice(controller->Pdo,
                             PCI_ENABLE_MEMORY |
                             PCI_ENABLE_BUSMASTER |
                             PCI_ENABLE_INTX);
    return status;
}

// Waits until queued or active interrupt work no longer references the controller.
static KSTATUS HdaWaitForControllerInterruptWorkIdle(HDA_CONTROLLER* controller) {
    if (!controller)
        return NTLSTATUS_INVALID_PARAMETER;
    for (uint32_t elapsed = 0; elapsed < 100000u; elapsed += 10u) {
        if (controller->IsrActive == 0 &&
            controller->Dpc.Inserted == 0 && controller->DpcActive == 0)
            return NTLSTATUS_SUCCESS;
        KeStallExecution(10);
    }
    return NTLSTATUS_TIMEOUT;
}

// Starts one PCI HDA controller and leaves endpoint registration to codec discovery.
static KSTATUS HdaStartDevice(HDA_CONTROLLER* controller) {
    if (!controller || controller->Magic != HDA_CONTROLLER_MAGIC ||
        !controller->Pdo || !controller->Fdo)
        return NTLSTATUS_INVALID_PARAMETER;
    if (controller->Started)
        return NTLSTATUS_SUCCESS;

    KSTATUS status = PciGetDeviceInfo(controller->Pdo, &controller->PciInfo);
    if (!KSUCCESS(status))
        return status;
    if (controller->PciInfo.ClassCode != 0x04u ||
        controller->PciInfo.Subclass != 0x03u ||
        controller->PciInfo.ProgIf != 0x00u)
        return NTLSTATUS_NOT_SUPPORTED;

    status = PciEnableDevice(controller->Pdo,
                             PCI_ENABLE_MEMORY |
                             PCI_ENABLE_BUSMASTER |
                             PCI_DISABLE_INTX);
    if (!KSUCCESS(status))
        return status;

    status = HdaMapRegisters(controller);
    if (!KSUCCESS(status))
        return status;

    status = HdaConfigureInterrupt(controller);
    if (!KSUCCESS(status))
        return status;

    status = HdaAllocateDma(controller);
    if (!KSUCCESS(status))
        return status;
    KeInitializeDpc(&controller->Dpc, HdaDpcRoutine, controller);

    status = HdaStartRecoveryWorker(controller);
    if (!KSUCCESS(status)) {
        HdaReleaseDma(controller);
        return status;
    }

    status = HdaInitializeController(controller);
    if (!KSUCCESS(status)) {
        HdaStopRecoveryWorker(controller);
        HdaQuiesceController(controller);
        HdaReleaseDma(controller);
        return status;
    }

    status = HdaRegisterController(controller);
    if (!KSUCCESS(status)) {
        HdaStopRecoveryWorker(controller);
        HdaQuiesceController(controller);
        HdaReleaseDma(controller);
        return status;
    }
    status = HdaRegisterIrq(controller->Irq);
    if (!KSUCCESS(status)) {
        HdaUnregisterController(controller);
        HdaStopRecoveryWorker(controller);
        HdaQuiesceController(controller);
        HdaReleaseDma(controller);
        return status;
    }

    controller->Started = true;
    controller->Fdo->PnpState = DevicePnpStarted;

    status = HdaCodecInitialize(controller);
    if (!KSUCCESS(status) && status != NTLSTATUS_NOT_SUPPORTED &&
        status != NTLSTATUS_NOT_FOUND) {
        controller->Started = false;
        controller->Fdo->PnpState = DevicePnpNotStarted;
        HdaUnregisterController(controller);
        HdaStopRecoveryWorker(controller);
        HdaQuiesceController(controller);
        HdaReleaseDma(controller);
        return status;
    }
    if (!KSUCCESS(status)) {
        DbgPrint(
            "HDA: no analog endpoint registered status=0x%x\n",
            status);
    }

    HdaArmInterrupts(controller);
    HdaHandleInterrupt(controller);

    return NTLSTATUS_SUCCESS;
}

// Handles PnP lifetime transitions for one HDA function device object.
static KSTATUS HdaDispatchPnp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp || !device->DeviceExtension)
        return NTLSTATUS_INVALID_PARAMETER;
    HDA_CONTROLLER* controller = (HDA_CONTROLLER*)device->DeviceExtension;
    if (controller->Magic != HDA_CONTROLLER_MAGIC)
        return NTLSTATUS_INVALID_PARAMETER;

    IO_STACK_LOCATION* stack = IoGetCurrentIrpStackLocation(irp);
    if (!stack || stack->MajorFunction != IRP_MJ_PNP)
        return NTLSTATUS_INVALID_PARAMETER;

    switch ((IO_PNP_MINOR_FUNCTION)stack->MinorFunction) {
        case IRP_MN_START_DEVICE:
            return HdaStartDevice(controller);

        case IRP_MN_QUERY_REMOVE_DEVICE:
        case IRP_MN_CANCEL_REMOVE_DEVICE:
            return NTLSTATUS_SUCCESS;

        case IRP_MN_SURPRISE_REMOVAL:
            controller->Started = false;
            controller->SurpriseRemoved = true;
            HdaUnregisterController(controller);
            HdaStopRecoveryWorker(controller);
            KSTATUS surpriseDpcStatus =
                HdaWaitForControllerInterruptWorkIdle(controller);
            if (!KSUCCESS(surpriseDpcStatus)) {
                KePanic(KE_PANIC_CODE_DRIVER,
                        "HDA: interrupt work did not quiesce during surprise removal status=0x%x\n",
                        surpriseDpcStatus);
            }
            HdaCodecStop(controller, IoAudioSurpriseRemoval);
            HdaCodecDestroy(controller);
            device->PnpState = DevicePnpSurpriseRemoved;
            return NTLSTATUS_SUCCESS;

        case IRP_MN_REMOVE_DEVICE:
            if (!controller->SurpriseRemoved && controller->Registers)
                HdaWrite32(controller, HDA_REG_INTCTL, 0);
            controller->Started = false;
            HdaUnregisterController(controller);
            HdaStopRecoveryWorker(controller);
            KSTATUS dpcStatus = HdaWaitForControllerInterruptWorkIdle(controller);
            if (!KSUCCESS(dpcStatus)) {
                KePanic(KE_PANIC_CODE_DRIVER,
                        "HDA: interrupt work did not quiesce during removal status=0x%x\n",
                        dpcStatus);
            }
            HdaCodecStop(controller,
                         controller->SurpriseRemoved
                             ? IoAudioSurpriseRemoval
                             : IoAudioOrderlyRemoval);
            HdaCodecDestroy(controller);
            if (!controller->SurpriseRemoved)
                HdaQuiesceController(controller);
            controller->Removed = true;
            HdaReleaseDma(controller);
            if (!controller->SurpriseRemoved)
                (void)PciEnableDevice(controller->Pdo, PCI_DISABLE_INTX);
            device->PnpState = DevicePnpDeleted;
            IoDetachDevice(device);
            IoDeleteDevice(device);
            return NTLSTATUS_SUCCESS;

        default:
            return IoForwardIrpToLower(device, irp);
    }
}

// Attaches one HDA function device object to a matching PCI PDO.
static KSTATUS HdaAddDevice(DRIVER_OBJECT* driver, DEVICE_OBJECT* physicalDeviceObject) {
    if (!driver || !physicalDeviceObject ||
        physicalDeviceObject->Type != DeviceObjectPhysical)
        return NTLSTATUS_INVALID_PARAMETER;

    DEVICE_OBJECT* fdo = (DEVICE_OBJECT*)0;
    KSTATUS status = IoCreateDevice(driver, sizeof(HDA_CONTROLLER), &fdo);
    if (!KSUCCESS(status))
        return status;
    (void)IoSetDeviceName(fdo, SYS_SV_LITERAL("hda"));

    HDA_CONTROLLER* controller = (HDA_CONTROLLER*)fdo->DeviceExtension;
    memset(controller, 0, sizeof(*controller));
    controller->Magic = HDA_CONTROLLER_MAGIC;
    controller->Fdo = fdo;
    controller->Pdo = physicalDeviceObject;
    controller->Irq = 0xFFu;
    controller->StateLock = (SPINLOCK)SPINLOCK_INIT;
    controller->Output.StreamTag = HDA_OUTPUT_STREAM_TAG;

    controller->LowerDevice = IoAttachDeviceToDeviceStack(fdo, physicalDeviceObject);
    if (!controller->LowerDevice) {
        IoDeleteDevice(fdo);
        return NTLSTATUS_FAILURE;
    }
    return NTLSTATUS_SUCCESS;
}

// Registers the Intel HDA transport driver with the I/O manager.
static KSTATUS DriverInitialize(NTL_MODULE* module) {
    (void)module;

    gHdaControllerCount = 0;
    memset(gHdaControllers, 0, sizeof(gHdaControllers));

    memset(&gHdaDriver, 0, sizeof(gHdaDriver));
    IoInitializeDriverObject(&gHdaDriver, SYS_SV_LITERAL("hda"));
    gHdaDriver.IdTable = &gHdaIdTable;
    gHdaDriver.DriverAddDevice = HdaAddDevice;
    gHdaDriver.MajorFunction[IRP_MJ_PNP] = HdaDispatchPnp;
    return NTLSTATUS_SUCCESS;
}

static DRIVER_OBJECT* const gModuleIoDrivers[] = { &gHdaDriver };
NTL_DECLARE_IO_MODULE(DriverInitialize, gModuleIoDrivers);