#![no_std]
#![no_main]

mod rng;
mod ssh;

use lancer_core::ipc::{MAX_IPC_BYTES, unpack_bytes};
use lancer_core::packet_ring::{PacketRingReader, PacketRingWriter};
use lancer_user::path;
use lancer_user::syscall;
use ssh::{SshProgress, SshSession};

const SHELL_RING_BASE_SLOT: u64 = 64;
const SHELL_RING_FRAME_COUNT: u64 = 4;
const OWN_NOTIF_SLOT: u64 = 3;
const NETSTACK_NOTIF_SLOT: u64 = 4;

const MAX_SESSIONS: usize = 4;

const TAG_ENDPOINT: u64 = 1;
const TAG_NOTIFICATION: u64 = 2;
const TAG_SCHED_CONTEXT: u64 = 5;
const TAG_FRAME: u64 = 11;
const RIGHTS_ALL: u64 = 0b1111;

const SLOT_UNTYPED: u64 = 28;

const VFS_RING_BASE: u64 = 128;
const VFS_RING_FRAMES: u64 = 16;
const VFS_CLIENT_NOTIF_SLOT: u64 = 12;
const VFS_NOTIF_SLOT: u64 = 13;
const VFS_PROC_SLOT: u64 = 15;

const BOOT_MODULE_COUNT: u64 = 7;

const RING_BASE: u64 = 0x4000_0000;
const RING_HALF: usize = 8192;

const MSG_DATA: u8 = 0;
const MSG_CONNECT: u8 = 1;
const MSG_DISCONNECT: u8 = 2;
const MSG_DNS_RESULT: u8 = 4;
const MSG_PING_RESULT: u8 = 6;
const MSG_PING_DONE: u8 = 7;
const MSG_UDP_RECV: u8 = 10;
const MSG_UDP_CLOSE: u8 = 11;
const MSG_UDP_BOUND: u8 = 12;
const MSG_DNS_RESULT6: u8 = 14;
const MSG_IFCONFIG_RESULT: u8 = 17;
const MSG_ARP_ENTRY: u8 = 19;
const MSG_ARP_DONE: u8 = 20;
const MSG_NETSTAT_ENTRY: u8 = 22;
const MSG_NETSTAT_ENTRY_UDP: u8 = 23;
const MSG_NETSTAT_DONE: u8 = 24;

const IDLE_TIMEOUT_MS: i64 = 300_000;
const NOTIFY_BIT_SHELL_DEATH_BASE: u64 = 0x10;
const NOTIFY_BIT_SHELL_NETSOCK: u64 = 0x04;

const INPUT_RING_PAGES: u64 = 1;
const INPUT_RING_SIZE: usize = 4096;
const INPUT_RING_SLOT_SIZE: u32 = 64;

const NETSOCK_RING_PAGES: u64 = 2;
const NETSOCK_RING_HALF: usize = 4096;
const NETSOCK_RING_SLOT_SIZE: u32 = 128;

static mut SSH_INBUFS: [[u8; 35000]; MAX_SESSIONS] = [[0u8; 35000]; MAX_SESSIONS];
static mut SSH_OUTBUFS: [[u8; 35000]; MAX_SESSIONS] = [[0u8; 35000]; MAX_SESSIONS];

fn slot_proc(i: usize) -> u64 {
    30 + i as u64 * 7
}
fn slot_sched(i: usize) -> u64 {
    31 + i as u64 * 7
}
fn slot_output_ep(i: usize) -> u64 {
    32 + i as u64 * 7
}
fn slot_input_ring_mem(i: usize) -> u64 {
    33 + i as u64 * 7
}
fn slot_netsock_mem(i: usize) -> u64 {
    34 + i as u64 * 7
}
fn slot_shell_notif(i: usize) -> u64 {
    36 + i as u64 * 7
}

fn input_ring_vaddr(i: usize) -> u64 {
    0x7000_0000 + i as u64 * 0x1000
}
fn netsock_ring_vaddr(i: usize) -> u64 {
    0x8000_0000 + i as u64 * 0x2000
}

struct Session<'a> {
    ssh: Option<SshSession<'a>>,
    shell_running: bool,
    input_tx: Option<PacketRingWriter>,
    netsock_tx: Option<PacketRingWriter>,
    netsock_rx: Option<PacketRingReader>,
    last_activity_ms: i64,
}

impl<'a> Session<'a> {
    const fn new() -> Self {
        Self {
            ssh: None,
            shell_running: false,
            input_tx: None,
            netsock_tx: None,
            netsock_rx: None,
            last_activity_ms: 0,
        }
    }
}

fn find_module_by_name(name: &[u8]) -> Option<u64> {
    let mut dummy = [0u8; 0];
    let total = syscall::module_info(u64::MAX, &mut dummy);
    if total < 0 {
        return None;
    }

    let mut path_buf = [0u8; 128];
    (BOOT_MODULE_COUNT..total as u64).find(|&i| {
        path_buf = [0u8; 128];
        let path_len = syscall::module_info(i, &mut path_buf);
        match path_len < 0 {
            true => false,
            false => {
                let used = (path_len as usize).min(path_buf.len());
                path::bytes_eq(path::extract_filename(&path_buf[..used]), name)
            }
        }
    })
}

fn send_ssh_bytes(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8, data: &[u8]) {
    let mut offset = 0usize;
    core::iter::from_fn(|| match offset < data.len() {
        true => {
            let written = ssh.write_plaintext(&data[offset..]);
            offset += written;
            flush_ssh_output(ssh, tx, conn_id);
            Some(())
        }
        false => None,
    })
    .take(64)
    .count();
}

fn flush_ssh_output(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8) -> bool {
    let mut flushed = false;
    let mut tmp = [0u8; 126];
    core::iter::from_fn(|| {
        let n = ssh.peek_output(&mut tmp[2..]);
        match n > 0 {
            true => {
                tmp[0] = conn_id;
                tmp[1] = MSG_DATA;
                match tx.try_push(&tmp[..2 + n]) {
                    true => {
                        ssh.consume_output(n);
                        flushed = true;
                        Some(())
                    }
                    false => None,
                }
            }
            false => None,
        }
    })
    .take(32)
    .count();
    if flushed {
        syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
    }
    flushed
}

fn drive_ssh(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8) -> bool {
    let mut any_flushed = false;
    let mut keep_going = true;
    core::iter::from_fn(|| match keep_going {
        true => {
            any_flushed |= flush_ssh_output(ssh, tx, conn_id);
            let r = ssh.progress();
            any_flushed |= flush_ssh_output(ssh, tx, conn_id);
            match r {
                SshProgress::Continue => Some(()),
                SshProgress::ShellReady => {
                    any_flushed = true;
                    None
                }
                SshProgress::Idle => {
                    keep_going = false;
                    None
                }
                SshProgress::Disconnected => {
                    keep_going = false;
                    None
                }
            }
        }
        false => None,
    })
    .take(128)
    .count();
    any_flushed
}

fn spawn_shell(conn_id: usize, module_idx: u64, session: &mut Session) -> bool {
    let proc_slot = slot_proc(conn_id);
    let sched_slot = slot_sched(conn_id);
    let output_ep = slot_output_ep(conn_id);
    let input_mem = slot_input_ring_mem(conn_id);
    let netsock_mem = slot_netsock_mem(conn_id);
    let shell_notif = slot_shell_notif(conn_id);

    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_ENDPOINT, 0, output_ep, 1);
    if r < 0 {
        return false;
    }

    let child_pid = syscall::proc_create(proc_slot);
    if child_pid < 0 {
        syscall::cap_revoke(output_ep);
        return false;
    }

    let r = syscall::proc_load_module(proc_slot, module_idx);
    if r < 0 {
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }

    let r = syscall::cap_grant(output_ep, proc_slot, 2, RIGHTS_ALL);
    if r < 0 {
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }

    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_FRAME, 0, input_mem, INPUT_RING_PAGES);
    if r < 0 {
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }
    let ir_vaddr = input_ring_vaddr(conn_id);
    let map_ok = (0..INPUT_RING_PAGES).all(|i| {
        syscall::frame_map(input_mem + i, ir_vaddr + i * 4096, 1) >= 0
    });
    if !map_ok {
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }
    let input_tx = unsafe {
        PacketRingWriter::init(ir_vaddr as *mut u8, INPUT_RING_SIZE, INPUT_RING_SLOT_SIZE)
    };
    let r = syscall::cap_grant(input_mem, proc_slot, 1, RIGHTS_ALL);
    if r < 0 {
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }

    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_FRAME, 0, netsock_mem, NETSOCK_RING_PAGES);
    if r < 0 {
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }
    let ns_vaddr = netsock_ring_vaddr(conn_id);
    let ns_map_ok = (0..NETSOCK_RING_PAGES).all(|i| {
        syscall::frame_map(netsock_mem + i, ns_vaddr + i * 4096, 1) >= 0
    });
    if !ns_map_ok {
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
        });
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(netsock_mem + i);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }
    let netsock_tx = unsafe {
        PacketRingWriter::init(
            ns_vaddr as *mut u8,
            NETSOCK_RING_HALF,
            NETSOCK_RING_SLOT_SIZE,
        )
    };
    let _ = unsafe {
        PacketRingWriter::init(
            (ns_vaddr + NETSOCK_RING_HALF as u64) as *mut u8,
            NETSOCK_RING_HALF,
            NETSOCK_RING_SLOT_SIZE,
        )
    };
    let netsock_rx = unsafe {
        PacketRingReader::attach(
            (ns_vaddr + NETSOCK_RING_HALF as u64) as *mut u8,
            NETSOCK_RING_HALF,
        )
    };
    let netsock_grant_ok = (0..NETSOCK_RING_PAGES).all(|i| {
        syscall::cap_grant(netsock_mem + i, proc_slot, 3 + i, RIGHTS_ALL) >= 0
    });
    if !netsock_grant_ok {
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
        });
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(netsock_mem + i);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }

    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_NOTIFICATION, 0, shell_notif, 1);
    if r < 0 {
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
        });
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(netsock_mem + i);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }
    let r = syscall::cap_grant(shell_notif, proc_slot, 9, RIGHTS_ALL);
    if r < 0 {
        syscall::cap_revoke(shell_notif);
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
        });
        (0..NETSOCK_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(netsock_mem + i);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
        });
        (0..INPUT_RING_PAGES).for_each(|i| {
            syscall::cap_revoke(input_mem + i);
        });
        syscall::proc_destroy(proc_slot);
        syscall::cap_revoke(output_ep);
        return false;
    }

    let _ = syscall::cap_grant(OWN_NOTIF_SLOT, proc_slot, 10, 0b0010);

    (0..VFS_RING_FRAMES).for_each(|i| {
        let _ = syscall::cap_grant(VFS_RING_BASE + i, proc_slot, 11 + i, RIGHTS_ALL);
    });
    [
        (VFS_CLIENT_NOTIF_SLOT, 28u64),
        (VFS_NOTIF_SLOT, 29),
        (VFS_PROC_SLOT, 15),
    ]
    .iter()
    .for_each(|&(src, dst)| {
        let _ = syscall::cap_grant(src, proc_slot, dst, RIGHTS_ALL);
    });

    let death_badge = NOTIFY_BIT_SHELL_DEATH_BASE << conn_id as u64;
    syscall::proc_bind_death_notif(proc_slot, OWN_NOTIF_SLOT, death_badge);

    let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_SCHED_CONTEXT, 0, sched_slot, 1);
    if r < 0 {
        cleanup_shell_caps(conn_id);
        return false;
    }

    let r = syscall::sched_configure(sched_slot, 10_000_000, 10_000_000, 100);
    if r < 0 {
        cleanup_shell_caps(conn_id);
        return false;
    }

    let r = syscall::sched_attach(sched_slot, child_pid as u64);
    if r < 0 {
        cleanup_shell_caps(conn_id);
        return false;
    }

    let r = syscall::proc_start(proc_slot, 2);
    if r < 0 {
        cleanup_shell_caps(conn_id);
        return false;
    }

    syscall::notify_signal(VFS_CLIENT_NOTIF_SLOT, 1);

    session.shell_running = true;
    session.input_tx = Some(input_tx);
    session.netsock_tx = Some(netsock_tx);
    session.netsock_rx = Some(netsock_rx);

    true
}

fn cleanup_shell_caps(conn_id: usize) {
    let ir_vaddr = input_ring_vaddr(conn_id);
    let ns_vaddr = netsock_ring_vaddr(conn_id);
    let input_mem = slot_input_ring_mem(conn_id);
    let netsock_mem = slot_netsock_mem(conn_id);

    (0..INPUT_RING_PAGES).for_each(|i| {
        syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
    });
    (0..INPUT_RING_PAGES).for_each(|i| {
        syscall::cap_revoke(input_mem + i);
    });

    (0..NETSOCK_RING_PAGES).for_each(|i| {
        syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
    });
    (0..NETSOCK_RING_PAGES).for_each(|i| {
        syscall::cap_revoke(netsock_mem + i);
    });

    syscall::cap_revoke(slot_shell_notif(conn_id));
    syscall::cap_revoke(slot_sched(conn_id));
    syscall::cap_revoke(slot_output_ep(conn_id));
}

fn drain_shell_output(conn_id: usize, ssh: &mut SshSession, tx: &PacketRingWriter) -> bool {
    let ep_slot = slot_output_ep(conn_id);
    let cid = conn_id as u8;
    let mut ipc_raw = [0u8; MAX_IPC_BYTES];
    let mut had_data = false;

    core::iter::from_fn(|| {
        let (status, msg) = syscall::nb_recv(ep_slot);
        match status >= 0 {
            true => {
                let count = unpack_bytes(&msg, &mut ipc_raw);
                const _: () = assert!(MAX_IPC_BYTES * 2 <= 80);
                if count > 0 {
                    had_data = true;
                    let mut translated = [0u8; MAX_IPC_BYTES * 2];
                    let tlen = (0..count).fold(0usize, |pos, i| match ipc_raw[i] {
                        b'\n' => {
                            translated[pos] = b'\r';
                            translated[pos + 1] = b'\n';
                            pos + 2
                        }
                        b => {
                            translated[pos] = b;
                            pos + 1
                        }
                    });
                    send_ssh_bytes(ssh, tx, cid, &translated[..tlen]);
                }
                Some(())
            }
            false => None,
        }
    })
    .take(64)
    .count();
    had_data
}

fn drain_shell_netsock(
    conn_id: usize,
    netsock_rx: &PacketRingReader,
    netstack_tx: &PacketRingWriter,
) -> bool {
    let cid = conn_id as u8;
    let mut buf = [0u8; 128];
    let mut relayed = false;

    core::iter::from_fn(|| match netsock_rx.try_pop(&mut buf) {
        Some(n) if n > 0 => {
            buf[0] = cid;
            match netstack_tx.try_push(&buf[..n]) {
                true => {
                    relayed = true;
                    Some(())
                }
                false => None,
            }
        }
        _ => None,
    })
    .take(16)
    .count();

    relayed
}

fn relay_netstack_to_shell(netsock_tx: &PacketRingWriter, data: &[u8], shell_notif: u64) {
    let mut buf = [0u8; 128];
    let copy_len = data.len().min(buf.len() - 1);
    buf[0] = 0;
    buf[1..1 + copy_len].copy_from_slice(&data[..copy_len]);
    if netsock_tx.try_push(&buf[..1 + copy_len]) {
        syscall::notify_signal(shell_notif, 4);
    }
}

fn teardown_session(conn_id: usize, session: &mut Session, tx: &PacketRingWriter) {
    if session.shell_running {
        syscall::proc_destroy(slot_proc(conn_id));
        cleanup_shell_caps(conn_id);
        session.shell_running = false;
        session.input_tx = None;
        session.netsock_tx = None;
        session.netsock_rx = None;
    }
    if let Some(ssh) = session.ssh.as_mut() {
        ssh.close();
    }
    let ctrl = [conn_id as u8, MSG_DISCONNECT];
    let _ = tx.try_push(&ctrl);
    drop(session.ssh.take());
}

#[unsafe(no_mangle)]
pub extern "C" fn lancer_main() -> ! {
    lancer_user::show!(rsh, "starting");

    if !(0..SHELL_RING_FRAME_COUNT).all(|i| syscall::frame_map(SHELL_RING_BASE_SLOT + i, RING_BASE + i * 4096, 1) >= 0) {
        lancer_user::show!(rsh, error, "ring frame_map failed");
        syscall::exit();
    }

    let rx_base = RING_BASE as *mut u8;
    let tx_base = (RING_BASE + RING_HALF as u64) as *mut u8;

    let rx = unsafe { PacketRingReader::attach(rx_base, RING_HALF) };
    let tx = unsafe { PacketRingWriter::init(tx_base, RING_HALF, 128) };

    let host_key = match sunset::SignKey::generate(sunset::KeyType::Ed25519, None) {
        Ok(k) => k,
        Err(_) => {
            lancer_user::show!(rsh, error, "keygen failed");
            syscall::exit();
        }
    };

    let shell_module_idx = find_module_by_name(b"shell");
    match shell_module_idx {
        Some(_) => lancer_user::show!(rsh, "shell module found"),
        None => lancer_user::show!(rsh, warn, "shell module not found"),
    }

    lancer_user::show!(rsh, "ready");

    let mut sessions: [Session; MAX_SESSIONS] = [
        Session::new(),
        Session::new(),
        Session::new(),
        Session::new(),
    ];
    let mut pop_buf = [0u8; 128];
    let mut plaintext_buf = [0u8; 256];

    syscall::ntfn_bind(OWN_NOTIF_SLOT);

    loop {
        let active_shell = (0..MAX_SESSIONS).find(|&i| sessions[i].shell_running);
        let mut any_output = false;

        let bits = match active_shell {
            Some(cid) => {
                let (status, msg) = syscall::recv(slot_output_ep(cid));
                match status == syscall::BOUND_NOTIFICATION_BADGE {
                    true => msg[1],
                    false => {
                        if status >= 0 {
                            let Session { ref mut ssh, .. } = sessions[cid];
                            if let Some(ssh) = ssh.as_mut() {
                                let mut ipc_raw = [0u8; MAX_IPC_BYTES];
                                let count = unpack_bytes(&msg, &mut ipc_raw);
                                if count > 0 {
                                    let mut translated = [0u8; MAX_IPC_BYTES * 2];
                                    let tlen = (0..count).fold(0usize, |pos, i| match ipc_raw[i] {
                                        b'\n' => {
                                            translated[pos] = b'\r';
                                            translated[pos + 1] = b'\n';
                                            pos + 2
                                        }
                                        b => {
                                            translated[pos] = b;
                                            pos + 1
                                        }
                                    });
                                    send_ssh_bytes(ssh, &tx, cid as u8, &translated[..tlen]);
                                    any_output = true;
                                }
                            }
                        }
                        let (_, poll_bits) = syscall::notify_poll(OWN_NOTIF_SLOT);
                        poll_bits
                    }
                }
            }
            None => {
                let (_, wait_bits) = syscall::notify_wait(OWN_NOTIF_SLOT);
                wait_bits
            }
        };

        let now_ms = syscall::clock_monotonic_ms().max(0);

        (0..MAX_SESSIONS).for_each(|i| {
            let death_bit = NOTIFY_BIT_SHELL_DEATH_BASE << i as u64;
            if bits & death_bit == 0 || !sessions[i].shell_running {
                return;
            }

            if let Some(ssh) = sessions[i].ssh.as_mut() {
                drain_shell_output(i, ssh, &tx);
            }
            syscall::proc_destroy(slot_proc(i));
            cleanup_shell_caps(i);
            sessions[i].shell_running = false;
            sessions[i].input_tx = None;
            sessions[i].netsock_tx = None;
            sessions[i].netsock_rx = None;
            any_output = true;

            if let Some(ssh) = sessions[i].ssh.as_mut() {
                ssh.close();
            }
            let ctrl = [i as u8, MSG_DISCONNECT];
            let _ = tx.try_push(&ctrl);
            drop(sessions[i].ssh.take());
        });

        if bits & NOTIFY_BIT_SHELL_NETSOCK != 0 {
            (0..MAX_SESSIONS).for_each(|i| {
                if let Some(ref netsock_rx) = sessions[i].netsock_rx {
                    let relayed = drain_shell_netsock(i, netsock_rx, &tx);
                    if relayed {
                        syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
                        any_output = true;
                    }
                }
            });
        }

        (0..MAX_SESSIONS).for_each(|i| {
            if let Some(ssh) = sessions[i].ssh.as_mut() {
                any_output |= drive_ssh(ssh, &tx, i as u8);
            }
        });

        let mut sessions_to_drop: u8 = 0;
        let mut need_spawn: u8 = 0;

        core::iter::from_fn(|| match rx.try_pop(&mut pop_buf) {
            Some(n) if n >= 2 => {
                let conn_id = pop_buf[0] as usize;
                let msg_type = pop_buf[1];
                if conn_id >= MAX_SESSIONS {
                    return Some(());
                }

                match msg_type {
                    MSG_CONNECT => {
                        if sessions[conn_id].shell_running {
                            syscall::proc_destroy(slot_proc(conn_id));
                            cleanup_shell_caps(conn_id);
                            sessions[conn_id].shell_running = false;
                            sessions[conn_id].input_tx = None;
                            sessions[conn_id].netsock_tx = None;
                            sessions[conn_id].netsock_rx = None;
                        }
                        drop(sessions[conn_id].ssh.take());
                        let inbuf: &'static mut [u8; 35000] = unsafe {
                            &mut *((&raw mut SSH_INBUFS) as *mut [u8; 35000]).add(conn_id)
                        };
                        let outbuf: &'static mut [u8; 35000] = unsafe {
                            &mut *((&raw mut SSH_OUTBUFS) as *mut [u8; 35000]).add(conn_id)
                        };
                        inbuf.fill(0);
                        outbuf.fill(0);
                        sessions[conn_id].ssh =
                            Some(SshSession::new(inbuf, outbuf, host_key.clone()));
                        sessions[conn_id].last_activity_ms = now_ms;
                        if let Some(ssh) = sessions[conn_id].ssh.as_mut() {
                            any_output |= drive_ssh(ssh, &tx, conn_id as u8);
                            syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
                        }
                    }
                    MSG_DISCONNECT => {
                        teardown_session(conn_id, &mut sessions[conn_id], &tx);
                    }
                    MSG_DATA if n > 2 => {
                        let Session {
                            ref mut ssh,
                            ref input_tx,
                            ref shell_running,
                            ref mut last_activity_ms,
                            ..
                        } = sessions[conn_id];
                        *last_activity_ms = now_ms;
                        if let Some(ssh) = ssh.as_mut() {
                            let data = &pop_buf[2..n];
                            let mut off = 0usize;
                            core::iter::from_fn(|| match off < data.len() {
                                true => {
                                    let accepted = ssh.feed_input(&data[off..]);
                                    match accepted {
                                        0 => {
                                            any_output |= drive_ssh(ssh, &tx, conn_id as u8);
                                            if ssh.shell_ready && *shell_running {
                                                let pt_n = ssh.read_plaintext(&mut plaintext_buf);
                                                if pt_n > 0
                                                    && let Some(it) = input_tx
                                                    && it.try_push(&plaintext_buf[..pt_n])
                                                {
                                                    syscall::notify_signal(
                                                        slot_shell_notif(conn_id),
                                                        2,
                                                    );
                                                }
                                            }
                                            any_output |= flush_ssh_output(ssh, &tx, conn_id as u8);
                                            let retry = ssh.feed_input(&data[off..]);
                                            match retry {
                                                0 => None,
                                                n2 => {
                                                    off += n2;
                                                    Some(())
                                                }
                                            }
                                        }
                                        n2 => {
                                            off += n2;
                                            Some(())
                                        }
                                    }
                                }
                                false => None,
                            })
                            .take(32)
                            .count();

                            any_output |= drive_ssh(ssh, &tx, conn_id as u8);
                            if ssh.shell_ready && *shell_running {
                                let pt_n = ssh.read_plaintext(&mut plaintext_buf);
                                if pt_n > 0
                                    && let Some(it) = input_tx
                                    && it.try_push(&plaintext_buf[..pt_n])
                                {
                                    syscall::notify_signal(slot_shell_notif(conn_id), 2);
                                }
                            }
                        }
                    }
                    MSG_DNS_RESULT
                    | MSG_DNS_RESULT6
                    | MSG_PING_RESULT
                    | MSG_PING_DONE
                    | MSG_UDP_RECV
                    | MSG_UDP_CLOSE
                    | MSG_UDP_BOUND
                    | MSG_IFCONFIG_RESULT
                    | MSG_ARP_ENTRY
                    | MSG_ARP_DONE
                    | MSG_NETSTAT_ENTRY
                    | MSG_NETSTAT_ENTRY_UDP
                    | MSG_NETSTAT_DONE
                        if sessions[conn_id].shell_running =>
                    {
                        if let Some(ref netsock_tx) = sessions[conn_id].netsock_tx {
                            relay_netstack_to_shell(
                                netsock_tx,
                                &pop_buf[1..n],
                                slot_shell_notif(conn_id),
                            );
                        }
                        any_output = true;
                    }
                    _ => {}
                }
                Some(())
            }
            _ => None,
        })
        .take(64)
        .count();

        (0..MAX_SESSIONS).for_each(|i| {
            let Session {
                ref mut ssh,
                ref shell_running,
                ref input_tx,
                ..
            } = sessions[i];
            if let Some(ssh) = ssh.as_mut() {
                any_output |= drive_ssh(ssh, &tx, i as u8);
                if ssh.shell_ready && !*shell_running {
                    need_spawn |= 1 << i;
                }
                if ssh.shell_ready && *shell_running {
                    let pt_n = ssh.read_plaintext(&mut plaintext_buf);
                    if pt_n > 0
                        && let Some(it) = input_tx
                        && it.try_push(&plaintext_buf[..pt_n])
                    {
                        syscall::notify_signal(slot_shell_notif(i), 2);
                    }
                }
                any_output |= flush_ssh_output(ssh, &tx, i as u8);
            }
        });

        (0..MAX_SESSIONS).for_each(|i| {
            if need_spawn & (1 << i) == 0 {
                return;
            }
            match shell_module_idx {
                Some(idx) => {
                    let ok = spawn_shell(i, idx, &mut sessions[i]);
                    if !ok {
                        if let Some(ssh) = sessions[i].ssh.as_mut() {
                            send_ssh_bytes(ssh, &tx, i as u8, b"Failed to start shell\r\n");
                            flush_ssh_output(ssh, &tx, i as u8);
                            ssh.close();
                        }
                        sessions_to_drop |= 1 << i;
                    }
                }
                None => {
                    if let Some(ssh) = sessions[i].ssh.as_mut() {
                        send_ssh_bytes(ssh, &tx, i as u8, b"Shell module not found\r\n");
                        flush_ssh_output(ssh, &tx, i as u8);
                        ssh.close();
                    }
                    sessions_to_drop |= 1 << i;
                }
            }
        });

        (0..MAX_SESSIONS).for_each(|i| {
            if sessions[i].shell_running
                && let Some(ssh) = sessions[i].ssh.as_mut()
            {
                any_output |= drain_shell_output(i, ssh, &tx);
            }
        });

        (0..MAX_SESSIONS).for_each(|i| {
            if sessions_to_drop & (1 << i) != 0 {
                teardown_session(i, &mut sessions[i], &tx);
            }
        });

        (0..MAX_SESSIONS).for_each(|i| {
            if sessions[i].ssh.is_none() {
                return;
            }
            if now_ms - sessions[i].last_activity_ms > IDLE_TIMEOUT_MS {
                teardown_session(i, &mut sessions[i], &tx);
            }
        });

        if any_output {
            syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
        }
    }
}