mod memory;

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

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

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

pub const PAGE_SIZE: usize = 4096;

extern {
    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;
}

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

#[no_mangle]
unsafe extern 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", None);
}

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

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