//! Init's spawn service (`types::spawn`): loads a program for a client and
//! hands it the process, unstarted. Each client has an allowance from the
//! manifest; a child's limits come out of it and return when the child dies.

use libkos::syscall;
use server_framework::{Decode, Encode, MALFORMED, Server};
use types::badge::Badge;
use types::bounded::CapSlot;
use types::budget::Limits;
use types::permission::{Permission, PermissionSet};
use types::spawn::{MAX_NAME, NONE, reply};

use crate::launch::Launcher;
use crate::load;

pub(crate) const MAX_CLIENTS: usize = 4;
/// Death signals carry the child's index as a bit.
const MAX_CHILDREN: usize = 8;

struct Child {
    client: usize,
    limits: Limits,
    process: CapSlot,
}

pub(crate) struct Spawner {
    /// What each client (badge = index + 1) may still give its children.
    allowances: [Option<Limits>; MAX_CLIENTS],
    children: [Option<Child>; MAX_CHILDREN],
}

fn fits(want: Limits, left: Limits) -> bool {
    want.pages <= left.pages
        && want.endpoints <= left.endpoints
        && want.notifications <= left.notifications
        && want.regions <= left.regions
}

fn add(a: Limits, b: Limits) -> Limits {
    Limits {
        pages: a.pages + b.pages,
        endpoints: a.endpoints + b.endpoints,
        notifications: a.notifications + b.notifications,
        regions: a.regions + b.regions,
    }
}

fn sub(a: Limits, b: Limits) -> Limits {
    Limits {
        pages: a.pages - b.pages,
        endpoints: a.endpoints - b.endpoints,
        notifications: a.notifications - b.notifications,
        regions: a.regions - b.regions,
    }
}

impl Spawner {
    pub(crate) const fn new() -> Self {
        Self { allowances: [None; MAX_CLIENTS], children: [const { None }; MAX_CHILDREN] }
    }

    /// The badge for a new client with `allowance`.
    pub(crate) fn add_client(&mut self, allowance: Limits) -> Option<Badge> {
        let i = self.allowances.iter().position(Option::is_none)?;
        self.allowances[i] = Some(allowance);
        Some(Badge::new(i + 1))
    }
}

pub(crate) struct Request {
    name: [u8; MAX_NAME],
    len: usize,
    limits: Limits,
}

impl Decode for Request {
    fn decode(words: [usize; 4]) -> Option<Self> {
        types::spawn::parse(words).map(|(name, len, limits)| Self { name, len, limits })
    }
}

pub(crate) struct Reply([usize; 4]);

impl Encode for Reply {
    fn encode(&self) -> [usize; 4] {
        self.0
    }
}

fn refused(tag: usize) -> ([usize; 4], Option<CapSlot>) {
    ([tag, 0, 0, 0], None)
}

impl Launcher<'_> {
    fn spawn(&mut self, badge: usize, r: &Request) -> ([usize; 4], Option<CapSlot>) {
        let Some(client) = badge.checked_sub(1).filter(|c| *c < MAX_CLIENTS) else {
            return refused(reply::ERROR);
        };
        let Some(left) = self.spawner.allowances[client] else {
            return refused(reply::ERROR);
        };
        if !fits(r.limits, left) {
            return refused(reply::OVER_ALLOWANCE);
        }
        let Some(index) = self.spawner.children.iter().position(Option::is_none) else {
            return refused(reply::ERROR);
        };
        let name = &r.name[..r.len];
        let (storage, scratch) = self.storage();
        match storage.open(name) {
            Ok(h) => {
                let _ = storage.close(h);
            }
            Err(_) => return refused(reply::NOT_FOUND),
        }
        let Some(loaded) = load::load_binary(storage, scratch, name, r.limits) else {
            return refused(reply::ERROR);
        };
        if !self.watch(loaded.process_cap, index) {
            let _ = syscall::kill(loaded.process_cap);
            syscall::drop_cap(loaded.process_cap);
            return refused(reply::ERROR);
        }
        let call_only = PermissionSet::empty().with(Permission::Execute);
        let log = libkos::start::log().and_then(|l| syscall::delegate(l, loaded.process_cap, call_only));
        self.spawner.allowances[client] = Some(sub(left, r.limits));
        self.spawner.children[index] = Some(Child { client, limits: r.limits, process: loaded.process_cap });
        let log = log.map_or(NONE, |s| s.value() as usize);
        ([reply::OK, loaded.entry_point.raw(), log, 0], Some(loaded.process_cap))
    }

    /// Has the child's death signal bit `index` on init's death notification.
    fn watch(&self, process: CapSlot, index: usize) -> bool {
        let signal = PermissionSet::empty().with(Permission::Write);
        let Ok(copy) = syscall::mint(self.deaths, Badge::new(1 << index), signal) else {
            return false;
        };
        let watched = syscall::watch_death(process, copy).is_ok();
        syscall::drop_cap(copy);
        watched
    }

    /// Children behind `bits` died: their limits go back to their clients.
    fn reap(&mut self, bits: usize) {
        for (i, slot) in self.spawner.children.iter_mut().enumerate() {
            if bits & (1 << i) == 0 {
                continue;
            }
            if let Some(child) = slot.take() {
                if let Some(Some(left)) = self.spawner.allowances.get_mut(child.client) {
                    *left = add(*left, child.limits);
                }
                syscall::drop_cap(child.process);
            }
        }
    }
}

impl Server for Launcher<'_> {
    type Request = Request;
    type Response = Reply;

    /// Without a badge there is no client, so no allowance.
    fn handle(&mut self, _request: Request) -> Reply {
        Reply(refused(reply::ERROR).0)
    }

    fn handle_call(&mut self, badge: usize, words: [usize; 4], _cap: Option<CapSlot>) -> ([usize; 4], Option<libkos::syscall::Give>) {
        match Request::decode(words) {
            Some(r) => {
                let (reply, process) = self.spawn(badge, &r);
                (reply, process.map(Into::into))
            }
            None => ([MALFORMED, 0, 0, 0], None),
        }
    }

    fn signal(&mut self, bits: usize) {
        self.reap(bits);
    }
}