#![no_std]

extern crate alloc;

use types::address::{EntryPoint, StackPointer, VirtualAddress};
use types::bounded::{CapSlot, PageCount};
use types::permission::{Permission, PermissionSet};
use libkos::syscall;

const PAGE_SIZE: usize = 4096;
const USER_CODE_BASE: usize = 0x10000;
const USER_STACK_TOP: StackPointer = StackPointer::from_raw(USER_CODE_BASE);
const USER_STACK_PAGES: usize = 4;
const CMD_READ_FILE: usize = 1;

const CREATE_CAP: CapSlot = CapSlot::const_new(0);
const STORAGE_EP: CapSlot = CapSlot::const_new(2);
const DISPLAY_EP: CapSlot = CapSlot::const_new(3);
const INPUT_EP: CapSlot = CapSlot::const_new(5);

fn request_file(shared_buf: *mut u8, filename: &[u8]) -> usize {
    unsafe {
        for (i, &b) in filename.iter().enumerate() {
            core::ptr::write_volatile(shared_buf.add(i), b);
        }
    }
    syscall::send(STORAGE_EP, CMD_READ_FILE, filename.len(), 0, 0);
    let (file_size, _, _, _) = syscall::receive(STORAGE_EP);
    file_size
}

struct LoadedProcess {
    process_cap: CapSlot,
    entry_point: EntryPoint,
}

fn load_binary(shared_buf: *mut u8, filename: &[u8]) -> Option<LoadedProcess> {
    let file_size = request_file(shared_buf, filename);
    if file_size == 0 {
        return None;
    }

    let elf_data = unsafe { core::slice::from_raw_parts(shared_buf as *const u8, file_size) };
    let elf_info = match elf::parse(elf_data) {
        Ok(info) => info,
        Err(_) => return None,
    };

    let process_cap = syscall::create_process(CREATE_CAP)?;

    for segment in elf_info.segments() {
        let page_start = segment.vaddr & !0xFFF;
        let offset_in_page = segment.vaddr & 0xFFF;
        let total_size = offset_in_page + segment.mem_size;
        let seg_pages = (total_size + PAGE_SIZE - 1) / PAGE_SIZE;

        let page_count = PageCount::new(seg_pages as u16)?;
        let local_va = syscall::request_pages(page_count)?;

        let local_ptr = local_va.raw() as *mut u8;
        unsafe {
            core::ptr::write_bytes(local_ptr, 0, seg_pages * PAGE_SIZE);
        }

        let copy_len = segment.file_size.min(segment.mem_size);
        if copy_len > 0 {
            let src = segment.file_offset;
            if src + copy_len <= file_size {
                unsafe {
                    core::ptr::copy_nonoverlapping(
                        shared_buf.add(src),
                        local_ptr.add(offset_in_page),
                        copy_len,
                    );
                }
            }
        }

        for p in 0..seg_pages {
            let src_page = VirtualAddress::from_raw(local_va.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);
        }
    }

    for i in 0..USER_STACK_PAGES {
        let page_count = PageCount::new(1)?;
        let stack_va = syscall::request_pages(page_count)?;
        let dst = VirtualAddress::from_raw(USER_CODE_BASE - (USER_STACK_PAGES - i) * PAGE_SIZE);
        syscall::map_page(process_cap, stack_va, dst, false);
    }

    Some(LoadedProcess {
        process_cap,
        entry_point: EntryPoint::from_raw(elf_info.entry_point),
    })
}

struct ServerInfo {
    filename: &'static [u8],
    name: &'static str,
}

const SERVERS: &[ServerInfo] = &[
    ServerInfo { filename: b"log.elf", name: "log" },
    ServerInfo { filename: b"permis.elf", name: "permis" },
    ServerInfo { filename: b"session.elf", name: "session" },
];

pub fn run(shared_buf_va: usize) -> ! {
    let shared_buf = shared_buf_va as *mut u8;

    libkos::println!("init: started (pid={:#x})", syscall::get_process_id().raw());

    let ipc_perms = PermissionSet::empty()
        .with(Permission::Read)
        .with(Permission::Execute);

    for server in SERVERS {
        libkos::print!("init: loading {}... ", server.name);

        let loaded = match load_binary(shared_buf, server.filename) {
            Some(l) => l,
            None => {
                libkos::println!("SKIP");
                continue;
            }
        };

        if let Some(ep_cap) = syscall::create_endpoint() {
            syscall::delegate(ep_cap, loaded.process_cap, ipc_perms);
            syscall::spawn_thread_in(loaded.process_cap, loaded.entry_point, USER_STACK_TOP);
        }

        libkos::println!("ok (entry={:#x})", loaded.entry_point.raw());
    }

    libkos::print!("init: loading snake... ");
    if let Some(snake) = load_binary(shared_buf, b"snake.elf") {
        syscall::delegate(INPUT_EP, snake.process_cap, ipc_perms);
        syscall::delegate(DISPLAY_EP, snake.process_cap, ipc_perms);
        syscall::spawn_thread_in(snake.process_cap, snake.entry_point, USER_STACK_TOP);
        libkos::println!("ok");
    } else {
        libkos::println!("SKIP");
    }

    libkos::println!("init: all loaded, supervising");

    loop {
        syscall::yield_now();
    }
}