#include <ntlite/ddk/ddk.h>
#include <kernel/hal/hal_api.h>
#include <kernel/io_ddk.h>
#include <kernel/io_dma.h>
#include <kernel/ke_audio.h>
#include <kernel/ke_irq.h>
#include <kernel/ke_vectors.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_irql.h>
#include <kernel/mm_pte.h>

static const HAL_API* gHal;
static uintptr_t gPhysicalWindowBase;

void IoInitializeDdk(const HAL_API* hal, uintptr_t physicalWindowBase) {
    gHal = hal;
    gPhysicalWindowBase = physicalWindowBase;
}

__declspec(dllexport) uintptr_t KeCurrentThreadIdentity(void) {
    return (uintptr_t)KeGetCurrentThread();
}

__declspec(dllexport) uint64_t KeQueryTicks(void) {
    return gHal && gHal->Timer.QueryTicks ? gHal->Timer.QueryTicks() : 0;
}

__declspec(dllexport) void KeStallExecution(uint32_t microseconds) {
    if (gHal && gHal->Timer.Stall) gHal->Timer.Stall(microseconds);
}

__declspec(dllexport) uint32_t KePciReadConfig32(uint16_t segment, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset) {
    return gHal && gHal->Pci.Read32 ? gHal->Pci.Read32(segment, bus, device, function, offset) : UINT32_MAX;
}

__declspec(dllexport) void KePciWriteConfig32(uint16_t segment, uint8_t bus, uint8_t device, uint8_t function, uint16_t offset, uint32_t value) {
    if (gHal && gHal->Pci.Write32) gHal->Pci.Write32(segment, bus, device, function, offset, value);
}

__declspec(dllexport) void* MmDdkPhysToVirt(uint64_t physicalAddress) {
    if (!gPhysicalWindowBase || physicalAddress > UINTPTR_MAX - gPhysicalWindowBase)
        return (void*)0;
    return (void*)(gPhysicalWindowBase + (uintptr_t)physicalAddress);
}

__declspec(dllexport) KSTATUS MmDdkMapPage(uintptr_t virtualAddress, uint64_t physicalAddress, uint64_t flags, NTL_DDK_MAP_MODE mode) {
    return MmPteMapPage(virtualAddress, physicalAddress, flags,
                        mode == NtlDdkMapReplaceExisting ? MmPteMapReplaceExisting : MmPteMapFailIfPresent);
}

// Each slot is a recyclable 2 MiB VA reservation; no pool allocation is needed
// when reserving a device mapping. This aperture stays separate from PCI BARs.
#define DDK_IO_VA_BASE  0xFFFF980000000000ULL
#define DDK_IO_SLOT_SIZE (2ULL * 1024ULL * 1024ULL)
#define DDK_IO_SLOT_COUNT 256u

typedef struct DDK_IO_MAPPING {
    uint8_t State;
    uint32_t Pages;
    uintptr_t Offset;
    size_t Length;
} DDK_IO_MAPPING;

static SPINLOCK gDdkIoMapLock = SPINLOCK_INIT;
static DDK_IO_MAPPING gDdkIoMaps[DDK_IO_SLOT_COUNT];

// The caller owns the mapping until an exact-length unmap completes.
__declspec(dllexport) KSTATUS MmDdkMapIoMemory(uint64_t physicalAddress, size_t length, void** outVirtual) {
    if (!outVirtual) return NTLSTATUS_INVALID_PARAMETER;
    *outVirtual = NULL;
    if (!KeCanWait()) return NTLSTATUS_NOT_SUPPORTED;
    if (!length || physicalAddress > UINT64_MAX - length) return NTLSTATUS_INVALID_PARAMETER;
    uintptr_t offset = (uintptr_t)(physicalAddress & (PAGE_SIZE - 1u));
    if (length > DDK_IO_SLOT_SIZE - offset) return NTLSTATUS_INVALID_PARAMETER;
    uint32_t pages = (uint32_t)((offset + length + PAGE_SIZE - 1u) / PAGE_SIZE);
    if (pages == 0 || pages > DDK_IO_SLOT_SIZE / PAGE_SIZE) return NTLSTATUS_INVALID_PARAMETER;

    // Reject a UC alias of RAM already covered by the firmware's WB/WC
    // direct map. Real device registers have no direct mapping or an UC one.
    for (uint32_t i = 0; i < pages; i++) {
        uint64_t frame = (physicalAddress - offset) + (uint64_t)i * PAGE_SIZE;
        if (gPhysicalWindowBase && frame <= UINTPTR_MAX - gPhysicalWindowBase) {
            PHYSICAL_ADDRESS existing = 0;
            uint64_t existingFlags = 0;
            if (KSUCCESS(MmPteQueryPage(gPhysicalWindowBase + (uintptr_t)frame,
                                        &existing, &existingFlags))) {
                if ((existing & PAGE_MASK) != frame ||
                    (existingFlags & PAGE_CACHE_MASK) != PAGE_CACHE_UC)
                    return NTLSTATUS_NOT_SUPPORTED;
            }
        }
    }

    uintptr_t flags = KeAcquireSpinlock(&gDdkIoMapLock);
    uint32_t slot;
    for (slot = 0; slot < DDK_IO_SLOT_COUNT; slot++) {
        if (!gDdkIoMaps[slot].State) {
            gDdkIoMaps[slot].State = 1;
            gDdkIoMaps[slot].Length = length;
            gDdkIoMaps[slot].Offset = offset;
            gDdkIoMaps[slot].Pages = pages;
            break;
        }
    }
    KeReleaseSpinlock(&gDdkIoMapLock, flags);
    if (slot == DDK_IO_SLOT_COUNT) return NTLSTATUS_NO_MEMORY;

    uintptr_t base = (uintptr_t)(DDK_IO_VA_BASE + (uint64_t)slot * DDK_IO_SLOT_SIZE);
    uint64_t physBase = physicalAddress - offset;
    uint32_t mapped = 0;
    KSTATUS status = NTLSTATUS_SUCCESS;
    for (; mapped < pages; mapped++) {
        status = MmPteMapPage(base + (uintptr_t)mapped * PAGE_SIZE,
                              physBase + (uint64_t)mapped * PAGE_SIZE,
                              PAGE_PRESENT | PAGE_WRITE | PAGE_CACHE_UC | PAGE_NX,
                              MmPteMapFailIfPresent);
        if (!KSUCCESS(status)) break;
    }
    if (mapped != pages) {
        while (mapped) (void)MmPteUnmapPage(base + (uintptr_t)(--mapped) * PAGE_SIZE);
        flags = KeAcquireSpinlock(&gDdkIoMapLock);
        gDdkIoMaps[slot].State = 0;
        KeReleaseSpinlock(&gDdkIoMapLock, flags);
        return status;
    }
    flags = KeAcquireSpinlock(&gDdkIoMapLock);
    gDdkIoMaps[slot].State = 2;
    KeReleaseSpinlock(&gDdkIoMapLock, flags);
    *outVirtual = (void*)(base + offset);
    return NTLSTATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS MmDdkUnmapIoMemory(void* virtualAddress, size_t length) {
    if (!KeCanWait()) return NTLSTATUS_NOT_SUPPORTED;
    if (!virtualAddress || !length) return NTLSTATUS_INVALID_PARAMETER;
    uintptr_t address = (uintptr_t)virtualAddress;
    if (address < DDK_IO_VA_BASE ||
        (uint64_t)(address - DDK_IO_VA_BASE) >= DDK_IO_SLOT_SIZE * DDK_IO_SLOT_COUNT)
        return NTLSTATUS_INVALID_PARAMETER;
    uint32_t slot = (uint32_t)((address - DDK_IO_VA_BASE) / DDK_IO_SLOT_SIZE);
    uintptr_t base = (uintptr_t)(DDK_IO_VA_BASE + (uint64_t)slot * DDK_IO_SLOT_SIZE);
    uintptr_t flags = KeAcquireSpinlock(&gDdkIoMapLock);
    DDK_IO_MAPPING* map = &gDdkIoMaps[slot];
    if (map->State != 2 || map->Length != length || address != base + map->Offset) {
        KeReleaseSpinlock(&gDdkIoMapLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    map->State = 3;
    uint32_t pages = map->Pages;
    KeReleaseSpinlock(&gDdkIoMapLock, flags);
    KSTATUS result = NTLSTATUS_SUCCESS;
    for (uint32_t i = 0; i < pages; i++) {
        KSTATUS status = MmPteUnmapPage(base + (uintptr_t)i * PAGE_SIZE);
        if (!KSUCCESS(status)) result = status;
    }
    flags = KeAcquireSpinlock(&gDdkIoMapLock);
    // Preserve the slot after a failed unmap to prevent stale PTE reuse.
    if (KSUCCESS(result)) map->State = 0;
    KeReleaseSpinlock(&gDdkIoMapLock, flags);
    return result;
}

__declspec(dllexport) KSTATUS IoDmaGetAdapter(struct DEVICE_OBJECT* device, const IO_DMA_REQUIREMENTS* requirements, IO_DMA_ADAPTER** outAdapter) {
    return IoGetDmaDdk()->GetAdapter(device, requirements, outAdapter);
}

__declspec(dllexport) void IoDmaPutAdapter(IO_DMA_ADAPTER* adapter) {
    IoGetDmaDdk()->PutAdapter(adapter);
}

__declspec(dllexport) KSTATUS IoDmaAllocateCommonBuffer(IO_DMA_ADAPTER* adapter, size_t size, size_t alignment, IO_DMA_BUFFER** outBuffer) {
    return IoGetDmaDdk()->AllocateCommonBuffer(adapter, size, alignment, outBuffer);
}

__declspec(dllexport) void IoDmaFreeCommonBuffer(IO_DMA_BUFFER* buffer) {
    IoGetDmaDdk()->FreeCommonBuffer(buffer);
}

__declspec(dllexport) void* IoDmaGetCommonBufferAddress(const IO_DMA_BUFFER* buffer) {
    return IoGetDmaDdk()->GetCommonBufferAddress(buffer);
}

__declspec(dllexport) uint64_t IoDmaGetCommonBufferDeviceAddress(const IO_DMA_BUFFER* buffer) {
    return IoGetDmaDdk()->GetCommonBufferDeviceAddress(buffer);
}

__declspec(dllexport) size_t IoDmaGetCommonBufferSize(const IO_DMA_BUFFER* buffer) {
    return IoGetDmaDdk()->GetCommonBufferSize(buffer);
}

__declspec(dllexport) KSTATUS IoDmaPinKernelMemory(void* address, size_t size, IO_DMA_DIRECTION direction, MM_PINNED_MEMORY** outMemory) {
    return IoGetDmaDdk()->PinKernelMemory(address, size, direction, outMemory);
}

__declspec(dllexport) KSTATUS IoDmaPinUserMemory(struct PROCESS* process, void* address, size_t size, IO_DMA_DIRECTION direction, MM_PINNED_MEMORY** outMemory) {
    return IoGetDmaDdk()->PinUserMemory(process, address, size, direction, outMemory);
}

__declspec(dllexport) void IoDmaUnpinMemory(MM_PINNED_MEMORY* memory) {
    IoGetDmaDdk()->UnpinMemory(memory);
}

__declspec(dllexport) KSTATUS IoDmaMapPinnedMemory(IO_DMA_ADAPTER* adapter, MM_PINNED_MEMORY* memory, size_t offset, size_t length, IO_DMA_DIRECTION direction, uint32_t flags, IO_DMA_MAPPING** outMapping) {
    return IoGetDmaDdk()->MapPinnedMemory(adapter, memory, offset, length, direction, flags, outMapping);
}

__declspec(dllexport) void IoDmaUnmapMemory(IO_DMA_MAPPING* mapping) {
    IoGetDmaDdk()->UnmapMemory(mapping);
}

__declspec(dllexport) KSTATUS IoDmaSynchronizeForDevice(IO_DMA_MAPPING* mapping) {
    return IoGetDmaDdk()->SynchronizeForDevice(mapping);
}

__declspec(dllexport) KSTATUS IoDmaSynchronizeForCpu(IO_DMA_MAPPING* mapping) {
    return IoGetDmaDdk()->SynchronizeForCpu(mapping);
}

__declspec(dllexport) uint32_t IoDmaGetSegmentCount(const IO_DMA_MAPPING* mapping) {
    return IoGetDmaDdk()->GetSegmentCount(mapping);
}

__declspec(dllexport) const IO_DMA_SEGMENT* IoDmaGetSegments(const IO_DMA_MAPPING* mapping) {
    return IoGetDmaDdk()->GetSegments(mapping);
}

__declspec(dllexport) void IoDmaQueryStats(IO_DMA_STATS* outStats) {
    IoGetDmaDdk()->QueryStats(outStats);
}

__declspec(dllexport) KSTATUS IoAudioRegisterEndpoint(const IO_AUDIO_ENDPOINT_REGISTRATION* registration, IO_AUDIO_ENDPOINT** outEndpoint) {
    return KeAudioGetDdk()->RegisterEndpoint(registration, outEndpoint);
}

__declspec(dllexport) void IoAudioUnregisterEndpoint(IO_AUDIO_ENDPOINT* endpoint, IO_AUDIO_REMOVAL_MODE mode) {
    KeAudioGetDdk()->UnregisterEndpoint(endpoint, mode);
}

__declspec(dllexport) void IoAudioNotifyPeriod(IO_AUDIO_ENDPOINT* endpoint, uint32_t completedPeriodIndex, uint64_t presentedFrames, uint64_t hardwarePositionFrames, uint64_t interruptCount, uint64_t timestampTicks) {
    KeAudioGetDdk()->NotifyPeriod(endpoint, completedPeriodIndex, presentedFrames, hardwarePositionFrames, interruptCount, timestampTicks);
}

__declspec(dllexport) void IoAudioSetEndpointAvailable(IO_AUDIO_ENDPOINT* endpoint, bool available) {
    KeAudioGetDdk()->SetEndpointAvailable(endpoint, available);
}

__declspec(dllexport) void IoAudioReportEndpointFailure(IO_AUDIO_ENDPOINT* endpoint, KSTATUS status, uint32_t flags) {
    KeAudioGetDdk()->ReportEndpointFailure(endpoint, status, flags);
}