//! Reading program files from storage and loading them into new processes.

use libkos::fs::Storage;
use libkos::syscall;
use types::address::{ByteOffset, EntryPoint, VirtualAddress};
use types::bounded::{CapSlot, PageCount};
use types::budget::Limits;

const PAGE_SIZE: usize = 4096;
/// A program's stack (see `types::layout`), unless init's table says more.
pub(crate) const STACK_PAGES: usize = 8;
/// An app's: drawing (tiny-skia, fontdue) takes more than 16 KiB.
pub(crate) const APP_STACK_PAGES: usize = 16;
/// The most the stack's band holds, leaving a guard page at its bottom.
const MAX_STACK_PAGES: usize = types::layout::STACK_BAND / PAGE_SIZE - 1;
const CREATE_CAP: CapSlot = CapSlot::const_new(0);
/// Where a program file is read through on its way into a process, a
/// chunk at a time, so a program may be any size.
const SCRATCH_VA: VirtualAddress = VirtualAddress::from_raw(0x0120_0000);
const SCRATCH_PAGES: u16 = 32;

pub(crate) struct LoadedProcess {
    pub(crate) process_cap: CapSlot,
    pub(crate) entry_point: EntryPoint,
}

/// A program file read forwards through a buffer (storage has no seek): a
/// segment's bytes are asked for in increasing file order.
struct Stream<'a> {
    storage: &'a mut Storage,
    handle: libkos::fs::FileHandle,
    buf: &'a mut [u8],
    /// The file offset of `buf[0]`, and how much of `buf` holds file.
    at: usize,
    len: usize,
    size: usize,
}

impl Stream<'_> {
    /// Reads the next chunk after what `buf` holds; false at the end or
    /// on an error.
    fn advance(&mut self) -> bool {
        let next = self.at + self.len;
        let want = self.buf.len().min(self.size - next.min(self.size));
        match self.storage.read(self.handle, &mut self.buf[..want]) {
            Ok(n) if n > 0 => {
                (self.at, self.len) = (next, n);
                true
            }
            _ => false,
        }
    }

    /// Hands the file's bytes `from..from + len` to `sink`, piece by piece,
    /// with each piece's place in that range.
    fn copy(&mut self, from: usize, len: usize, mut sink: impl FnMut(usize, &[u8])) -> bool {
        if from < self.at || from.checked_add(len).is_none_or(|end| end > self.size) {
            return false;
        }
        let mut done = 0;
        while done < len {
            let want = from + done;
            while want >= self.at + self.len {
                if !self.advance() {
                    return false;
                }
            }
            let start = want - self.at;
            let n = (self.len - start).min(len - done);
            sink(done, &self.buf[start..start + n]);
            done += n;
        }
        true
    }
}

pub(crate) fn load_binary(
    storage: &mut Storage,
    scratch: &mut [u8],
    filename: &[u8],
    limits: Limits,
    stack_pages: usize,
) -> Option<LoadedProcess> {
    let handle = storage.open(filename).ok()?;
    let loaded = storage
        .file_size(handle)
        .ok()
        .and_then(|size| load_from(&mut Stream { storage, handle, buf: scratch, at: 0, len: 0, size }, limits, stack_pages));
    let _ = storage.close(handle);
    loaded
}

fn load_from(file: &mut Stream, limits: Limits, stack_pages: usize) -> Option<LoadedProcess> {
    let stack_pages = stack_pages.min(MAX_STACK_PAGES);
    // The headers are in the first chunk.
    if !file.advance() {
        return None;
    }
    let elf_info = elf::parse(&file.buf[..file.len]).ok()?;
    let process_cap = syscall::create_process(CREATE_CAP, limits)?;
    let filled = elf_info.segments().all(|segment| load_segment(file, process_cap, segment).is_some())
        && (0..stack_pages).all(|i| {
            let dst = VirtualAddress::from_raw(types::layout::STACK_TOP - (stack_pages - i) * PAGE_SIZE);
            PageCount::new(1).and_then(syscall::request_pages).is_some_and(|va| {
                syscall::map_page(process_cap, va, dst, false);
                true
            })
        });
    if !filled {
        let _ = syscall::kill(process_cap);
        syscall::drop_cap(process_cap);
        return None;
    }
    Some(LoadedProcess { process_cap, entry_point: EntryPoint::from_raw(elf_info.entry_point) })
}

/// Fresh pages with the segment's bytes, zero past them, moved into the
/// process at its address.
fn load_segment(file: &mut Stream, process_cap: CapSlot, segment: &elf::Segment) -> Option<()> {
    let page_start = segment.vaddr & !0xFFF;
    let offset_in_page = segment.vaddr & 0xFFF;
    let seg_pages = (offset_in_page + segment.mem_size).div_ceil(PAGE_SIZE);
    let local = libkos::window::request(PageCount::new(u16::try_from(seg_pages).ok()?)?)?;
    local.fill_u32(ByteOffset::ZERO, seg_pages * PAGE_SIZE / 4, 0);
    let copy_len = segment.file_size.min(segment.mem_size);
    if copy_len > 0 && !file.copy(segment.file_offset, copy_len, |at, bytes| local.copy_from(ByteOffset::new(offset_in_page + at), bytes)) {
        return None;
    }
    for p in 0..seg_pages {
        let src_page = VirtualAddress::from_raw(local.address().raw() + p * PAGE_SIZE);
        let dst_page = VirtualAddress::from_raw(page_start + p * PAGE_SIZE);
        syscall::map_page(process_cap, src_page, dst_page, segment.flags.execute);
    }
    Some(())
}

/// A region of init's own to read program files into.
pub(crate) fn scratch() -> Option<&'static mut [u8]> {
    let region = syscall::create_region(PageCount::new(SCRATCH_PAGES)?).ok()?;
    libkos::window::map_buffer(region, SCRATCH_VA).ok()
}