//! echoserver -- an event-loop echo server over loopback (127.0.0.1).
//!
//! One process serves several clients at once without threads: it never
//! blocks on any single socket. Every socket is nonblocking, and the loop
//! asks the kernel "which of my sockets are ready?" (epoll_wait, or plain
//! poll), then only accepts/reads/writes those:
//!
//!     loop {
//!         ready = epoll_wait(...)        // sleep until something happens
//!         for each ready fd:
//!             listener -> accept new clients, watch them too
//!             client   -> read; echo it back; 0 bytes = client left
//!     }
//!
//! The self-test forks 3 client processes that talk to it over loopback,
//! once with epoll and once with poll, then checks UDP over loopback.
//! With an argument (`echoserver PORT`) it just serves forever.
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

extern crate alloc;

use alloc::format;
use alloc::vec::Vec;
use ulib::*;

const LOOPBACK: Ip = Ip([127, 0, 0, 1]);
const CLIENTS: usize = 3;
const ROUNDS: usize = 3;

/// A client process: connect, send ROUNDS messages, check each echo.
fn client(port: u16, id: usize) -> i64 {
    let fd = socket(SOCK_STREAM, 0);
    if connect(fd, LOOPBACK, port) < 0 {
        return 10;
    }
    for round in 0..ROUNDS {
        let msg = format!("hello {} from client {}", round, id);
        write_all(fd, msg.as_bytes());
        let mut buf = [0u8; 64];
        let mut got = 0;
        while got < msg.len() {
            let n = read(fd, &mut buf[got..]);
            if n <= 0 {
                return 11;
            }
            got += n as usize;
        }
        if &buf[..got] != msg.as_bytes() {
            return 12;
        }
    }
    close(fd); // the server sees end-of-file
    0
}

fn listener(port: u16) -> i64 {
    let fd = socket(SOCK_STREAM | O_NONBLOCK, 0); // SOCK_NONBLOCK
    if bind(fd, port) < 0 || listen(fd, 4) < 0 {
        return -1;
    }
    fd
}

/// Accept everyone waiting (nonblocking: stop at EAGAIN).
fn accept_all(lfd: i64, mut add: impl FnMut(i64)) {
    loop {
        let (fd, _) = accept(lfd);
        if fd < 0 {
            break; // -EAGAIN: nobody else waiting
        }
        set_nonblocking(fd, true);
        add(fd);
    }
}

/// Read what's there and echo it back. Returns false when the client left.
fn echo(fd: i64, served: &mut usize) -> bool {
    let mut buf = [0u8; 512];
    loop {
        let n = read(fd, &mut buf);
        if n == -EAGAIN {
            return true; // drained for now
        }
        if n <= 0 {
            return false; // 0 = end of file (or an error)
        }
        write_all(fd, &buf[..n as usize]);
        *served += 1;
    }
}

/// The epoll flavour: register fds once, then wait for just the ready ones.
/// Returns the number of clients served to the end (or -1 on a stall).
fn serve_epoll(lfd: i64, want: usize) -> i64 {
    let ep = epoll_create1(0);
    epoll_ctl(ep, EPOLL_CTL_ADD, lfd, EPOLLIN, lfd as u64);
    let (mut finished, mut served) = (0, 0);
    let mut events = [EpollEvent::default(); 8];
    while finished < want {
        let n = epoll_wait(ep, &mut events, 5000);
        if n <= 0 {
            close(ep);
            return -1; // 5 s with nothing happening: something's wrong
        }
        for ev in &events[..n as usize] {
            let fd = ev.data as i64;
            if fd == lfd {
                accept_all(lfd, |c| {
                    epoll_ctl(ep, EPOLL_CTL_ADD, c, EPOLLIN, c as u64);
                });
            } else if !echo(fd, &mut served) {
                epoll_ctl(ep, EPOLL_CTL_DEL, fd, 0, 0);
                close(fd);
                finished += 1;
            }
        }
    }
    close(ep);
    println!("echoserver: epoll: {} clients, {} echoes", finished, served);
    finished as i64
}

/// The poll flavour: pass the whole fd list on every call.
fn serve_poll(lfd: i64, want: usize) -> i64 {
    let mut fds = alloc::vec![PollFd { fd: lfd as i32, events: POLLIN, revents: 0 }];
    let (mut finished, mut served) = (0, 0);
    while finished < want {
        if poll(&mut fds, 5000) <= 0 {
            return -1;
        }
        let mut new = Vec::new();
        let mut gone = Vec::new();
        for p in fds.iter() {
            if p.revents == 0 {
                continue;
            }
            if p.fd as i64 == lfd {
                accept_all(lfd, |c| new.push(c));
            } else if !echo(p.fd as i64, &mut served) {
                close(p.fd as i64);
                gone.push(p.fd);
                finished += 1;
            }
        }
        fds.retain(|p| !gone.contains(&p.fd));
        fds.extend(new.into_iter().map(|c| PollFd { fd: c as i32, events: POLLIN, revents: 0 }));
    }
    println!("echoserver: poll: {} clients, {} echoes", finished, served);
    finished as i64
}

/// Start the clients, run the server, collect the clients' results.
fn round(c: &mut Checks, port: u16, use_epoll: bool) {
    let lfd = listener(port);
    c.check(lfd >= 0, "listen on a loopback port");
    let mut pids = Vec::new();
    for id in 0..CLIENTS {
        let pid = fork();
        if pid == 0 {
            close(lfd);
            exit(client(port, id));
        }
        pids.push(pid);
    }
    let finished = if use_epoll { serve_epoll(lfd, CLIENTS) } else { serve_poll(lfd, CLIENTS) };
    close(lfd);
    let ok = pids.iter().all(|&pid| wait(pid) == 0);
    let what = if use_epoll { "epoll" } else { "poll" };
    c.check(finished == CLIENTS as i64, &format!("{}: served {} clients to the end", what, CLIENTS));
    c.check(ok, &format!("{}: every client got its echoes back", what));
}

fn main() -> i64 {
    if let Some(port) = arg(1).and_then(|a| a.parse().ok()) {
        let lfd = listener(port);
        if lfd < 0 {
            eprintln!("echoserver: can't listen on port {}", port);
            return 1;
        }
        println!("echoserver: listening on port {} (epoll)", port);
        loop {
            serve_epoll(lfd, usize::MAX);
        }
    }
    let mut c = Checks::new("echoserver");
    round(&mut c, 7070, true);
    round(&mut c, 7071, false);

    // UDP over loopback: a datagram to ourselves.
    let a = socket(SOCK_DGRAM, 0);
    let b = socket(SOCK_DGRAM, 0);
    bind(b, 7072);
    c.check(sendto(a, b"ping", LOOPBACK, 7072) == 4, "UDP sendto 127.0.0.1");
    set_recv_timeout(b, 2000);
    let mut buf = [0u8; 16];
    let (n, from) = recvfrom(b, &mut buf);
    c.check(n == 4 && &buf[..4] == b"ping" && from.ip() == LOOPBACK, "UDP datagram arrives, from 127.0.0.1");
    close(a);
    close(b);
    c.done()
}