/*
    * UserRange.cpp
    * Cross-CPU invalidation and teardown of user address-space mappings
    * Copyright (c) 2026 Daniel Hammer
    *
    * Split out of Ipc.cpp: this is paging and SMP work with no dependency on
    * any IPC object or handle pool, and it lived there only by history.
*/

#include "UserRange.hpp"

#include <Sched/Scheduler.hpp>
#include <Memory/PageFrameAllocator.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/Paging.hpp>
#include <CppLib/Spinlock.hpp>
#include <Hal/Apic/Apic.hpp>
#include <Hal/Apic/Interrupts.hpp>
#include <Hal/SmpBoot.hpp>
#include <Common/Panic.hpp>

namespace Memory {

    // MUST be a Mutex, never a Spinlock. ShootdownUserRange holds this while
    // waiting for remote CPUs to acknowledge the shootdown IPI, so a CPU that
    // is queued behind the holder has to stay interruptible long enough to
    // service that IPI itself. An interrupt-disabling Spinlock here deadlocks
    // every CPU contending for the lock, and the bounded-retry logic below
    // then reports it as a "target failed to acknowledge" Panic -- which reads
    // like a hardware fault rather than a lock-type regression.
    static kcp::Mutex g_tlbShootdownLock;
    static volatile uint64_t g_tlbShootdownSeq = 0;
    static volatile uint64_t g_tlbShootdownPml4 = 0;
    static volatile uint64_t g_tlbShootdownStartVa = 0;
    static volatile uint32_t g_tlbShootdownPages = 0;
    static volatile uint64_t g_tlbShootdownDone[Smp::MaxCPUs] = {};

    static bool CpuCurrentlyUsesPml4(Smp::CpuData* cpu, uint64_t pml4Phys) {
        if (cpu == nullptr || pml4Phys == 0 || cpu->currentSlot < 0) return false;

        Sched::Process* proc = Sched::GetProcessSlot(cpu->currentSlot);
        if (proc == nullptr) return false;
        if (proc->state == Sched::ProcessState::Free) return false;
        return proc->pml4Phys == pml4Phys;
    }

    static void InvalidateLocalUserRange(uint64_t startVa, uint32_t pages) {
        if (pages == 0) return;

        if (pages > 1024) {
            Memory::VMM::FlushTLB();
            return;
        }

        for (uint32_t p = 0; p < pages; p++) {
            uint64_t va = startVa + (uint64_t)p * 0x1000ULL;
            asm volatile("invlpg (%0)" :: "r"(va) : "memory");
        }
    }

    static void TlbShootdownIpiHandler(uint8_t, bool) {
        Smp::CpuData* cpu = Smp::GetCurrentCpuData();
        uint64_t seq = g_tlbShootdownSeq;
        uint64_t pml4 = g_tlbShootdownPml4;
        uint64_t startVa = g_tlbShootdownStartVa;
        uint32_t pages = g_tlbShootdownPages;

        if (CpuCurrentlyUsesPml4(cpu, pml4)) {
            InvalidateLocalUserRange(startVa, pages);
        }

        if (cpu != nullptr && cpu->cpuIndex >= 0 && cpu->cpuIndex < Smp::MaxCPUs) {
            asm volatile("" ::: "memory");
            g_tlbShootdownDone[cpu->cpuIndex] = seq;
        }
    }

    void ShootdownUserRange(uint64_t pml4Phys, uint64_t startVa, uint32_t pages) {
        if (pml4Phys == 0 || pages == 0) return;

        bool targets[Smp::MaxCPUs] = {};
        Smp::CpuData* currentCpu = Smp::GetCurrentCpuData();
        int currentCpuIndex = currentCpu ? currentCpu->cpuIndex : -1;

        g_tlbShootdownLock.Acquire();
        uint64_t seq = g_tlbShootdownSeq + 1;
        g_tlbShootdownPml4 = pml4Phys;
        g_tlbShootdownStartVa = startVa;
        g_tlbShootdownPages = pages;
        asm volatile("" ::: "memory");
        g_tlbShootdownSeq = seq;

        for (int i = 0; i < Smp::GetCpuCount(); i++) {
            Smp::CpuData* cpu = Smp::GetCpuData(i);
            if (cpu == nullptr || !cpu->started) continue;

            if (i == currentCpuIndex) {
                if (CpuCurrentlyUsesPml4(cpu, pml4Phys)) {
                    InvalidateLocalUserRange(startVa, pages);
                }
                g_tlbShootdownDone[i] = seq;
                continue;
            }

            if (!CpuCurrentlyUsesPml4(cpu, pml4Phys)) {
                g_tlbShootdownDone[i] = seq;
                continue;
            }

            targets[i] = true;
            (void)Hal::LocalApic::SendFixedIpi(cpu->lapicId,
                Hal::IRQ_VECTOR_BASE + Hal::IRQ_TLB_SHOOTDOWN);
        }

        for (int i = 0; i < Smp::GetCpuCount(); i++) {
            if (!targets[i]) continue;
            uint32_t spins = 0;
            uint32_t retries = 0;
            while (g_tlbShootdownDone[i] != seq) {
                asm volatile("pause");
                if (++spins < 1000000) continue;

                // Delivery normally completes in a handful of cycles. Retry a
                // bounded number of times in case the first IPI was lost while
                // the target changed interrupt state. Continuing without an
                // acknowledgement would let the caller free frames still
                // reachable through a remote stale TLB entry, so fail loudly
                // instead of either corrupting memory or spinning forever.
                spins = 0;
                if (++retries > 4) {
                    Panic("TLB shootdown target failed to acknowledge", nullptr);
                }
                Smp::CpuData* cpu = Smp::GetCpuData(i);
                if (cpu != nullptr && cpu->started) {
                    (void)Hal::LocalApic::SendFixedIpi(cpu->lapicId,
                        Hal::IRQ_VECTOR_BASE + Hal::IRQ_TLB_SHOOTDOWN);
                }
            }
        }

        g_tlbShootdownLock.Release();
    }

    void UnmapAndFreeUserRange(uint64_t pml4Phys, uint64_t startVa, uint64_t pages) {
        static constexpr uint32_t PagesPerChunk = 64;
        uint64_t physPages[PagesPerChunk];

        for (uint64_t base = 0; base < pages; base += PagesPerChunk) {
            uint32_t count = (uint32_t)((pages - base > PagesPerChunk)
                ? PagesPerChunk : pages - base);

            for (uint32_t i = 0; i < count; i++) {
                uint64_t pageVa = startVa + (base + i) * 0x1000ULL;
                physPages[i] = Memory::VMM::Paging::GetPhysAddr(pml4Phys, pageVa);
                Memory::VMM::Paging::UnmapUserIn(pml4Phys, pageVa);
            }

            ShootdownUserRange(pml4Phys, startVa + base * 0x1000ULL, count);

            for (uint32_t i = 0; i < count; i++) {
                if (physPages[i] != 0) {
                    Memory::g_pfa->Free((void*)Memory::HHDM(physPages[i]));
                }
            }
        }
    }

    void InitUserRange() {
        Hal::RegisterIrqHandler(Hal::IRQ_TLB_SHOOTDOWN, TlbShootdownIpiHandler);
    }

}