// hal_main.c - HAL initialisation and HalGetApi
#include <ntlite/ddi/hal_api.h>
#include <kernel/hal/hal_apic.h>
#include <string.h>
#include <stdint.h>
#include <stdarg.h>
#include <intrin.h>
#pragma intrinsic(_disable)
#include <common/printf.h>

void HalScreenInitialize(void* (*AllocatePool)(size_t size), uintptr_t fb, int w, int h, int stride,
                        HAL_SCREEN* out);
KSTATUS HalTimerInitialize(uintptr_t lapicPhys, uintptr_t hpetPhys, HAL_TIMER *out);
void HalPciInitialize(HAL_PCI* out);
extern void HalSetDebugEnabled(int v);

// Defined in mm_virt.c - HAL needs it for MMIO window access.
static uint64_t gPhysBase = 0xFFFF880000000000ULL;

HAL_API gApi;
static int gInitialised = 0;
extern void (*gTimerCallback)(void);

// ---------------------------------------------------------------------------
// ACPI table walker - minimal, just enough to find APIC and HPET addresses.
// ---------------------------------------------------------------------------
#define ACPI_SIG(a,b,c,d) ((uint32_t)(a)|((uint32_t)(b)<<8)|((uint32_t)(c)<<16)|((uint32_t)(d)<<24))

#pragma pack(push, 1)
#pragma warning(disable : 4200)

// ACPI RSDP (Root System Description Pointer)
typedef struct {
    char Signature[8];
    uint8_t Checksum;
    char OEMID[6];
    uint8_t Revision;
    uint32_t RsdtAddress;
    uint32_t Length;
    uint64_t XsdtAddress;
    uint8_t ExtendedChecksum;
    uint8_t Padding[3];
} RSDP;

// ACPI SDT Header
typedef struct {
    char Signature[4];
    uint32_t Length;
    uint8_t Revision;
    uint8_t Checksum;
    char OEMID[6];
    char OEMTableID[8];
    uint32_t OEMRevision;
    uint32_t CreatorID;
    uint32_t CreatorRevision;
} ACPISDTHeader;

// RSDT (Root System Description Table)
typedef struct {
    ACPISDTHeader Header;
    uint32_t PointerToOtherSDT[];
} RSDT;

// XSDT (Extended System Description Table)
typedef struct {
    ACPISDTHeader Header;
    uint64_t PointerToOtherSDT[];
} XSDT;

// MADT (Multiple APIC Description Table)
typedef struct {
    ACPISDTHeader Header;
    uint32_t LocalAPICAddress;
    uint32_t Flags;
    uint8_t Entries[];
} MADT;

#pragma pack(pop)

static uintptr_t gLapicPhys  = 0xFEE00000;
static uintptr_t gIoapicPhys = 0xFEC00000;
static uintptr_t gHpetPhys   = 0xFED00000;
static HAL_CPU_TOPOLOGY gCpuTopology;

// ACPI power-off state (filled in by ParseFadt / ScanDsdtForS5).
// Do not guess SLP_TYP on real hardware: these values are firmware-specific.
static uint16_t gPm1aCntBlk = 0;
static uint16_t gPm1bCntBlk = 0;
static uint8_t  gSlpTypA    = 0;
static uint8_t  gSlpTypB    = 0;
static int      gS5Found    = 0;

int debugEnabled = 0;
int _fltused; // MSVC floating-point helper anchor
static void hal_putchar(char c) {
    gApi.Screen.PrintChar(c, 0xCCCCCC, 0x012456);
}
void HalDbgPrint(const char* fmt, ...) {
    if (!debugEnabled) return;

    va_list args;
    va_start(args, fmt);
    vprintf(fmt, args);
    va_end(args);
}

static void *HalPhysToVirt(uint64_t phys) {
    return (void *)(phys + gPhysBase);
}

// FADT (Fixed ACPI Description Table) offsets used here.  ACPI 2.0+
// systems may put the real DSDT and PM1 control registers in the extended
// X_* fields, so do not rely only on the legacy 32-bit fields.
#define FADT_DSDT_OFFSET           40u
#define FADT_PM1A_CNT_BLK_OFFSET   64u
#define FADT_PM1B_CNT_BLK_OFFSET   68u
#define FADT_X_DSDT_OFFSET         140u
#define FADT_X_PM1A_CNT_BLK_OFFSET 172u
#define FADT_X_PM1B_CNT_BLK_OFFSET 184u

#pragma pack(push, 1)
typedef struct {
    uint8_t  AddressSpaceId;
    uint8_t  RegisterBitWidth;
    uint8_t  RegisterBitOffset;
    uint8_t  AccessSize;
    uint64_t Address;
} ACPI_GAS;
#pragma pack(pop)

static uint32_t AcpiReadU32(const void* table, uint32_t tableLen, uint32_t off) {
    uint32_t value = 0;
    if (off + sizeof(value) <= tableLen)
        memcpy(&value, (const uint8_t*)table + off, sizeof(value));
    return value;
}

static uint64_t AcpiReadU64(const void* table, uint32_t tableLen, uint32_t off) {
    uint64_t value = 0;
    if (off + sizeof(value) <= tableLen)
        memcpy(&value, (const uint8_t*)table + off, sizeof(value));
    return value;
}

static uint16_t AcpiReadGasIoPort(const void* table, uint32_t tableLen, uint32_t off) {
    ACPI_GAS gas;
    if (off + sizeof(gas) > tableLen) return 0;

    memcpy(&gas, (const uint8_t*)table + off, sizeof(gas));

    // 0x01 == System I/O.  HalShutdown currently writes PM1_CNT with outw(),
    // so ignore MMIO GAS entries until the HAL grows MMIO register access here.
    if (gas.AddressSpaceId != 0x01 || gas.Address == 0 || gas.Address > 0xFFFF)
        return 0;

    return (uint16_t)gas.Address;
}

static int AmlParsePkgLength(const uint8_t* p, const uint8_t* end,
                             const uint8_t** body, const uint8_t** pkgEnd) {
    if (p >= end) return 0;

    const uint8_t* start = p;
    uint8_t first = *p++;
    uint32_t follow = first >> 6;
    uint32_t length;

    if (follow == 0) {
        length = first & 0x3F;
    } else {
        length = first & 0x0F;
        for (uint32_t i = 0; i < follow; i++) {
            if (p >= end) return 0;
            length |= ((uint32_t)*p++) << (4 + i * 8);
        }
    }

    if (start + length < p || start + length > end) return 0;
    *body = p;
    *pkgEnd = start + length;
    return 1;
}

static const uint8_t* AmlParseNameString(const uint8_t* p, const uint8_t* end,
                                         int* lastSegmentIsS5) {
    uint32_t segments = 1;
    *lastSegmentIsS5 = 0;

    // Optional root or parent prefixes: '\\', '^', '^^', ...
    while (p < end && (*p == 0x5C || *p == 0x5E)) p++;
    if (p >= end) return NULL;

    if (*p == 0x00) {          // NullName
        return p + 1;
    } else if (*p == 0x2E) {   // DualNamePrefix
        p++;
        segments = 2;
    } else if (*p == 0x2F) {   // MultiNamePrefix
        p++;
        if (p >= end) return NULL;
        segments = *p++;
    }

    for (uint32_t s = 0; s < segments; s++) {
        if (p + 4 > end) return NULL;
        if (s == segments - 1 && memcmp(p, "_S5_", 4) == 0)
            *lastSegmentIsS5 = 1;
        p += 4;
    }

    return p;
}

static int AmlReadInteger(const uint8_t** pp, const uint8_t* end, uint8_t* out) {
    const uint8_t* p = *pp;
    if (p >= end) return 0;

    switch (*p++) {
    case 0x00: // ZeroOp
        *out = 0;
        break;
    case 0x01: // OneOp
        *out = 1;
        break;
    case 0x0A: // BytePrefix
        if (p >= end) return 0;
        *out = *p++;
        break;
    case 0x0B: // WordPrefix
        if (p + 2 > end) return 0;
        *out = p[0];
        p += 2;
        break;
    case 0x0C: // DWordPrefix
        if (p + 4 > end) return 0;
        *out = p[0];
        p += 4;
        break;
    case 0x0E: // QWordPrefix
        if (p + 8 > end) return 0;
        *out = p[0];
        p += 8;
        break;
    default:
        return 0;
    }

    *pp = p;
    return 1;
}

// Scan DSDT AML for the _S5_ Name object and extract SLP_TYP values.  Real
// firmware commonly encodes the name as "\\_S5_" instead of just "_S5_", and
// package integers may be ZeroOp/OneOp/Byte/Word/DWord/QWord constants.
static void ScanDsdtForS5(const uint8_t* aml, uint32_t len) {
    const uint8_t* end = aml + len;

    for (const uint8_t* cur = aml; cur < end; cur++) {
        if (*cur != 0x08) continue; // NameOp

        const uint8_t* p = cur + 1;
        int isS5 = 0;
        p = AmlParseNameString(p, end, &isS5);
        if (!p || !isS5 || p >= end || *p != 0x12) continue; // PackageOp

        const uint8_t* body;
        const uint8_t* pkgEnd;
        if (!AmlParsePkgLength(p + 1, end, &body, &pkgEnd)) continue;
        if (body >= pkgEnd) continue;

        uint8_t elements = *body++;
        if (elements < 2) continue;

        uint8_t typA, typB;
        if (!AmlReadInteger(&body, pkgEnd, &typA)) continue;
        if (!AmlReadInteger(&body, pkgEnd, &typB)) continue;

        gSlpTypA = typA & 0x07;
        gSlpTypB = typB & 0x07;
        gS5Found = 1;

        HalDbgPrint("HAL: _S5_ SLP_TYP_A=0x%02x SLP_TYP_B=0x%02x\n",
                    gSlpTypA, gSlpTypB);
        return;
    }

    HalDbgPrint("HAL: _S5_ not found in DSDT; ACPI shutdown disabled\n");
}

static int validate_checksum(void* table, uint32_t length) {
    uint8_t sum = 0;
    uint8_t* bytes = (uint8_t*)table;
    for (uint32_t i = 0; i < length; i++) {
        sum += bytes[i];
    }
    return sum == 0;
}

static void ParseFadt(const ACPISDTHeader* header) {
    const void* fadt = header;
    uint32_t len = header->Length;

    gPm1aCntBlk = (uint16_t)AcpiReadU32(fadt, len, FADT_PM1A_CNT_BLK_OFFSET);
    gPm1bCntBlk = (uint16_t)AcpiReadU32(fadt, len, FADT_PM1B_CNT_BLK_OFFSET);

    if (!gPm1aCntBlk)
        gPm1aCntBlk = AcpiReadGasIoPort(fadt, len, FADT_X_PM1A_CNT_BLK_OFFSET);
    if (!gPm1bCntBlk)
        gPm1bCntBlk = AcpiReadGasIoPort(fadt, len, FADT_X_PM1B_CNT_BLK_OFFSET);

    uint64_t dsdtPhys = AcpiReadU64(fadt, len, FADT_X_DSDT_OFFSET);
    if (!dsdtPhys)
        dsdtPhys = AcpiReadU32(fadt, len, FADT_DSDT_OFFSET);

    HalDbgPrint("HAL: FADT PM1a_CNT=0x%04x PM1b_CNT=0x%04x DSDT=0x%016llx\n",
                gPm1aCntBlk, gPm1bCntBlk, dsdtPhys);

    if (dsdtPhys) {
        const ACPISDTHeader* dsdt = (const ACPISDTHeader*)(dsdtPhys + gPhysBase);
        if (memcmp(dsdt->Signature, "DSDT", 4) == 0 &&
            validate_checksum((void*)dsdt, dsdt->Length)) {
            const uint8_t* aml = (const uint8_t*)(dsdt + 1);
            ScanDsdtForS5(aml, dsdt->Length - (uint32_t)sizeof(ACPISDTHeader));
        } else {
            HalDbgPrint("HAL: DSDT invalid or checksum failed\n");
        }
    }
}


static void HalAddCpuRecord(uint32_t uid, uint32_t lapicId, uint32_t flags) {
    if (gCpuTopology.CpuCount >= HAL_MAX_CPUS)
        return;

    for (uint32_t i = 0; i < gCpuTopology.CpuCount; i++) {
        if (gCpuTopology.Cpus[i].LapicId == lapicId)
            return;
    }

    HAL_CPU_INFO* cpu = &gCpuTopology.Cpus[gCpuTopology.CpuCount++];
    memset(cpu, 0, sizeof(*cpu));
    cpu->ProcessorUid = uid;
    cpu->LapicId = lapicId;
    cpu->Flags = flags;
    cpu->Enabled = (flags & 1u) ? 1u : 0u;
}

static void ParseMadt(MADT* madt) {
    gLapicPhys = madt->LocalAPICAddress;
    uint8_t* p   = (uint8_t *)(madt + 1);
    uint8_t* end = (uint8_t *)madt + madt->Header.Length;
    while (p + 2 <= end) {
        uint8_t type = p[0], len = p[1];
        if (!len || p + len > end) break;

        if (type == 0 && len >= 8) {  // Processor Local APIC
            uint32_t uid = p[2];
            uint32_t lapicId = p[3];
            uint32_t flags = *(uint32_t *)(p + 4);
            if (flags & 1u)
                HalAddCpuRecord(uid, lapicId, flags);
        } else if (type == 1 && len >= 12) {  // I/O APIC entry
            gIoapicPhys = *(uint32_t *)(p + 4);
        } else if (type == 5 && len >= 12) {  // Local APIC Address Override
            gLapicPhys = (uintptr_t)*(uint64_t *)(p + 4);
        } else if (type == 9 && len >= 16) {  // Processor Local x2APIC
            uint32_t lapicId = *(uint32_t *)(p + 4);
            uint32_t flags = *(uint32_t *)(p + 8);
            uint32_t uid = *(uint32_t *)(p + 12);
            if (flags & 1u)
                HalAddCpuRecord(uid, lapicId, flags);
        }

        p += len;
    }
}

static void WalkAcpi(uintptr_t rsdpPhys) {
    if (!rsdpPhys) return;

    // Map RSDP into virtual memory
    RSDP* rsdp = (RSDP*)(rsdpPhys + gPhysBase);

    if (memcmp(rsdp->Signature, "RSD PTR ", 8) != 0) {
        HalDbgPrint("HAL: Invalid RSDP signature!\n");
        return;
    }

    ACPISDTHeader* sdtHeader = NULL;
    uint32_t count = 0;
    int isXsdt = 0;

    // Safely find and validate the XSDT (ACPI 2.0+)
    if (rsdp->Revision >= 2 && rsdp->XsdtAddress != 0) {
        XSDT* xsdt = (XSDT*)(rsdp->XsdtAddress + gPhysBase);
        
        if (validate_checksum(xsdt, xsdt->Header.Length) && 
            memcmp(xsdt->Header.Signature, "XSDT", 4) == 0) {
            
            sdtHeader = &xsdt->Header;
            count = (xsdt->Header.Length - sizeof(ACPISDTHeader)) / 8;
            isXsdt = 1;
        }
    }

    // Fallback to safely finding and validating the RSDT (ACPI 1.0)
    if (sdtHeader == NULL && rsdp->RsdtAddress != 0) {
        RSDT* rsdt = (RSDT*)((uintptr_t)rsdp->RsdtAddress + gPhysBase);
        
        if (validate_checksum(rsdt, rsdt->Header.Length) && 
            memcmp(rsdt->Header.Signature, "RSDT", 4) == 0) {
            
            sdtHeader = &rsdt->Header;
            count = (rsdt->Header.Length - sizeof(ACPISDTHeader)) / 4;
            isXsdt = 0;
        }
    }

    if (sdtHeader == NULL) {
        HalDbgPrint("HAL: No valid XSDT or RSDT found!\n");
        return;
    }

    // Walk the tables
    for (uint32_t i = 0; i < count; i++) {
        uintptr_t tablePhys = 0;
        
        if (isXsdt) {
            tablePhys = ((XSDT*)sdtHeader)->PointerToOtherSDT[i];
        } else {
            tablePhys = ((RSDT*)sdtHeader)->PointerToOtherSDT[i];
        }

        // Paranoid alignment check to reject firmware garbage pointers
        if (!tablePhys || (tablePhys & 3) != 0) continue;

        ACPISDTHeader* header = (ACPISDTHeader*)(tablePhys + gPhysBase);

        // Ensure the table isn't corrupt before parsing
        if (!validate_checksum(header, header->Length)) continue;

        // Parse MADT
        if (memcmp(header->Signature, "APIC", 4) == 0) {
            ParseMadt((MADT*)header);
        }
        // Parse FADT (signature "FACP") for PM1 control ports and DSDT pointer
        else if (memcmp(header->Signature, "FACP", 4) == 0) {
            ParseFadt(header);
        }
        // Parse HPET
        else if (memcmp(header->Signature, "HPET", 4) == 0) {
            // HPET base address block is located at offset 44 
            // (which is 8 bytes past the 36-byte ACPISDTHeader)
            gHpetPhys = *(uint64_t*)((uint8_t*)header + 44) & ~0xFFFULL;
        }
    }
}

// ---------------------------------------------------------------------------
// HalShutdown - ACPI S5 soft power-off. Does not return.
// ---------------------------------------------------------------------------
static void HalShutdown(void) {
    // Silence! Stop the world!
    _disable();

    if (!gS5Found || !gPm1aCntBlk) {
        // Refuse to guess.  Guessing SLP_TYP is exactly how real hardware gets
        // put into S3/S4 sleep instead of S5 soft-off.
        while (true) __halt();
    }

    // SLP_EN (bit 13) | SLP_TYPx (bits 12:10)
    uint16_t valA = (uint16_t)(((uint16_t)gSlpTypA << 10) | (1 << 13));
    uint16_t valB = (uint16_t)(((uint16_t)gSlpTypB << 10) | (1 << 13));

    __outword(gPm1aCntBlk, valA);
    if (gPm1bCntBlk) __outword(gPm1bCntBlk, valB);

    // Die!
    while (true) __halt();
}

// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS HalGetProcessorTopology(HAL_CPU_TOPOLOGY* out) {
    if (!out) return STATUS_INVALID_PARAMETER;

    memset(out, 0, sizeof(*out));
    *out = gCpuTopology;

    if (!out->CpuCount) {
        out->CpuCount = 1;
        out->BspLapicId = HalGetBspLapicId();
        out->Cpus[0].LapicId = out->BspLapicId;
        out->Cpus[0].ProcessorUid = 0;
        out->Cpus[0].Flags = 1;
        out->Cpus[0].Enabled = 1;
        out->Cpus[0].IsBsp = 1;
        return STATUS_SUCCESS;
    }

    uint32_t bsp = HalGetBspLapicId();
    out->BspLapicId = bsp;
    int foundBsp = 0;
    for (uint32_t i = 0; i < out->CpuCount; i++) {
        out->Cpus[i].IsBsp = (out->Cpus[i].LapicId == bsp) ? 1u : 0u;
        if (out->Cpus[i].IsBsp) foundBsp = 1;
    }

    if (!foundBsp && out->CpuCount < HAL_MAX_CPUS) {
        HAL_CPU_INFO* cpu = &out->Cpus[out->CpuCount++];
        memset(cpu, 0, sizeof(*cpu));
        cpu->LapicId = bsp;
        cpu->ProcessorUid = bsp;
        cpu->Flags = 1;
        cpu->Enabled = 1;
        cpu->IsBsp = 1;
    }
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS HalInitialize(const HAL_INIT_PARAMS* initParams) {
    if (gInitialised) return STATUS_ALREADY_EXISTS;
    gPhysBase = initParams->PhysicalWindowBase;
    gApi.PhysToVirt = HalPhysToVirt;
    debugEnabled = initParams->Debug;

    // Screen first so DbgPrint works immediately.
    HalScreenInitialize(initParams->AllocatePool, initParams->FramebufferBase + initParams->PhysicalWindowBase,
                        (int)initParams->FramebufferWidth,
                        (int)initParams->FramebufferHeight,
                        (int)initParams->FramebufferStride,
                        &gApi.Screen);

    printf_set_putchar(hal_putchar);

    HalDbgPrint("HAL: screen %dx%d\n",
            initParams->FramebufferWidth, initParams->FramebufferHeight);

    // Parse ACPI to find LAPIC / IOAPIC / HPET addresses.
    if (initParams->RsdpPhysical) WalkAcpi(initParams->RsdpPhysical);
    HalDbgPrint("HAL: LAPIC=%016llx IOAPIC=%016llx HPET=%016llx\n",
            gLapicPhys, gIoapicPhys, gHpetPhys);

    HalApicInitialize(gLapicPhys, gIoapicPhys);
    HalDbgPrint("HAL: local APIC mode=%s BSP_ID=%u\n",
                HalApicUsingX2Apic() ? "x2APIC" : "xAPIC",
                HalGetBspLapicId());

    KSTATUS ts = HalTimerInitialize(gLapicPhys, gHpetPhys, &gApi.Timer);
    if (!KSUCCESS(ts)) HalDbgPrint("HAL: HPET unavailable, using fallback\n");

    HalPciInitialize(&gApi.Pci);
    gApi.Intc.RouteIrq    = HalApicRouteIrq;
    gApi.Intc.MaskIrq     = HalApicMaskIrq;
    gApi.Intc.UnmaskIrq   = HalApicUnmaskIrq;
    gApi.Intc.SendEoi     = HalApicSendEoi;

    gApi.Power.Shutdown   = HalShutdown;

    gApi.Version = HAL_API_VERSION;
    gInitialised = 1;

    // The 8042 can leave output bytes pending while the keyboard/mouse drivers
    // are being brought up with interrupts temporarily disabled. Drain the
    // controller before returning from HAL init so stale status/data bytes do
    // not suppress the next keyboard or mouse interrupt edge.
    for (int i = 0; i < 256; i++) {
        if (__inbyte(0x64) & 0x01) {
            __inbyte(0x60); // Consume
        }
        gApi.Timer.Stall(1); // Micro-stalls
    }

    HalDbgPrint("HAL: initialisation complete.\n");
    return STATUS_SUCCESS;
}

__declspec(dllexport) void HalRegisterScheduler(void (*pKeSchedulerTick)(void)) {
    gTimerCallback = pKeSchedulerTick;
}

__declspec(dllexport) const HAL_API* HalGetApi(void) { return &gApi; }