use crate::manifest::{BinarySource, BootManifest};
use types::id::ProcessId;

const MAX_SERVERS: usize = 16;
const MAX_RESTARTS: u8 = 5;
const BACKOFF_TICKS: [usize; 5] = [0, 10_000_000, 50_000_000, 300_000_000, 300_000_000];

struct ServerState {
    name: &'static str,
    pid: Option<ProcessId>,
    restartable: bool,
    restart_count: u8,
    binary: BinarySource,
}

pub struct Supervisor {
    servers: [Option<ServerState>; MAX_SERVERS],
    count: usize,
    storage_endpoint: usize,
    create_cap: usize,
}

impl Supervisor {
    pub fn new(manifest: &BootManifest, storage_endpoint: usize, create_cap: usize) -> Self {
        let mut servers = [const { None }; MAX_SERVERS];
        for (i, entry) in manifest.servers.iter().enumerate() {
            if i < MAX_SERVERS {
                servers[i] = Some(ServerState {
                    name: entry.name,
                    pid: None,
                    restartable: entry.restartable,
                    restart_count: 0,
                    binary: entry.binary,
                });
            }
        }
        Self {
            servers,
            count: manifest.servers.len().min(MAX_SERVERS),
            storage_endpoint,
            create_cap,
        }
    }

    pub fn register_pid(&mut self, index: usize, pid: ProcessId) {
        if let Some(state) = &mut self.servers[index] {
            state.pid = Some(pid);
        }
    }

    pub fn should_restart(&self, index: usize) -> bool {
        self.servers[index]
            .as_ref()
            .map(|s| s.restartable && s.restart_count < MAX_RESTARTS)
            .unwrap_or(false)
    }

    fn find_by_pid(&self, pid: ProcessId) -> Option<usize> {
        for i in 0..self.count {
            if let Some(state) = &self.servers[i] {
                if state.pid == Some(pid) {
                    return Some(i);
                }
            }
        }
        None
    }

    fn backoff_delay(&self, restart_count: u8) -> usize {
        let idx = (restart_count as usize).min(BACKOFF_TICKS.len() - 1);
        BACKOFF_TICKS[idx]
    }

    fn wait_ticks(&self, ticks: usize) {
        if ticks == 0 { return; }
        let start = userspace::syscall::get_time();
        let deadline = start + ticks;
        while userspace::syscall::get_time() < deadline {
            userspace::syscall::yield_now();
        }
    }

    fn restart_server(&mut self, index: usize) {
        let (binary, restart_count) = match &self.servers[index] {
            Some(state) => (state.binary, state.restart_count),
            None => return,
        };

        let delay = self.backoff_delay(restart_count);
        self.wait_ticks(delay);

        match binary {
            BinarySource::Ram { addr: _, len: _ } => {
                let new_pid = userspace::syscall::create_process(self.create_cap);
                if let Some(state) = &mut self.servers[index] {
                    state.pid = Some(ProcessId::from_raw(new_pid as u16));
                    state.restart_count += 1;
                }
            }
            BinarySource::Storage { path } => {
                self.load_from_storage_and_restart(index, path);
            }
        }
    }

    fn load_from_storage_and_restart(&mut self, index: usize, path: &str) {
        use crate::protocol_helpers;

        let handle = protocol_helpers::storage_open(self.storage_endpoint, path.as_bytes());
        if handle == usize::MAX { return; }

        let new_pid = userspace::syscall::create_process(self.create_cap);

        protocol_helpers::storage_close(self.storage_endpoint, handle);

        if let Some(state) = &mut self.servers[index] {
            state.pid = Some(ProcessId::from_raw(new_pid as u16));
            state.restart_count += 1;
        }
    }

    pub fn run(&mut self, fault_endpoint: usize) -> ! {
        loop {
            let (crashed_pid_raw, _fault_code, _fault_addr, _fault_pc) =
                userspace::syscall::receive(fault_endpoint);

            let crashed_pid = ProcessId::from_raw(crashed_pid_raw as u16);

            if let Some(index) = self.find_by_pid(crashed_pid) {
                if self.should_restart(index) {
                    self.restart_server(index);
                }
            }
        }
    }
}