mod memory;

use {
    alloc::{boxed::Box, vec::Vec},
    core::{
        arch::{asm, global_asm},
        fmt::Write,
    },
    sbi::system_reset::{system_reset, ResetReason, ResetType},
    spin::{Mutex, Once},
    uart_16550::MmioSerialPort,
};

use crate::{
    allocator,
    arch::riscv64::memory::{PageEntryFlags, PageSize, PageTable, PAGE_TABLE},
    memory::PhysicalAddress,
};

global_asm!(include_str!("entry.s"));
global_asm!(include_str!("memory_regions.s"));

pub const PAGE_SIZE: usize = 4096;

extern "C" {
    static TEXT_START: PhysicalAddress;
    static TEXT_END: PhysicalAddress;

    static RODATA_START: PhysicalAddress;
    static RODATA_END: PhysicalAddress;

    static DATA_START: PhysicalAddress;
    static DATA_END: PhysicalAddress;

    static SDATA_START: PhysicalAddress;
    static SDATA_END: PhysicalAddress;

    static BSS_START: PhysicalAddress;
    static BSS_END: PhysicalAddress;

    static INITIAL_KERNEL_HEAP_START: PhysicalAddress;
    static INITIAL_KERNEL_HEAP_SIZE: usize;

    static USABLE_MEMORY_START: PhysicalAddress;
    static USABLE_MEMORY_SIZE: usize;
}

pub static SERIAL_CONSOLE: Once<Mutex<MmioSerialPort>> = Once::new();

#[no_mangle]
unsafe extern "C" fn _kernel_start() -> ! {
    SERIAL_CONSOLE.call_once(|| Mutex::new(unsafe { MmioSerialPort::new(0x1000_0000) }));
    crate::logger::init().expect("failed to set logger");
    log::info!("Initialising AKern {}", crate::VERSION);

    allocator::init(
        INITIAL_KERNEL_HEAP_START.as_mut_ptr::<u8>(),
        INITIAL_KERNEL_HEAP_SIZE,
    );
    memory::init(USABLE_MEMORY_START.into(), USABLE_MEMORY_SIZE / PAGE_SIZE);

    let mut page_table_addr = PAGE_TABLE.get().unwrap().lock();
    let mut page_table = page_table_addr.as_mut_ptr::<PageTable>().as_mut().unwrap();

    // Map text (executable) section
    page_table.identity_map_range(TEXT_START, TEXT_END, PageEntryFlags::ReadExecute);
    // Map rodata section
    page_table.identity_map_range(RODATA_START, RODATA_END, PageEntryFlags::Read);
    // Map data section
    page_table.identity_map_range(DATA_START, DATA_END, PageEntryFlags::ReadWrite);
    // Map sdata section
    page_table.identity_map_range(SDATA_START, SDATA_END, PageEntryFlags::ReadWrite);
    // Map bss section (includes stack and initial kernel heap)
    page_table.identity_map_range(BSS_START, BSS_END, PageEntryFlags::ReadWrite);
    // Map usable memory range (as rw so not executable)
    page_table.identity_map_range(
        USABLE_MEMORY_START,
        USABLE_MEMORY_START + USABLE_MEMORY_SIZE.into(),
        PageEntryFlags::ReadWrite,
    );
    // Map Uart so we can continue using serial
    page_table.identity_map(
        0x1000_0000_usize.into(),
        PageEntryFlags::ReadWrite,
        PageSize::Size4KiB,
    );

    let table_ppn = page_table_addr.as_addr() as usize >> 12;
    let satp_value = 8 << 60 | table_ppn;
    log::info!("Enabling MMU");

    asm!(
        "csrw satp, {}",
        "sfence.vma",
        in(reg) satp_value,
    );

    crate::kmain::kmain("baka=9", Vec::new());
}

/// Spin loop
pub fn spin_loop() -> ! {
    loop {
        unsafe { asm!("wfi") }
    }
}

pub fn hardware_random_u64() -> u64 {
    0
}

pub fn register_dump() {}

pub fn log(args: core::fmt::Arguments<'_>) -> core::fmt::Result {
    SERIAL_CONSOLE.get().unwrap().lock().write_fmt(args)
}