/*
    * SmpBoot.cpp
    * Symmetric Multiprocessing bootstrap and AP entry
    * Copyright (c) 2026 Daniel Hammer
*/

#include "SmpBoot.hpp"
#include <ACPI/CpuIdle.hpp>
#include <Hal/Apic/Apic.hpp>
#include <Hal/Apic/Interrupts.hpp>
#include <Hal/IDT.hpp>
#include <Hal/MSR.hpp>
#include <Hal/Cpu.hpp>
#include <Hal/CpuPower.hpp>
#include <Memory/Paging.hpp>
#include <Memory/PageFrameAllocator.hpp>
#include <Memory/HHDM.hpp>
#include <Terminal/Terminal.hpp>
#include <CppLib/Stream.hpp>
#include <Timekeeping/ApicTimer.hpp>
#include <Boot/BootInfo.hpp>
#include <Libraries/Memory.hpp>

// Verify assembly offsets match struct layout
static_assert(__builtin_offsetof(Smp::CpuData, selfPtr)       == CPUDATA_SELF_PTR,     "CpuData offset mismatch: selfPtr");
static_assert(__builtin_offsetof(Smp::CpuData, kernelRsp)     == CPUDATA_KERNEL_RSP,   "CpuData offset mismatch: kernelRsp");
static_assert(__builtin_offsetof(Smp::CpuData, userRspScratch)== CPUDATA_USER_RSP,     "CpuData offset mismatch: userRspScratch");
static_assert(__builtin_offsetof(Smp::CpuData, currentSlot)   == CPUDATA_CURRENT_SLOT, "CpuData offset mismatch: currentSlot");

extern "C" void SyscallEntry();

// Assembly helpers (GDT.asm)
extern "C" void LoadGDT(Hal::GDTPointer* ptr);
extern "C" void ReloadSegments();
extern "C" void LoadTR();

using namespace Kt;

namespace Smp {

    static CpuData g_cpus[MaxCPUs];
    static int g_cpuCount = 0;

    CpuData* GetCpuData(int index) {
        if (index < 0 || index >= g_cpuCount) return nullptr;
        return &g_cpus[index];
    }

    int GetCpuCount() {
        return g_cpuCount;
    }

    // ====================================================================
    // Per-CPU GDT/TSS setup
    // ====================================================================

    static void SetupPerCpuGdtTss(CpuData& cpu) {
        // Zero the TSS
        memset(&cpu.cpuTss, 0, sizeof(Hal::TSS64));
        cpu.cpuTss.iopbOffset = sizeof(Hal::TSS64);

        // Allocate a 4KB IST1 stack for Double Fault and NMI.
        // These exceptions need a known-good stack to avoid triple
        // faults when the normal kernel stack overflows.
        void* istPage = Memory::g_pfa->AllocateZeroed();
        if (istPage) {
            cpu.cpuTss.ist1 = (uint64_t)istPage + 0x1000;  // top of stack
        }

        // Copy the standard GDT layout
        cpu.cpuGdt = {
            {0xFFFF, 0, 0, 0x00, 0x00, 0},    // 0x00 Null
            {0xFFFF, 0, 0, 0x9A, 0xA0, 0},    // 0x08 KernelCode
            {0xFFFF, 0, 0, 0x92, 0xA0, 0},    // 0x10 KernelData
            {0xFFFF, 0, 0, 0xF2, 0xA0, 0},    // 0x18 UserData
            {0xFFFF, 0, 0, 0xFA, 0xA0, 0},    // 0x20 UserCode
            {0, 0, 0, 0, 0, 0},               // 0x28 TSS low
            {0, 0, 0, 0, 0, 0},               // 0x30 TSS high
        };

        // Encode the 16-byte TSS descriptor
        uint64_t base = (uint64_t)&cpu.cpuTss;
        uint32_t limit = sizeof(Hal::TSS64) - 1;

        // Low 8 bytes
        cpu.cpuGdt.TSS.LimitLow       = limit & 0xFFFF;
        cpu.cpuGdt.TSS.BaseLow        = base & 0xFFFF;
        cpu.cpuGdt.TSS.BaseMiddle      = (base >> 16) & 0xFF;
        cpu.cpuGdt.TSS.AccessByte      = 0x89;  // Present, 64-bit TSS Available
        cpu.cpuGdt.TSS.GranularityByte = (limit >> 16) & 0x0F;
        cpu.cpuGdt.TSS.BaseHigh        = (base >> 24) & 0xFF;

        // High 8 bytes (base[63:32])
        uint32_t baseUpper = (uint32_t)(base >> 32);
        cpu.cpuGdt.TSSHigh.LimitLow       = baseUpper & 0xFFFF;
        cpu.cpuGdt.TSSHigh.BaseLow        = (baseUpper >> 16) & 0xFFFF;
        cpu.cpuGdt.TSSHigh.BaseMiddle      = 0;
        cpu.cpuGdt.TSSHigh.AccessByte      = 0;
        cpu.cpuGdt.TSSHigh.GranularityByte = 0;
        cpu.cpuGdt.TSSHigh.BaseHigh        = 0;

        cpu.tss = &cpu.cpuTss;
    }

    // ====================================================================
    // Set GS base MSRs for per-CPU data
    // ====================================================================

    static constexpr uint32_t IA32_GS_BASE        = 0xC0000101;
    static constexpr uint32_t IA32_KERNEL_GS_BASE = 0xC0000102;

    static void SetGSBase(CpuData* cpu) {
        // In kernel mode, GSBASE = per-CPU data pointer.
        // KernelGSBASE = 0 (user GS base, swapped in on swapgs).
        Hal::WriteMSR(IA32_GS_BASE, (uint64_t)cpu);
        Hal::WriteMSR(IA32_KERNEL_GS_BASE, 0);
    }

    // ====================================================================
    // BSP initialization
    // ====================================================================

    void InitBsp() {
        memset(g_cpus, 0, sizeof(g_cpus));

        CpuData& bsp = g_cpus[0];
        bsp.selfPtr = (uint64_t)&bsp;
        bsp.cpuIndex = 0;
        bsp.lapicId = Hal::LocalApic::GetId();
        bsp.currentSlot = -1;
        bsp.started = true;
        bsp.hasMwait = Hal::HasMwait();

        // BSP uses the global TSS (already set up in PrepareGDT)
        bsp.tss = &Hal::g_tss;

        // Allocate IST1 stack for the BSP (for Double Fault / NMI)
        void* bspIstPage = Memory::g_pfa->AllocateZeroed();
        if (bspIstPage) {
            Hal::g_tss.ist1 = (uint64_t)bspIstPage + 0x1000;
        }

        // Set GS base for BSP
        SetGSBase(&bsp);

        g_cpuCount = 1;

        KernelLogStream(OK, "SMP") << "BSP initialized (LAPIC ID " << (uint64_t)bsp.lapicId << ")";
    }

    // ====================================================================
    // AP entry point
    // Invoked by the boot contract (montauk::boot::SmpInfo::startCpu) when
    // the AP is woken.  Receives the BootCpu for this processor; the kernel
    // stashed this CPU's CpuData* in BootCpu::extraArgument before starting
    // it.  Runs on a bootloader-provided stack.
    // ====================================================================

    static void ApEntry(montauk::boot::BootCpu* bootCpu) {
        // Find our CpuData (stashed in extraArgument by BootAPs)
        CpuData* cpu = (CpuData*)bootCpu->extraArgument;

        // --- Load per-CPU GDT ---
        Hal::GDTPointer gdtPtr {
            .Size = sizeof(Hal::BasicGDT) - 1,
            .GDTAddress = (uint64_t)&cpu->cpuGdt,
        };
        LoadGDT(&gdtPtr);
        ReloadSegments();

        // --- Load TSS ---
        LoadTR();

        // --- Load shared IDT ---
        Hal::IDTReload();

        // --- Switch to kernel page tables ---
        Memory::VMM::LoadCR3(Memory::VMM::g_paging->PML4);

        // --- Enable SSE ---
        Hal::EnableSSE();

        // --- Set GS base ---
        SetGSBase(cpu);

        // --- Initialize local APIC ---
        Hal::LocalApic::InitializeAP();

        // --- Program SYSCALL MSRs ---
        uint64_t efer = Hal::ReadMSR(Hal::IA32_EFER);
        efer |= 1;  // SCE
        Hal::WriteMSR(Hal::IA32_EFER, efer);

        uint64_t star = (0x0010ULL << 48) | (0x0008ULL << 32);
        Hal::WriteMSR(Hal::IA32_STAR, star);
        Hal::WriteMSR(Hal::IA32_LSTAR, (uint64_t)SyscallEntry);
        Hal::WriteMSR(Hal::IA32_FMASK, 0x200);

        // --- Program PAT (entry 1 = WC) ---
        Hal::InitializePAT();

        // --- Calibrate and start APIC timer ---
        Timekeeping::ApicTimerInitializeAP();

        // --- Check MWAIT support ---
        cpu->hasMwait = Hal::HasMwait();

        // --- Per-CPU power setup (C1E promotion, HWP request) ---
        Hal::CpuPower::InitializeAp();

        // --- Signal that we are online ---
        cpu->started = true;

        // --- Enable interrupts and enter idle loop ---
        asm volatile("sti");

        static volatile uint64_t s_idleMonitor = 0;
        for (;;) {
            // Pick up thermal-governor frequency changes decided by the BSP.
            Hal::CpuPower::ApplyPolicyIfChanged();
            Hal::CpuIdle::Wait(10, cpu->hasMwait, &s_idleMonitor);
        }
    }

    // ====================================================================
    // Boot all APs
    // ====================================================================

    void BootAPs(const montauk::boot::SmpInfo& smp) {
        if (smp.cpus == nullptr || smp.startCpu == nullptr) {
            KernelLogStream(WARNING, "SMP") << "No SMP info from bootloader - single CPU mode";
            return;
        }

        uint64_t cpuCount = smp.cpuCount;

        KernelLogStream(INFO, "SMP") << "Bootloader reports " << cpuCount << " CPU(s), BSP LAPIC ID "
            << (uint64_t)smp.bspLapicId;

        if (cpuCount <= 1) {
            KernelLogStream(INFO, "SMP") << "Single CPU system - no APs to boot";
            return;
        }

        if (cpuCount > (uint64_t)MaxCPUs) {
            KernelLogStream(WARNING, "SMP") << "Clamping CPU count from " << cpuCount << " to " << (uint64_t)MaxCPUs;
            cpuCount = MaxCPUs;
        }

        // Prepare all APs, then wake them all at once. This is safe because:
        // - APs use BSP's timer calibration (no PIT contention)
        // - APs don't log (no terminal contention)
        // - Each AP's init is purely local (GDT, TSS, APIC, MSRs)
        int apIndex = 1;  // BSP is index 0
        for (uint64_t i = 0; i < cpuCount; i++) {
            montauk::boot::BootCpu& info = smp.cpus[i];

            if (info.lapicId == smp.bspLapicId) continue;
            if (apIndex >= MaxCPUs) break;

            CpuData& ap = g_cpus[apIndex];
            ap.selfPtr = (uint64_t)&ap;
            ap.cpuIndex = apIndex;
            ap.lapicId = info.lapicId;
            ap.currentSlot = -1;
            ap.started = false;

            SetupPerCpuGdtTss(ap);

            // Stash this AP's CpuData* where ApEntry will recover it, then
            // ask the bootloader to wake the AP into ApEntry.  APs run in
            // parallel with each other.
            info.extraArgument = (uint64_t)&ap;
            smp.startCpu(&info, ApEntry);

            apIndex++;
        }

        g_cpuCount = apIndex;

        // Wait for all APs to come online
        for (int i = 1; i < g_cpuCount; i++) {
            volatile bool* flag = &g_cpus[i].started;
            uint64_t timeout = 100000000;
            while (!*flag && timeout > 0) {
                asm volatile("pause");
                timeout--;
            }

            if (!*flag) {
                KernelLogStream(ERROR, "SMP") << "AP " << i << " (LAPIC "
                    << (uint64_t)g_cpus[i].lapicId << ") failed to start";
            }
        }

        // Count how many actually started
        int onlineCount = 1;  // BSP
        for (int i = 1; i < g_cpuCount; i++) {
            if (g_cpus[i].started) onlineCount++;
        }

        KernelLogStream(OK, "SMP") << onlineCount << " CPU(s) online";
    }
}