/*
    * Cpu.hpp
    * aarch64 implementation of the Arch/Cpu.hpp primitives
    * Copyright (c) 2026 Daniel Hammer
*/

#pragma once
#include <cstdint>

namespace Arch {

    // PSTATE.I in the DAIF view
    static constexpr uint64_t DaifIrqMask = 1ULL << 7;

    inline void Pause() {
        asm volatile("yield" ::: "memory");
    }

    // Device-nGnRE / Normal-NC (write-combining) stores complete in order
    // once a DSB has been issued.
    inline void StoreFence() {
        asm volatile("dsb st" ::: "memory");
    }

    inline void MemoryFence() {
        asm volatile("dsb sy" ::: "memory");
    }

    // Devices are outer-shareable observers.
    inline void DmaWriteFence() {
        asm volatile("dmb oshst" ::: "memory");
    }

    inline void DmaReadFence() {
        asm volatile("dmb oshld" ::: "memory");
    }

    inline void EnableInterrupts() {
        asm volatile("msr daifclr, #2" ::: "memory");
    }

    inline void DisableInterrupts() {
        asm volatile("msr daifset, #2" ::: "memory");
    }

    inline bool InterruptsEnabled() {
        uint64_t daif;
        asm volatile("mrs %0, daif" : "=r"(daif));
        return (daif & DaifIrqMask) == 0;
    }

    inline IrqFlags SaveAndDisableInterrupts() {
        uint64_t daif;
        asm volatile("mrs %0, daif\n\tmsr daifset, #2" : "=r"(daif) :: "memory");
        return daif;
    }

    inline void RestoreInterrupts(IrqFlags flags) {
        asm volatile("msr daif, %0" :: "r"(flags) : "memory");
    }

    inline void Halt() {
        asm volatile("wfi" ::: "memory");
    }

    // WFI wakes on a pending interrupt even while PSTATE.I masks it, so
    // checking for work and then waiting with IRQs masked cannot lose a
    // wakeup. Unmask briefly afterwards so the pending IRQ is taken.
    inline void HaltWithInterruptsDisabled() {
        asm volatile("wfi\n\tmsr daifclr, #2\n\tisb\n\tmsr daifset, #2" ::: "memory");
    }

    [[noreturn]] inline void HaltForever() {
        for (;;) {
            asm volatile("msr daifset, #0xf\n\twfi" ::: "memory");
        }
    }

    inline uint64_t ReadTimestamp() {
        uint64_t value;
        asm volatile("isb\n\tmrs %0, cntvct_el0" : "=r"(value) :: "memory");
        return value;
    }

    // The generic timer runs from boot at CNTFRQ_EL0 Hz, so it times short
    // delays without calibration.
    inline void IoDelay() {
        uint64_t frequency;
        asm volatile("mrs %0, cntfrq_el0" : "=r"(frequency));
        uint64_t ticks = frequency / 1000000 + 1;
        uint64_t start = ReadTimestamp();
        while (ReadTimestamp() - start < ticks) {
            Pause();
        }
    }

    // ARMv8 has no "clean the whole cache" instruction: set/way maintenance
    // is only safe from firmware, which PSCI performs on power transitions.
    // Order outstanding accesses so the firmware call sees them.
    inline void WritebackInvalidateCaches() {
        asm volatile("dsb sy\n\tisb" ::: "memory");
    }

    // The TLBI operand is VA[55:12] in bits [43:0]; the bits above are RES0.
    inline uint64_t DataCacheLineSize() {
        uint64_t ctr;
        asm volatile("mrs %0, ctr_el0" : "=r"(ctr));
        return 4ULL << ((ctr >> 16) & 0xF);  // DminLine
    }

    // To the point of coherency: where non-snooping device DMA looks.
    inline void CleanDataCache(const void* address, uint64_t size) {
        uint64_t line = DataCacheLineSize();
        uint64_t end = (uint64_t)address + size;
        for (uint64_t p = (uint64_t)address & ~(line - 1); p < end; p += line) {
            asm volatile("dc cvac, %0" :: "r"(p) : "memory");
        }
        asm volatile("dsb sy" ::: "memory");
    }

    // Clean too: a partial line at either end may hold CPU data next to the
    // buffer, which a plain invalidate would throw away.
    inline void InvalidateDataCache(const void* address, uint64_t size) {
        uint64_t line = DataCacheLineSize();
        uint64_t end = (uint64_t)address + size;
        for (uint64_t p = (uint64_t)address & ~(line - 1); p < end; p += line) {
            asm volatile("dc civac, %0" :: "r"(p) : "memory");
        }
        asm volatile("dsb sy" ::: "memory");
    }

    // Clean the data cache to the point of unification, then invalidate
    // the whole instruction cache: invalidating by address would only
    // reach the alias it names on an aliasing (VIPT) instruction cache.
    inline void SyncInstructionCache(const void* address, uint64_t size) {
        uint64_t line = DataCacheLineSize();
        uint64_t start = (uint64_t)address & ~(line - 1);
        uint64_t end = (uint64_t)address + size;
        for (uint64_t p = start; p < end; p += line) {
            asm volatile("dc cvau, %0" :: "r"(p) : "memory");
        }
        asm volatile("dsb ish\n\tic ialluis\n\tdsb ish\n\tisb" ::: "memory");
    }

    inline void FlushTlbPage(uint64_t virtualAddress) {
        asm volatile("dsb ishst\n\ttlbi vaae1is, %0\n\tdsb ish\n\tisb"
                     :: "r"((virtualAddress >> 12) & ((1ULL << 44) - 1)) : "memory");
    }

    inline void FlushTlbAll() {
        asm volatile("dsb ishst\n\ttlbi vmalle1is\n\tdsb ish\n\tisb" ::: "memory");
    }

    // The process (lower) half of the address space is translated through
    // TTBR0_EL1; the kernel half through TTBR1_EL1, which never changes.
    static constexpr uint64_t TtbrAddressMask = 0x0000FFFFFFFFFFFEULL;

    inline uint64_t ReadPageTableRoot() {
        uint64_t ttbr;
        asm volatile("mrs %0, ttbr0_el1" : "=r"(ttbr));
        return ttbr & TtbrAddressMask;
    }

    // No ASIDs yet: every switch invalidates the old process's entries.
    inline void WritePageTableRoot(uint64_t rootPhys) {
        asm volatile("dsb ishst\n\tmsr ttbr0_el1, %0\n\tisb\n\ttlbi vmalle1\n\tdsb nsh\n\tisb"
                     :: "r"(rootPhys) : "memory");
    }

}