//! Networking: the virtio-net card + the smoltcp TCP/IP stack.
//!
//! The layers, bottom to top (see the README's networking chapter):
//!   * virtio_net.rs  moves raw Ethernet frames in and out of the card.
//!   * device.rs      what smoltcp sees: the card plus a loopback queue,
//!                    so 127.0.0.1 works (even with no card at all).
//!   * smoltcp        does everything protocol-shaped: Ethernet addressing,
//!                    ARP (who has 10.0.2.2?), IPv4, ICMP (ping), UDP, TCP
//!                    (connections, retransmits, windows), DHCP and DNS.
//!   * this file      owns the one smoltcp `Interface` and its sockets, runs
//!                    the `netd` kernel thread that keeps smoltcp going, and
//!                    reacts to DHCP (configure our IP address and route).
//!   * socket.rs      the BSD-style sockets user programs get as fds.
//!
//! How packets keep flowing: smoltcp does nothing by itself. Someone must
//! call `Interface::poll` whenever a frame arrives, whenever a program sent
//! data, and whenever a protocol timer (TCP retransmit, DHCP renew) is due.
//!   * frame arrives -> the card raises its IRQ -> `handle_irq` wakes netd
//!   * a program sends -> socket.rs calls `poll()` right away
//!   * timers -> netd sleeps only until smoltcp's `poll_delay`
//! After every poll we wake `CHANGED`, the wait queue that blocked socket
//! calls (recv, accept, connect...) sleep on; they re-check their socket.
//!
//! Locking: `NET` is taken with interrupts off (wait-queue conditions look
//! at sockets with interrupts off; see the locking rule in thread.rs).

pub mod device;
pub mod socket;
pub mod virtio_net;

use crate::thread::WaitQueue;
use alloc::string::String;
use alloc::vec;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicBool, AtomicU16, Ordering};
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet};
use smoltcp::socket::{dhcpv4, dns, tcp};
use smoltcp::time::Instant;
use smoltcp::wire::{DnsQueryType, EthernetAddress, HardwareAddress, IpAddress, IpCidr, Ipv4Address};
use spin::Mutex;
use device::NetDevice;
use virtio_net::VirtioNet;
use x86_64::instructions::interrupts::without_interrupts;
use x86_64::instructions::port::Port;

/// Our IPv4 configuration, as handed out by DHCP.
#[derive(Clone, Default)]
pub struct IpConfig {
    pub address: [u8; 4],
    pub prefix_len: u8,
    pub gateway: Option<[u8; 4]>,
    pub dns: Vec<[u8; 4]>,
}

pub struct Net {
    pub iface: Interface,
    pub device: NetDevice,
    pub sockets: SocketSet<'static>,
    dhcp: SocketHandle,
    dns: SocketHandle,
    pub config: Option<IpConfig>,
    /// TCP sockets whose fd was closed: they live on until the FIN
    /// handshake is done, then `poll` removes them.
    closing: Vec<SocketHandle>,
    next_port: u16,
}

static NET: Mutex<Option<Net>> = Mutex::new(None);
/// Socket state may have changed: blocked socket calls wait on this.
pub static CHANGED: WaitQueue = WaitQueue::new();
/// netd sleeps here; the IRQ handler wakes it.
static NETD_WAKE: WaitQueue = WaitQueue::new();
static IRQ_PENDING: AtomicBool = AtomicBool::new(false);
/// The card's I/O base, so the IRQ handler can read its ISR register
/// without taking the NET lock.
static IO_BASE: AtomicU16 = AtomicU16::new(0);

/// Run `f` on the network state (None before `init`).
pub fn with_net<R>(f: impl FnOnce(&mut Net) -> R) -> Option<R> {
    without_interrupts(|| NET.lock().as_mut().map(f))
}

/// smoltcp's idea of "now".
pub fn now() -> Instant {
    Instant::from_micros(crate::time::now_us() as i64)
}

/// The loopback address, 127.0.0.1/8.
fn loopback_cidr() -> IpCidr {
    IpCidr::new(IpAddress::Ipv4(Ipv4Address::new(127, 0, 0, 1)), 8)
}

/// Is this address on our loopback network (127.x.x.x)?
pub fn is_loopback(ip: [u8; 4]) -> bool {
    ip[0] == 127
}

/// Set up smoltcp on the first virtio-net card (if any) plus loopback,
/// start DHCP and the netd thread. Returns lines for the boot log.
pub fn init(devices: &[crate::pci::PciDevice]) -> Vec<String> {
    let nic = devices.iter().find(|d| virtio_net::is_virtio_net(d)).and_then(VirtioNet::new);
    let irq = nic.as_ref().map(|n| n.irq);
    let io_base = nic.as_ref().map_or(0, |n| n.io_base);
    let mut device = NetDevice::new(nic);
    let mac = device.mac;
    let mut config = Config::new(HardwareAddress::Ethernet(EthernetAddress(mac)));
    // Random seed: TCP's initial sequence numbers and DHCP transaction ids.
    // From virtio-rng (rng.rs) -- guessable ones would let an attacker
    // inject packets into our connections.
    config.random_seed = crate::virtio::rng::u64();
    let mut iface = Interface::new(config, &mut device, now());
    iface.update_ip_addrs(|addrs| {
        addrs.push(loopback_cidr()).ok();
    });
    let mut sockets = SocketSet::new(Vec::new());
    let dhcp = sockets.add(dhcpv4::Socket::new());
    let dns = sockets.add(dns::Socket::new(&[], vec![None, None, None, None]));
    // Ephemeral (client) ports start at a random spot in 49152..65000, so
    // they're not predictable either.
    let first_port = 49152 + (crate::virtio::rng::u64() % (65000 - 49152)) as u16;
    IO_BASE.store(io_base, Ordering::Relaxed);
    without_interrupts(|| {
        *NET.lock() = Some(Net { iface, device, sockets, dhcp, dns, config: None, closing: Vec::new(), next_port: first_port });
    });
    crate::thread::spawn("netd", netd);
    let mut log = alloc::vec![String::from("lo: loopback, 127.0.0.1/8")];
    if let Some(irq) = irq {
        crate::interrupts::enable_pci_irq(irq);
        log.push(alloc::format!(
            "eth0: virtio-net, MAC {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}, IRQ {}",
            mac[0], mac[1], mac[2], mac[3], mac[4], mac[5], irq
        ));
    }
    log
}

/// Interrupt handler part (called by interrupts.rs for PCI IRQ lines).
/// Reading the ISR register acknowledges the card's interrupt; bit 0 means
/// "I used a queue". We don't touch smoltcp here (interrupt handlers must
/// be short); we just wake netd.
pub fn handle_irq() {
    let io = IO_BASE.load(Ordering::Relaxed);
    if io == 0 {
        return;
    }
    let isr = unsafe { Port::<u8>::new(io + virtio_net::REG_ISR_STATUS).read() };
    if isr & 1 != 0 {
        IRQ_PENDING.store(true, Ordering::Release);
        NETD_WAKE.wake_all();
    }
}

/// Ask netd to poll soon (safe to call while holding other locks).
pub fn kick() {
    IRQ_PENDING.store(true, Ordering::Release);
    NETD_WAKE.wake_all();
}

/// Let smoltcp process everything that's pending: received frames,
/// queued data to send, timers. Then wake blocked socket calls (and
/// poll/epoll waiters).
pub fn poll() {
    with_net(|n| {
        // Frames smoltcp sent to itself (loopback) are only received on
        // the *next* poll, so go round until lo is empty. Bounded, so
        // two sockets chattering away can't keep us here forever.
        for _ in 0..64 {
            n.iface.poll(now(), &mut n.device, &mut n.sockets);
            if !n.device.lo_pending() {
                break;
            }
        }
        n.check_dhcp();
        n.reap_closing();
    });
    CHANGED.wake_all();
    crate::poll::notify();
}

/// The network thread: poll, sleep until a frame arrives or a timer is
/// due, repeat. Never spins: at most it wakes every PIT tick.
fn netd() {
    loop {
        poll();
        let delay_us = with_net(|n| n.iface.poll_delay(now(), &n.sockets)).flatten().map(|d| d.total_micros());
        // Wake up at least once a second anyway (cheap insurance against
        // a lost interrupt: we also re-read the ISR then).
        let deadline = crate::time::deadline_after_us(delay_us.unwrap_or(1_000_000).min(1_000_000));
        NETD_WAKE.wait_until_deadline(deadline, || IRQ_PENDING.swap(false, Ordering::Acquire).then_some(()));
        handle_irq();
    }
}

impl Net {
    /// React to DHCP news: configure (or unconfigure) our address, default
    /// route and DNS servers.
    fn check_dhcp(&mut self) {
        let event = self.sockets.get_mut::<dhcpv4::Socket>(self.dhcp).poll();
        let new_config = match event {
            None => return,
            Some(dhcpv4::Event::Deconfigured) => None,
            Some(dhcpv4::Event::Configured(c)) => Some(IpConfig {
                address: c.address.address().octets(),
                prefix_len: c.address.prefix_len(),
                gateway: c.router.map(|r| r.octets()),
                dns: c.dns_servers.iter().map(|d| d.octets()).collect(),
            }),
        };
        // DHCP's address first: smoltcp uses the first address as the
        // source for destinations outside both networks.
        self.iface.update_ip_addrs(|addrs| {
            addrs.clear();
            if let Some(c) = &new_config {
                let ip = Ipv4Address::from(c.address);
                addrs.push(IpCidr::new(IpAddress::Ipv4(ip), c.prefix_len)).ok();
            }
            addrs.push(loopback_cidr()).ok();
        });
        self.device.own_ip = new_config.as_ref().map(|c| c.address);
        self.iface.routes_mut().remove_default_ipv4_route();
        let mut servers = Vec::new();
        if let Some(c) = &new_config {
            if let Some(gw) = c.gateway {
                self.iface.routes_mut().add_default_ipv4_route(Ipv4Address::from(gw)).ok();
            }
            servers = c.dns.iter().map(|d| IpAddress::Ipv4(Ipv4Address::from(*d))).collect();
            crate::kprintln!(
                "[net] DHCP lease: {}/{}, gateway {}, DNS {}",
                Ipv4Address::from(c.address), c.prefix_len,
                c.gateway.map(Ipv4Address::from).unwrap_or(Ipv4Address::UNSPECIFIED),
                c.dns.first().map(|d| Ipv4Address::from(*d)).unwrap_or(Ipv4Address::UNSPECIFIED),
            );
        } else if self.config.is_some() {
            crate::kprintln!("[net] DHCP lease lost");
        }
        self.sockets.get_mut::<dns::Socket>(self.dns).update_servers(&servers);
        self.config = new_config;
    }

    fn reap_closing(&mut self) {
        let sockets = &mut self.sockets;
        self.closing.retain(|&h| {
            let state = sockets.get::<tcp::Socket>(h).state();
            let done = matches!(state, tcp::State::Closed | tcp::State::TimeWait);
            if done {
                sockets.remove(h);
            }
            !done
        });
    }

    /// Hand a closed TCP socket to the reaper (after its FIN goes out).
    pub fn close_later(&mut self, handle: SocketHandle) {
        self.closing.push(handle);
    }

    /// A free local port for outgoing connections (49152.. like Linux's
    /// ephemeral range).
    pub fn ephemeral_port(&mut self) -> u16 {
        let port = self.next_port;
        self.next_port = if port >= 65000 { 49152 } else { port + 1 };
        port
    }
}

/// Is there a network card at all?
pub fn present() -> bool {
    with_net(|n| n.device.nic.is_some()).unwrap_or(false)
}

/// Wait (up to `timeout_us`) until DHCP has given us an address.
pub fn wait_configured(timeout_us: u64) -> Result<IpConfig, crate::errno::Errno> {
    use crate::errno::Errno;
    if !present() {
        return Err(Errno::ENETDOWN);
    }
    let deadline = crate::time::deadline_after_us(timeout_us);
    CHANGED.wait_until_deadline(deadline, || with_net(|n| n.config.clone()).flatten()).ok_or(Errno::ENETUNREACH)
}

/// Number of live smoltcp sockets (DHCP and DNS included); tests use it to
/// check that closed sockets really go away.
pub fn socket_count() -> usize {
    with_net(|n| n.sockets.iter().count()).unwrap_or(0)
}

/// Resolve a host name to IPv4 addresses with DNS (smoltcp's DNS socket
/// sends the UDP query to the DHCP-provided server and parses the answer).
pub fn resolve(name: &str, timeout_us: u64) -> Result<Vec<[u8; 4]>, crate::errno::Errno> {
    use crate::errno::Errno;
    wait_configured(5_000_000)?;
    let query = with_net(|n| {
        let dns = n.sockets.get_mut::<dns::Socket>(n.dns);
        dns.start_query(n.iface.context(), name, DnsQueryType::A)
    })
    .ok_or(Errno::ENETDOWN)?
    .map_err(|_| Errno::EINVAL)?;
    poll(); // send the query now
    let deadline = crate::time::deadline_after_us(timeout_us);
    let result = CHANGED.wait_until_deadline(deadline, || {
        with_net(|n| match n.sockets.get_mut::<dns::Socket>(n.dns).get_query_result(query) {
            Err(dns::GetQueryResultError::Pending) => None,
            Err(dns::GetQueryResultError::Failed) => Some(Err(Errno::ENOENT)),
            Ok(addrs) => Some(Ok(addrs
                .iter()
                .filter_map(|a| match a {
                    IpAddress::Ipv4(v4) => Some(v4.octets()),
                    #[allow(unreachable_patterns)]
                    _ => None,
                })
                .collect())),
        })
        .flatten()
    });
    match result {
        Some(r) => r,
        None => {
            with_net(|n| n.sockets.get_mut::<dns::Socket>(n.dns).cancel_query(query));
            Err(Errno::ETIMEDOUT)
        }
    }
}

/// What `ifconfig` shows (syscall 504 copies it to user space).
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct NetInfo {
    pub present: u8,
    pub configured: u8,
    pub mac: [u8; 6],
    pub address: [u8; 4],
    pub prefix_len: u8,
    pub irq: u8,
    pub _pad: [u8; 2],
    pub gateway: [u8; 4],
    pub dns: [u8; 4],
    pub rx_packets: u64,
    pub rx_bytes: u64,
    pub tx_packets: u64,
    pub tx_bytes: u64,
}

pub fn info() -> NetInfo {
    with_net(|n| {
        let Some(nic) = &n.device.nic else { return NetInfo::default() };
        let s = nic.stats;
        let mut info = NetInfo {
            present: 1,
            mac: nic.mac,
            irq: nic.irq,
            rx_packets: s.rx_packets,
            rx_bytes: s.rx_bytes,
            tx_packets: s.tx_packets,
            tx_bytes: s.tx_bytes,
            ..Default::default()
        };
        if let Some(c) = &n.config {
            info.configured = 1;
            info.address = c.address;
            info.prefix_len = c.prefix_len;
            info.gateway = c.gateway.unwrap_or_default();
            info.dns = c.dns.first().copied().unwrap_or_default();
        }
        info
    })
    .unwrap_or_default()
}