//! BSD-style sockets: what `socket()` returns, as a file descriptor.
//!
//! A socket is a `File` (read = recv, write = send, close = close), plus
//! the socket-only calls connect/bind/listen/accept/sendto/recvfrom. Each
//! wraps one smoltcp socket (kept in the shared `SocketSet` in net/mod.rs):
//!   * TCP  (SOCK_STREAM): a reliable byte stream to one peer.
//!   * UDP  (SOCK_DGRAM): single datagrams, to/from anyone.
//!   * ICMP (SOCK_DGRAM + IPPROTO_ICMP): Linux's "ping socket": you send an
//!     ICMP echo request (8-byte header + payload), you receive the replies.
//!     The kernel fills in the identifier, so programs needn't be root.
//!
//! Blocking: smoltcp never blocks, it just says "no data yet". So e.g.
//! `recv` sleeps on `net::CHANGED` (woken after every `net::poll`) until
//! its socket has data, the peer closed, or the timeout (SO_RCVTIMEO) ran
//! out. No busy-waiting: a blocked program costs no CPU. A *nonblocking*
//! socket (O_NONBLOCK) never sleeps: "would block" becomes EAGAIN, and a
//! TCP connect returns EINPROGRESS while the handshake goes on (poll for
//! POLLOUT to learn when it's done). `File::poll` below says what's ready.
//!
//! A listening TCP socket is a small *backlog* of smoltcp sockets all
//! listening on the port (a smoltcp socket handles exactly one
//! connection). `accept` takes one that got connected, makes it a new fd,
//! and puts a fresh listener in its place.

use super::{Net, poll, with_net};
use crate::errno::Errno;
use crate::fs::file::{File, KIND_SOCKET, Stat};
use crate::poll::{POLLERR, POLLHUP, POLLIN, POLLOUT};
use core::sync::atomic::{AtomicBool, Ordering};
use alloc::vec;
use alloc::vec::Vec;
use smoltcp::iface::SocketHandle;
use smoltcp::socket::{icmp, tcp, udp};
use smoltcp::wire::{IpAddress, IpEndpoint, Ipv4Address};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

const TCP_BUFFER: usize = 16 * 1024;
const DGRAM_BUFFER: usize = 8 * 1024;
const DGRAM_PACKETS: usize = 16;
const MAX_BACKLOG: usize = 4;
const CONNECT_TIMEOUT_US: u64 = 10_000_000;
/// A closed connection gives up on an unresponsive peer after this long.
const CLOSE_TIMEOUT_MS: u64 = 10_000;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Kind {
    Tcp,
    Udp,
    Icmp,
}

struct State {
    handle: Option<SocketHandle>,
    listeners: Vec<SocketHandle>,
    local_port: u16,
    /// Default destination (UDP/ICMP after `connect`).
    remote: Option<IpEndpoint>,
    /// SO_RCVTIMEO: how long recv/accept may block (None = forever).
    timeout_us: Option<u64>,
    /// Stage 17: shutdown(SHUT_RD) / shutdown(SHUT_WR) done.
    rd_shut: bool,
    wr_shut: bool,
    /// SO_ERROR already reported a refused connection (it's said once).
    error_taken: bool,
    /// Boolean options set with setsockopt (bit n = option n, e.g.
    /// SO_REUSEADDR = 2): remembered so getsockopt reads them back.
    flags: u32,
}

impl State {
    fn new(handle: Option<SocketHandle>, local_port: u16) -> State {
        State { handle, listeners: Vec::new(), local_port, remote: None, timeout_us: None, rd_shut: false, wr_shut: false, error_taken: false, flags: 0 }
    }
}

pub struct Socket {
    pub kind: Kind,
    state: Mutex<State>,
    /// O_NONBLOCK: never sleep, say EAGAIN instead.
    nonblock: AtomicBool,
}

/// An IPv4 address + port, as plain numbers (the syscall layer converts
/// from/to `struct sockaddr_in`).
pub type Addr = ([u8; 4], u16);

fn endpoint((ip, port): Addr) -> IpEndpoint {
    IpEndpoint::new(IpAddress::Ipv4(Ipv4Address::from(ip)), port)
}

fn addr(ep: IpEndpoint) -> Addr {
    #[allow(irrefutable_let_patterns)]
    let IpAddress::Ipv4(ip) = ep.addr else { return ([0; 4], ep.port) };
    (ip.octets(), ep.port)
}

/// Wait for a usable route to `to`: loopback always has one, anything
/// else needs DHCP to have configured us.
fn route_ready(to: Addr) -> Result<(), Errno> {
    if super::is_loopback(to.0) { Ok(()) } else { super::wait_configured(5_000_000).map(|_| ()) }
}

fn new_tcp() -> tcp::Socket<'static> {
    tcp::Socket::new(tcp::SocketBuffer::new(vec![0; TCP_BUFFER]), tcp::SocketBuffer::new(vec![0; TCP_BUFFER]))
}

impl Socket {
    /// socket(): a new, unconnected socket.
    pub fn new(kind: Kind) -> Result<Socket, Errno> {
        if with_net(|_| ()).is_none() {
            return Err(Errno::ENETDOWN); // net::init not run
        }
        let mut state = State::new(None, 0);
        match kind {
            Kind::Tcp => {} // created by connect/listen
            Kind::Udp => {
                let buf = || udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY; DGRAM_PACKETS], vec![0; DGRAM_BUFFER]);
                state.handle = with_net(|n| n.sockets.add(udp::Socket::new(buf(), buf())));
            }
            Kind::Icmp => {
                let buf = || icmp::PacketBuffer::new(vec![icmp::PacketMetadata::EMPTY; DGRAM_PACKETS], vec![0; DGRAM_BUFFER]);
                state.handle = with_net(|n| {
                    let mut sock = icmp::Socket::new(buf(), buf());
                    let ident = n.ephemeral_port();
                    sock.bind(icmp::Endpoint::Ident(ident)).ok();
                    state.local_port = ident; // our echo identifier
                    n.sockets.add(sock)
                });
            }
        }
        Ok(Socket::with_state(kind, state))
    }

    fn with_state(kind: Kind, state: State) -> Socket {
        Socket { kind, state: Mutex::new(state), nonblock: AtomicBool::new(false) }
    }

    /// Lock our state and the network (always in that order, interrupts
    /// off) and run `f`.
    fn locked<R>(&self, f: impl FnOnce(&mut State, &mut Net) -> R) -> Result<R, Errno> {
        without_interrupts(|| {
            let mut st = self.state.lock();
            with_net(|n| f(&mut st, n)).ok_or(Errno::ENETDOWN)
        })
    }

    /// Block until `f` returns Some (re-checked after every network poll),
    /// or the deadline passes (-> EAGAIN, or `on_timeout`).
    fn wait<R>(
        &self,
        timeout_us: Option<u64>,
        mut f: impl FnMut(&mut State, &mut Net) -> Option<Result<R, Errno>>,
    ) -> Result<R, Errno> {
        if self.nonblock.load(Ordering::Relaxed) {
            // One look, no sleeping.
            return self.locked(&mut f)?.unwrap_or(Err(Errno::EAGAIN));
        }
        let deadline = timeout_us.map_or(u64::MAX, crate::time::deadline_after_us);
        // Interruptible: Ctrl+C must be able to stop a program waiting
        // for the network.
        super::CHANGED
            .wait_until_deadline_intr(deadline, || match self.locked(&mut f) {
                Ok(r) => r,
                Err(e) => Some(Err(e)),
            })?
            .unwrap_or(Err(Errno::EAGAIN))
    }

    fn recv_timeout(&self) -> Option<u64> {
        without_interrupts(|| self.state.lock().timeout_us)
    }

    pub fn set_recv_timeout(&self, us: Option<u64>) {
        without_interrupts(|| self.state.lock().timeout_us = us);
    }

    pub fn bind(&self, port: u16) -> Result<(), Errno> {
        self.locked(|st, n| {
            if st.local_port != 0 && self.kind != Kind::Icmp {
                return Err(Errno::EINVAL); // already bound
            }
            match self.kind {
                Kind::Tcp => st.local_port = port,
                Kind::Udp => {
                    let port = if port == 0 { n.ephemeral_port() } else { port };
                    n.sockets.get_mut::<udp::Socket>(st.handle.unwrap()).bind(port).map_err(|_| Errno::EADDRINUSE)?;
                    st.local_port = port;
                }
                Kind::Icmp => {} // ICMP has no ports; we keep our identifier
            }
            Ok(())
        })?
    }

    /// connect(): TCP does the three-way handshake (SYN, SYN-ACK, ACK) and
    /// waits for it; UDP/ICMP just remember the default destination.
    pub fn connect(&self, to: Addr) -> Result<(), Errno> {
        route_ready(to)?;
        if self.kind != Kind::Tcp {
            if self.kind == Kind::Udp && without_interrupts(|| self.state.lock().local_port) == 0 {
                self.bind(0)?;
            }
            return self.locked(|st, _| st.remote = Some(endpoint(to)));
        }
        self.locked(|st, n| {
            if st.handle.is_some() || !st.listeners.is_empty() {
                return Err(Errno::EISCONN);
            }
            let local = if st.local_port != 0 { st.local_port } else { n.ephemeral_port() };
            let mut sock = new_tcp();
            sock.connect(n.iface.context(), endpoint(to), local).map_err(|_| Errno::EINVAL)?;
            st.handle = Some(n.sockets.add(sock));
            st.local_port = local;
            Ok(())
        })??;
        poll(); // send the SYN now
        if self.nonblock.load(Ordering::Relaxed) {
            return Err(Errno::EINPROGRESS); // poll for POLLOUT to see it finish
        }
        let result = self.wait(Some(CONNECT_TIMEOUT_US), |st, n| {
            match n.sockets.get::<tcp::Socket>(st.handle?).state() {
                tcp::State::SynSent | tcp::State::SynReceived => None,
                tcp::State::Closed => Some(Err(Errno::ECONNREFUSED)), // got a RST
                _ => Some(Ok(())),
            }
        });
        if let Err(e) = result {
            // Failed: forget the half-open socket so connect can be retried.
            self.locked(|st, n| {
                if let Some(h) = st.handle.take() {
                    n.sockets.remove(h);
                }
            })?;
            return Err(if e == Errno::EAGAIN { Errno::ETIMEDOUT } else { e });
        }
        Ok(())
    }

    /// listen(): start `backlog` (1..=4) smoltcp listeners on our port.
    pub fn listen(&self, backlog: usize) -> Result<(), Errno> {
        if self.kind != Kind::Tcp {
            return Err(Errno::EOPNOTSUPP);
        }
        self.locked(|st, n| {
            if st.local_port == 0 || st.handle.is_some() {
                return Err(Errno::EINVAL);
            }
            while st.listeners.len() < backlog.clamp(1, MAX_BACKLOG) {
                let mut sock = new_tcp();
                sock.listen(st.local_port).map_err(|_| Errno::EINVAL)?;
                st.listeners.push(n.sockets.add(sock));
            }
            Ok(())
        })?
    }

    /// accept(): wait for one of our listeners to get a connection; it
    /// becomes the returned socket, a new listener takes its place.
    pub fn accept(&self) -> Result<(Socket, Addr), Errno> {
        if self.kind != Kind::Tcp {
            return Err(Errno::EOPNOTSUPP);
        }
        self.wait(self.recv_timeout(), |st, n| {
            if st.listeners.is_empty() {
                return Some(Err(Errno::EINVAL)); // not listening
            }
            let i = st.listeners.iter().position(|&h| {
                !matches!(n.sockets.get::<tcp::Socket>(h).state(), tcp::State::Listen | tcp::State::SynReceived)
            })?;
            let handle = st.listeners[i];
            let mut fresh = new_tcp();
            if fresh.listen(st.local_port).is_ok() {
                st.listeners[i] = n.sockets.add(fresh);
            } else {
                st.listeners.remove(i);
            }
            let peer = n.sockets.get::<tcp::Socket>(handle).remote_endpoint().map(addr).unwrap_or(([0; 4], 0));
            let state = State::new(Some(handle), st.local_port);
            Some(Ok((Socket::with_state(Kind::Tcp, state), peer)))
        })
    }

    /// send()/sendto()/write(). Blocks while the send buffer is full.
    pub fn send(&self, data: &[u8], to: Option<Addr>) -> Result<usize, Errno> {
        if self.kind == Kind::Udp && without_interrupts(|| self.state.lock().local_port) == 0 {
            if let Some(to) = to {
                route_ready(to)?;
            }
            self.bind(0)?; // first send picks a local port, like Linux
        }
        let sent = self.wait(None, |st, n| {
            if st.wr_shut {
                return Some(Err(Errno::EPIPE)); // shutdown(SHUT_WR)
            }
            let Some(h) = st.handle else { return Some(Err(Errno::ENOTCONN)) };
            match self.kind {
                Kind::Tcp => {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    if !sock.may_send() {
                        return Some(Err(Errno::EPIPE)); // closed
                    }
                    sock.can_send().then(|| sock.send_slice(data).map_err(|_| Errno::EPIPE))
                }
                Kind::Udp => {
                    let Some(dest) = to.map(endpoint).or(st.remote) else { return Some(Err(Errno::EDESTADDRREQ)) };
                    match n.sockets.get_mut::<udp::Socket>(h).send_slice(data, dest) {
                        Ok(()) => Some(Ok(data.len())),
                        Err(udp::SendError::BufferFull) => None,
                        Err(_) => Some(Err(Errno::EMSGSIZE)),
                    }
                }
                Kind::Icmp => {
                    let Some(dest) = to.map(endpoint).or(st.remote) else { return Some(Err(Errno::EDESTADDRREQ)) };
                    if data.len() < 8 {
                        return Some(Err(Errno::EINVAL)); // not even an ICMP header
                    }
                    // Our identifier goes in bytes 4..6 (smoltcp computes
                    // the checksum when it sends).
                    let mut packet = data.to_vec();
                    packet[4..6].copy_from_slice(&st.local_port.to_be_bytes());
                    match n.sockets.get_mut::<icmp::Socket>(h).send_slice(&packet, dest.addr) {
                        Ok(()) => Some(Ok(data.len())),
                        Err(icmp::SendError::BufferFull) => None,
                        Err(_) => Some(Err(Errno::EMSGSIZE)),
                    }
                }
            }
        })?;
        poll(); // put it on the wire now, not at netd's next wake-up
        Ok(sent)
    }

    /// recv()/recvfrom()/read(). Returns 0 at end of stream (TCP peer
    /// closed). A datagram bigger than `buf` is cut short (like Linux).
    pub fn recv(&self, buf: &mut [u8]) -> Result<(usize, Option<Addr>), Errno> {
        let result = self.wait(self.recv_timeout(), |st, n| {
            if st.rd_shut {
                return Some(Ok((0, None))); // shutdown(SHUT_RD): end of stream
            }
            let Some(h) = st.handle else {
                return Some(Err(if self.kind == Kind::Tcp { Errno::ENOTCONN } else { Errno::EINVAL }));
            };
            match self.kind {
                Kind::Tcp => {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    if sock.can_recv() {
                        Some(sock.recv_slice(buf).map(|len| (len, None)).map_err(|_| Errno::ECONNRESET))
                    } else if !sock.may_recv() {
                        Some(Ok((0, None))) // FIN received (or reset): end of stream
                    } else {
                        None
                    }
                }
                Kind::Udp => {
                    let sock = n.sockets.get_mut::<udp::Socket>(h);
                    if !sock.is_open() {
                        return Some(Err(Errno::EINVAL)); // never bound: nothing can arrive
                    }
                    let (data, meta) = sock.recv().ok()?;
                    let len = data.len().min(buf.len());
                    buf[..len].copy_from_slice(&data[..len]);
                    Some(Ok((len, Some(addr(meta.endpoint)))))
                }
                Kind::Icmp => {
                    let (data, from) = n.sockets.get_mut::<icmp::Socket>(h).recv().ok()?;
                    let len = data.len().min(buf.len());
                    buf[..len].copy_from_slice(&data[..len]);
                    Some(Ok((len, Some(addr(IpEndpoint::new(from, 0))))))
                }
            }
        })?;
        if self.kind == Kind::Tcp {
            poll(); // our receive window opened up: tell the peer
        }
        Ok(result)
    }

    /// getsockname (stage 17): our address and port. Connected TCP knows
    /// its local IP; otherwise it's 0.0.0.0 ("any of ours"), like an
    /// unconnected socket on Linux.
    pub fn local(&self) -> Addr {
        self.locked(|st, n| match (self.kind, st.handle) {
            (Kind::Tcp, Some(h)) => n.sockets.get::<tcp::Socket>(h).local_endpoint().map(addr).unwrap_or(([0; 4], st.local_port)),
            _ => ([0; 4], st.local_port),
        })
        .unwrap_or(([0; 4], 0))
    }

    /// shutdown (stage 17): 0 = SHUT_RD (reads say end-of-file), 1 =
    /// SHUT_WR (TCP sends its FIN; writes get EPIPE), 2 = both.
    pub fn shutdown(&self, how: u64) -> Result<(), Errno> {
        if how > 2 {
            return Err(Errno::EINVAL);
        }
        let kind = self.kind;
        self.locked(|st, n| {
            let connected = match kind {
                Kind::Tcp => st.handle.is_some(),
                _ => st.remote.is_some(),
            };
            if !connected {
                return Err(Errno::ENOTCONN);
            }
            if how != 1 {
                st.rd_shut = true;
            }
            if how != 0 {
                st.wr_shut = true;
                if let (Kind::Tcp, Some(h)) = (kind, st.handle) {
                    n.sockets.get_mut::<tcp::Socket>(h).close(); // FIN; we can still receive
                }
            }
            Ok(())
        })??;
        poll();
        Ok(())
    }

    /// SO_ERROR (stage 17): a pending error, reported once. That's how a
    /// program learns that a nonblocking connect failed.
    pub fn take_error(&self) -> i32 {
        let kind = self.kind;
        self.locked(|st, n| match (kind, st.handle) {
            (Kind::Tcp, Some(h)) if !st.error_taken && !st.wr_shut && n.sockets.get::<tcp::Socket>(h).state() == tcp::State::Closed => {
                st.error_taken = true;
                Errno::ECONNREFUSED as i32
            }
            _ => 0,
        })
        .unwrap_or(0)
    }

    /// Is it listening (SO_ACCEPTCONN)?
    pub fn listening(&self) -> bool {
        without_interrupts(|| !self.state.lock().listeners.is_empty())
    }

    /// The buffer sizes (SO_RCVBUF / SO_SNDBUF): fixed here.
    pub fn buffer_size(&self) -> usize {
        if self.kind == Kind::Tcp { TCP_BUFFER } else { DGRAM_BUFFER }
    }

    /// Remember / read back a boolean option (SO_REUSEADDR & co.). They
    /// change nothing: smoltcp lets a port be reused at once anyway, and
    /// there are no keepalives to send.
    pub fn set_flag(&self, opt: u32, on: bool) {
        without_interrupts(|| {
            let mut st = self.state.lock();
            if on { st.flags |= 1 << opt } else { st.flags &= !(1 << opt) }
        });
    }
    pub fn flag(&self, opt: u32) -> bool {
        without_interrupts(|| self.state.lock().flags & 1 << opt != 0)
    }

    /// The SO_RCVTIMEO set, if any (µs).
    pub fn recv_timeout_us(&self) -> Option<u64> {
        self.recv_timeout()
    }

    /// Local and remote address (for `accept` callers and tests).
    pub fn peer(&self) -> Option<Addr> {
        self.locked(|st, n| match (self.kind, st.handle) {
            (Kind::Tcp, Some(h)) => n.sockets.get::<tcp::Socket>(h).remote_endpoint().map(addr),
            _ => st.remote.map(addr),
        })
        .ok()
        .flatten()
    }
}

impl Drop for Socket {
    /// close(): TCP sends its FIN and lingers (net::reap_closing frees it
    /// once the goodbye handshake is done); the rest is freed at once.
    fn drop(&mut self) {
        let kind = self.kind;
        let _ = self.locked(|st, n| {
            for h in st.listeners.drain(..) {
                n.sockets.remove(h);
            }
            if let Some(h) = st.handle.take() {
                if kind == Kind::Tcp {
                    let sock = n.sockets.get_mut::<tcp::Socket>(h);
                    sock.close();
                    sock.set_timeout(Some(smoltcp::time::Duration::from_millis(CLOSE_TIMEOUT_MS)));
                    n.close_later(h);
                } else {
                    n.sockets.remove(h);
                }
            }
        });
        super::kick(); // let netd send the FIN (we may hold other locks here)
    }
}

impl File for Socket {
    fn read(&self, buf: &mut [u8]) -> Result<usize, Errno> {
        self.recv(buf).map(|(n, _)| n)
    }
    fn write(&self, buf: &[u8]) -> Result<usize, Errno> {
        self.send(buf, None)
    }
    fn stat(&self) -> Result<Stat, Errno> {
        Ok(Stat { kind: KIND_SOCKET, size: 0, mode: 0 })
    }
    fn as_socket(&self) -> Option<&Socket> {
        Some(self)
    }
    fn set_nonblocking(&self, on: bool) {
        self.nonblock.store(on, Ordering::Relaxed);
    }
    fn nonblocking(&self) -> bool {
        self.nonblock.load(Ordering::Relaxed)
    }
    /// What's ready right now (for poll/select/epoll).
    fn poll(&self) -> u16 {
        let kind = self.kind;
        self.locked(|st, n| {
            let mut ready = 0;
            match kind {
                Kind::Tcp if !st.listeners.is_empty() => {
                    // Listening: "readable" = a connection to accept.
                    let waiting = st.listeners.iter().any(|&h| {
                        !matches!(n.sockets.get::<tcp::Socket>(h).state(), tcp::State::Listen | tcp::State::SynReceived)
                    });
                    if waiting {
                        ready |= POLLIN;
                    }
                }
                Kind::Tcp => match st.handle {
                    // Never connected: nothing to read or write.
                    None => ready |= POLLOUT | POLLHUP,
                    Some(h) => {
                        let sock = n.sockets.get::<tcp::Socket>(h);
                        match sock.state() {
                            tcp::State::SynSent | tcp::State::SynReceived => {} // handshake going on
                            tcp::State::Closed => ready |= POLLIN | POLLERR | POLLHUP, // refused or reset
                            _ => {
                                if sock.can_recv() || !sock.may_recv() {
                                    ready |= POLLIN; // data, or EOF
                                }
                                if sock.can_send() || !sock.may_send() {
                                    ready |= POLLOUT; // room (or EPIPE waiting)
                                }
                                if !sock.may_recv() && !sock.may_send() {
                                    ready |= POLLHUP;
                                }
                            }
                        }
                    }
                },
                Kind::Udp => {
                    if let Some(h) = st.handle {
                        let sock = n.sockets.get::<udp::Socket>(h);
                        if sock.can_recv() {
                            ready |= POLLIN;
                        }
                        if sock.can_send() {
                            ready |= POLLOUT;
                        }
                    }
                }
                Kind::Icmp => {
                    if let Some(h) = st.handle {
                        let sock = n.sockets.get::<icmp::Socket>(h);
                        if sock.can_recv() {
                            ready |= POLLIN;
                        }
                        if sock.can_send() {
                            ready |= POLLOUT;
                        }
                    }
                }
            }
            ready
        })
        .unwrap_or(POLLERR)
    }
}