use crate::hw::uart::Uart;
use crate::mmu;
use types::address::VirtualAddress;
use types::budget::Limits;
use types::id::{ProcessId, ThreadId};
use types::mapping::PagePerms;

const PAGE_SIZE: usize = mmu::PAGE_SIZE;
const USER_VA_BASE: usize = 0x10000;
/// The stack sits just below the image and stops short of the null page.
const USER_STACK_PAGES: usize = USER_VA_BASE / PAGE_SIZE - 1;

pub struct LoadedProcess {
    pub pid: ProcessId,
}

fn kern() -> &'static mut kernel::Kernel<crate::platform::RiscVPlatform> {
    unsafe { &mut *crate::kernel_ptr() }
}

/// The two programs boot starts, packed into the kernel image.
pub static INIT_ELF: &[u8] = include_bytes!(env!("KOS_INIT_ELF"));
pub static STORAGE_ELF: &[u8] = include_bytes!(env!("KOS_STORAGE_ELF"));

pub fn load_init(uart: &Uart) -> LoadedProcess {
    let (pid, tid) = kern()
        .boot(
            types::bounded::Priority::DEFAULT,
            types::time::TimeSlice::DEFAULT,
        )
        .unwrap_or_else(|_| crate::halt(uart, "boot failed"));

    let entry = load_elf(uart, pid, INIT_ELF);
    setup_thread(tid, entry);
    uart.write_str("k-os: init loaded\n");

    LoadedProcess { pid }
}

/// Storage's budget, carved from init's.
const STORAGE: Limits = {
    use types::storage::{BUFFER_PAGES, MAX_CLIENTS};
    Limits {
        // A buffer per client, and its image, heap and cache.
        pages: MAX_CLIENTS * BUFFER_PAGES + 256,
        endpoints: 2,
        notifications: 2,
        // A buffer per client, and the disk's DMA.
        regions: MAX_CLIENTS + 4,
    }
};

/// Storage, which init needs to read anything else, its budget carved
/// from init's.
pub fn load_storage(uart: &Uart, init: ProcessId) -> LoadedProcess {
    let pid = kern()
        .create_process_internal(init, STORAGE)
        .unwrap_or_else(|| crate::halt(uart, "process create failed"));
    let entry = load_elf(uart, pid, STORAGE_ELF);
    let tid = kern()
        .spawn_thread_internal(pid)
        .unwrap_or_else(|| crate::halt(uart, "thread spawn failed"));
    setup_thread(tid, entry);
    uart.write_str("k-os: storage loaded\n");

    LoadedProcess { pid }
}

fn setup_thread(tid: ThreadId, entry: usize) {
    if let Some(ctx) = kern().thread_context_mut(tid) {
        ctx.set_pc(entry);
        ctx.registers_mut()[2] = USER_VA_BASE;
        ctx.set_status(1 << 5);
    }
}

/// Copies `data` into `pid`'s already mapped pages starting at `va`.
fn copy_to_user(pid: ProcessId, va: usize, data: &[u8]) -> bool {
    let mut done = 0;
    while done < data.len() {
        let addr = va + done;
        let page = addr & !(PAGE_SIZE - 1);
        let offset = addr - page;
        let Some(pa) = kern().user_page_frame(pid, VirtualAddress::from_raw(page)) else {
            return false;
        };
        let len = (PAGE_SIZE - offset).min(data.len() - done);
        // SAFETY: `pa` is a pool frame the kernel just mapped for `pid`;
        // pool frames are RAM, mapped at `phys_to_virt`.
        unsafe {
            core::ptr::copy_nonoverlapping(
                data.as_ptr().add(done),
                (mmu::phys_to_virt(pa.raw()) as *mut u8).add(offset),
                len,
            );
        }
        done += len;
    }
    true
}

fn segment_perms(flags: elf::SegmentFlags) -> Option<PagePerms> {
    match (flags.write, flags.execute) {
        (true, true) => None,
        (false, true) => Some(PagePerms::ReadExecute),
        (true, false) => Some(PagePerms::ReadWrite),
        (false, false) => Some(PagePerms::ReadOnly),
    }
}

/// Loads `filename` into `pid` through the kernel's own mapping API, so
/// every page is tracked, charged and checked like any other. Halts on any
/// malformed or overlapping image: boot servers are trusted to be sane.
fn load_elf(uart: &Uart, pid: ProcessId, file: &[u8]) -> usize {
    if file.is_empty() {
        crate::halt(uart, "program missing from the kernel image (make build)");
    }
    let elf_info = elf::parse(file).unwrap_or_else(|_| crate::halt(uart, "bad ELF"));

    for segment in elf_info.segments() {
        let perms =
            segment_perms(segment.flags).unwrap_or_else(|| crate::halt(uart, "W+X segment"));
        let page_start = segment.vaddr & !(PAGE_SIZE - 1);
        let offset_in_page = segment.vaddr - page_start;
        let pages = offset_in_page
            .checked_add(segment.mem_size)
            .map(|total| total.div_ceil(PAGE_SIZE))
            .unwrap_or_else(|| crate::halt(uart, "segment too large"));
        if pages == 0 {
            continue;
        }

        kern()
            .map_fresh_range(pid, VirtualAddress::from_raw(page_start), pages, perms)
            .unwrap_or_else(|_| crate::halt(uart, "cannot map segment"));

        let copy_len = segment.file_size.min(segment.mem_size);
        let data = segment
            .file_offset
            .checked_add(copy_len)
            .and_then(|end| file.get(segment.file_offset..end))
            .unwrap_or_else(|| crate::halt(uart, "segment outside file"));
        if !copy_to_user(pid, segment.vaddr, data) {
            crate::halt(uart, "segment copy failed");
        }
    }

    let stack_bottom = USER_VA_BASE - USER_STACK_PAGES * PAGE_SIZE;
    kern()
        .map_fresh_range(
            pid,
            VirtualAddress::from_raw(stack_bottom),
            USER_STACK_PAGES,
            PagePerms::ReadWrite,
        )
        .unwrap_or_else(|_| crate::halt(uart, "cannot map stack"));

    elf_info.entry_point
}