#![allow(static_mut_refs)]

use alloc::vec::Vec;
use core::sync::atomic::{AtomicU64, Ordering};

use moa_bitops::{bitmap::BitMap, make_64bit_mask};
use moa_cpumask::{CpuId, Cpumask};

use crate::{
    println,
    processor::{nr_cpus, this_processor_id},
    space::mm::{MmStruct, entries::Pgd, tlb::local_flush_tlb_all},
    sync::spinlock::Spinlock,
};

pub(crate) struct MmContext(pub(crate) AtomicU64);

impl MmContext {
    pub(crate) const fn new() -> Self {
        MmContext(AtomicU64::new(0))
    }
}

const ID_AA64MMFR0_ASID_SHIFT: usize = 4;

static mut ASID_BITS: u64 = 0;
static CPU_ASID_LOCK: Spinlock<usize> = Spinlock::new(0);

static ASID_GENERATION: AtomicU64 = AtomicU64::new(0);
static mut ASID_MAP: BitMap = BitMap::new();

static mut ACTIVE_ASIDS: Vec<AtomicU64> = Vec::new();
static mut RESERVED_ASIDS: Vec<u64> = Vec::new();
static TLB_FLUSH_PENDING: Cpumask = Cpumask::new();

#[inline(always)]
fn asid_mask() -> u64 {
    unsafe { !make_64bit_mask(0, ASID_BITS as usize) }
}

#[inline(always)]
fn asid_first_version() -> usize {
    unsafe { 1 << ASID_BITS }
}

#[inline(always)]
fn num_user_asids() -> usize {
    asid_first_version()
}

#[inline(always)]
fn asid2idx(asid: u64) -> u64 {
    asid & !asid_mask()
}

#[inline(always)]
fn idx2asid(idx: u64) -> u64 {
    asid2idx(idx)
}

fn get_cpu_asid_bits() -> u32 {
    let value: u64;
    unsafe {
        core::arch::asm!(
            "mrs {0}, ID_AA64MMFR0_EL1",
            out(reg) value
        );
    }
    let fld = (value << (64 - 4 - ID_AA64MMFR0_ASID_SHIFT)) >> (64 - 4);
    match fld {
        0 => 8,
        2 => 16,
        _ => {
            println!(
                "CPU{}: Unknown ASID size ({}); assuming 8-bit",
                this_processor_id().raw(),
                fld
            );
            8
        },
    }
}

fn verify_cpu_asid_bits() {
    let asid = u64::from(get_cpu_asid_bits());
    let boot_cpu_asid = unsafe { ASID_BITS };

    assert!(
        asid >= boot_cpu_asid,
        "CPU{}: smaller ASID size({}) than boot CPU ({})",
        this_processor_id().raw(),
        asid,
        boot_cpu_asid
    );
}

pub(crate) fn context_init() {
    unsafe {
        for _ in 0..nr_cpus() {
            RESERVED_ASIDS.push(0);
            ACTIVE_ASIDS.push(AtomicU64::new(0));
        }

        ASID_BITS = u64::from(get_cpu_asid_bits());
        assert!(num_user_asids() - 1 > nr_cpus());
        ASID_MAP.init(num_user_asids());
    }
    ASID_GENERATION.store(asid_first_version() as u64, Ordering::Relaxed);

    println!("ASID allocator initialised with {} entries", num_user_asids());
}

pub(crate) fn context_init_cpu() {
    verify_cpu_asid_bits();
}

// function

fn flush_context() {
    unsafe {
        ASID_MAP.clear_range(0, num_user_asids());
        for i in 0..nr_cpus() {
            let active_asids = &ACTIVE_ASIDS[i];
            let mut asid = active_asids.swap(0, Ordering::Relaxed);
            let reserved_asids = &mut RESERVED_ASIDS[i];
            if asid == 0 {
                asid = *reserved_asids;
            }
            ASID_MAP.set_bit(asid2idx(asid) as usize);
            *reserved_asids = asid;

            TLB_FLUSH_PENDING.set(CpuId::from_raw(i));
        }
    }
}

fn check_update_reserved_asid(asid: u64, newasid: u64) -> bool {
    unsafe {
        let mut hit = false;
        for reserved_asids in RESERVED_ASIDS.iter_mut().take(nr_cpus()) {
            if *reserved_asids == asid {
                hit = true;
                *reserved_asids = newasid;
            }
        }
        hit
    }
}

fn set_aisd(asid: u64, generation: u64) -> u64 {
    unsafe {
        ASID_MAP.set_bit(asid as usize);
        CUR_IDX = asid as usize;
        idx2asid(asid) | generation
    }
}

static mut CUR_IDX: usize = 1;

impl MmStruct {
    fn new_context(&self) -> u64 {
        unsafe {
            let mut asid = self.context.0.load(Ordering::Relaxed);
            let mut generation = ASID_GENERATION.load(Ordering::Relaxed);

            if asid != 0 {
                let newasid = generation | (asid & !asid_mask());

                if check_update_reserved_asid(asid, newasid) {
                    return newasid;
                }

                if !ASID_MAP.test_and_set_bit(asid2idx(asid) as usize) {
                    return newasid;
                }
            }

            asid = ASID_MAP.find_next_zero_bit(CUR_IDX) as u64;
            if asid != num_user_asids() as u64 {
                return set_aisd(asid, generation);
            }

            generation = ASID_GENERATION.fetch_add(asid_first_version() as u64, Ordering::Relaxed)
                + asid_first_version() as u64;
            flush_context();

            asid = ASID_MAP.find_next_zero_bit(1) as u64;
            set_aisd(asid, generation)
        }
    }

    // 检查并且映射上下文
    pub(crate) fn check_and_switch_context(&self) {
        unsafe {
            let cpu = this_processor_id();
            let mut asid = self.context.0.load(Ordering::Relaxed);
            let active_asids = &ACTIVE_ASIDS[cpu.raw()];
            let old_active_asid = active_asids.load(Ordering::Relaxed);
            if old_active_asid != 0
                && (((asid ^ ASID_GENERATION.load(Ordering::Relaxed)) >> ASID_BITS) == 0)
                && active_asids
                    .compare_exchange(old_active_asid, asid, Ordering::Relaxed, Ordering::Relaxed)
                    .is_ok()
            {
                cpu_switch_mm(&self.pgd, self.context.0.load(Ordering::Relaxed));
                return;
            }
            {
                let mut lock = CPU_ASID_LOCK.lock_irq_save();
                asid = self.context.0.load(Ordering::Relaxed);
                if (asid ^ ASID_GENERATION.load(Ordering::Relaxed) >> ASID_BITS) != 0 {
                    asid = self.new_context();
                    self.context.0.store(asid, Ordering::Relaxed);
                }

                if TLB_FLUSH_PENDING.test_and_clear(cpu) {
                    local_flush_tlb_all();
                }
                active_asids.store(asid, Ordering::Relaxed);
                *lock = 1;
            }
            cpu_switch_mm(&self.pgd, self.context.0.load(Ordering::Relaxed));
        }
    }
}

#[doc(hidden)]
#[unsafe(no_mangle)]
#[allow(unreachable_pub)]
pub extern "C" fn post_ttbr_update_workaround() {
    unsafe {
        core::arch::asm!("nop; nop; nop");
    }
}

#[inline(always)]
fn cpu_switch_mm(pgd: &Pgd, asid: u64) {
    unsafe extern "C" {
        fn cpu_do_switch_mm(phys: usize, asid: u64);
    }
    unsafe {
        cpu_do_switch_mm(pgd.to_phys().to_value(), asid);
    }
}

core::arch::global_asm!(
    r"
TTBR_CNP_BIT = 1

.global cpu_do_switch_mm; .align 2; cpu_do_switch_mm:
    mrs	x2, ttbr1_el1
    mov x3, x0

    cbz     x1, 1f                          // skip CNP for reserved ASID
    orr     x3, x3, #TTBR_CNP_BIT
1:
    bfi	x2, x1, #48, #16		// set the ASID
    msr	ttbr1_el1, x2			// in TTBR1 (since TCR.A1 is set)
    isb
    msr	ttbr0_el1, x3			// now update TTBR0
    isb
    b	post_ttbr_update_workaround	// Back to C code...
.type cpu_do_switch_mm, @function; .size cpu_do_switch_mm, .-cpu_do_switch_mm
"
);
