tempest-io 0.0.1

use std::{
    collections::{HashMap, VecDeque},
    io,
    path::{Path, PathBuf},
};

use bytes::BytesMut;

use crate::{
    Completions, DirEntry, FstatHandle, Io, OpHandle, OpenOptions, ReadHandle, Statx, WriteHandle,
};

struct WasmFile {
    data: Vec<u8>,
    open_count: i32,
}

#[derive(Clone, Copy)]
pub struct FileStat {
    pub size: u64,
}

impl Statx for FileStat {
    fn stx_size(&self) -> u64 {
        self.size
    }
}

/// A fault to inject into [`WasmIo`] for testing purposes.
///
/// Faults are queued via [`WasmIo::inject`] and consumed in order,
/// one per matching operation. An empty queue means normal behavior.
pub enum Fault {
    /// Cause the next `read_at` to return fewer bytes than requested,
    /// simulating a short read from the kernel. The state machine must
    /// resubmit for the remainder - use this to test [`ReadExact`] and
    /// any phase that handles partial reads.
    PartialRead { bytes: usize },

    /// Cause the next `write_at` to return fewer bytes than requested,
    /// simulating a short write from the kernel. Use this to test
    /// [`WriteExact`] and any phase that handles partial writes.
    PartialWrite { bytes: usize },

    /// Cause the next submitted operation to fail with the given error kind.
    /// Use [`io::ErrorKind::WouldBlock`] to simulate a full submission queue,
    /// forcing state machines to restore their phase and retry next tick.
    FailNext { kind: io::ErrorKind },

    /// Reverse the order of completions returned by the next [`WasmIo::poll`],
    /// simulating out-of-order CQE delivery. State machines must handle
    /// completions arriving in any order within a tick.
    ReorderCompletions,
}

pub const BLOCK_SIZE: usize = 4096;

pub struct WasmIoConfig {
    /// Number of registered buffers that will be preallocated.
    pub buf_count: u16,
    pub report_unclosed_fds: bool,
}

impl WasmIoConfig {
    pub fn buf_count(mut self, c: u16) -> Self {
        self.buf_count = c;
        self
    }

    pub fn report_unclosed_fds(mut self, b: bool) -> Self {
        self.report_unclosed_fds = b;
        self
    }
}

impl Default for WasmIoConfig {
    fn default() -> Self {
        Self {
            buf_count: 64,
            report_unclosed_fds: true,
        }
    }
}

pub struct WasmIo {
    files: HashMap<PathBuf, WasmFile>,
    /// Maps a file descriptor as the index to the path.
    fds: Vec<Option<PathBuf>>,
    pending: VecDeque<(OpHandle, io::Result<u32>)>,
    completions: Completions,
    faults: VecDeque<Fault>,
    registered_bufs: VecDeque<BytesMut>,
    config: WasmIoConfig,
}

impl WasmIo {
    pub fn new(config: WasmIoConfig) -> Self {
        let mut registered_bufs = VecDeque::with_capacity(config.buf_count as usize);
        for _ in 0..config.buf_count {
            registered_bufs.push_back(BytesMut::with_capacity(BLOCK_SIZE));
        }
        Self {
            files: HashMap::new(),
            fds: Vec::new(),
            pending: VecDeque::new(),
            completions: Vec::new(),
            faults: VecDeque::new(),
            registered_bufs,
            config,
        }
    }

    pub fn inject(&mut self, fault: Fault) {
        self.faults.push_back(fault);
    }

    pub fn create_file_sync(&mut self, path: impl Into<PathBuf>) -> u32 {
        let path = path.into();
        self.files.insert(
            path.clone(),
            WasmFile {
                data: Vec::new(),
                open_count: 1,
            },
        );
        self.allocate_fd(path)
    }

    /// Synchronously close an fd without going through the async tick loop.
    ///
    /// Use this in test teardown to release fds that were obtained outside of
    /// a [`CloseFile`] state machine, so [`WasmIo`]'s drop does not warn
    /// about leaked descriptors.
    pub fn close_fd_sync(&mut self, fd: u32) {
        let path = self.fds[fd as usize]
            .take()
            .expect("close_fd_sync called on invalid fd");
        if let Some(file) = self.files.get_mut(&path) {
            file.open_count -= 1;
        }
    }

    fn allocate_fd(&mut self, path: PathBuf) -> u32 {
        if let Some(idx) = self.fds.iter().position(|slot| slot.is_none()) {
            self.fds[idx] = Some(path);
            idx as u32
        } else {
            let fd = self.fds.len() as u32;
            self.fds.push(Some(path));
            fd
        }
    }

    pub fn all_registered_bufs_released(&self) -> bool {
        self.registered_bufs.len() == self.config.buf_count as usize
    }

    pub fn all_fds_closed(&self) -> bool {
        self.fds.iter().filter(|&o| o.is_some()).count() == 0
    }
}

impl Default for WasmIo {
    fn default() -> Self {
        Self::new(WasmIoConfig::default())
    }
}

impl Io for WasmIo {
    fn block_size(&self) -> usize {
        BLOCK_SIZE
    }

    fn now(&self) -> std::time::Duration {
        // TODO: this is quite imprecise, but should not matter for the Wasm playground for now
        std::time::Duration::from_millis(js_sys::Date::now() as u64)
    }

    type Fd = u32;

    #[inline]
    unsafe fn into_fd(result: u32) -> Self::Fd {
        result
    }

    type RegisteredBuf = BytesMut;

    fn acquire_buf(&mut self) -> Option<Self::RegisteredBuf> {
        self.registered_bufs.pop_front()
    }

    fn release_buf(&mut self, buf: Self::RegisteredBuf) {
        self.registered_bufs.push_back(buf);
    }

    type Statx = FileStat;

    fn fstat(&mut self, fd: Self::Fd, handle: OpHandle) -> io::Result<FstatHandle<Self::Statx>> {
        let path = self.fds[fd as usize]
            .as_ref()
            .expect("fstat called on invalid fd");

        let file = self
            .files
            .get_mut(path)
            .expect("fd points to a file that no longer exists");

        // SAFETY: we only access stx_size, which gets initialized below
        let size = file.data.len() as u64;
        let statx = Box::new(FileStat { size });

        self.pending.push_back((handle, Ok(0)));
        Ok(FstatHandle::new(statx))
    }

    fn open(&mut self, path: &Path, opts: OpenOptions, handle: OpHandle) -> io::Result<()> {
        if opts.create_new && self.files.contains_key(path) {
            self.pending.push_back((
                handle,
                Err(io::Error::new(
                    io::ErrorKind::AlreadyExists,
                    "file already exists",
                )),
            ));
            return Ok(());
        }

        if !opts.create && !opts.create_new && !self.files.contains_key(path) {
            self.pending.push_back((
                handle,
                Err(io::Error::new(io::ErrorKind::NotFound, "file not found")),
            ));
            return Ok(());
        }

        let file = self
            .files
            .entry(path.to_owned())
            .or_insert_with(|| WasmFile {
                data: Vec::new(),
                open_count: 0,
            });

        if opts.truncate {
            file.data.clear();
        }

        file.open_count += 1;
        let fd = self.allocate_fd(path.to_owned());
        self.pending.push_back((handle, Ok(fd)));
        Ok(())
    }

    fn close(&mut self, fd: Self::Fd, handle: OpHandle) -> io::Result<()> {
        let path = self.fds[fd as usize]
            .take()
            .expect("close called on invalid fd");

        if let Some(file) = self.files.get_mut(&path) {
            file.open_count -= 1;
        }

        self.pending.push_back((handle, Ok(0)));
        Ok(())
    }

    fn read_at<B: crate::IoBufMut>(
        &mut self,
        fd: Self::Fd,
        mut buf: B,
        offset: u64,
        handle: OpHandle,
    ) -> Result<ReadHandle<B>, (std::io::Error, B)> {
        let path = self.fds[fd as usize]
            .as_ref()
            .expect("write_at called on invalid fd");

        let file = self
            .files
            .get_mut(path)
            .expect("fd points to a file that no longer exists");

        let offset = offset as usize;
        let available = file.data.len().saturating_sub(offset);
        // ensure we do not overflow the valid buffer lengths,
        // but instead just do a short read instead
        let mut len = buf.bytes_total().min(available);

        // clamp to fault if injected
        if let Some(&Fault::PartialRead { bytes }) = self.faults.front() {
            self.faults.pop_front();
            len = len.min(bytes)
        }

        // SAFETY: the two buffers can not overlap, they are valid for creating pointers,
        // and buf fits at least `len` bytes at the end, due to clamping above if required
        unsafe {
            std::ptr::copy_nonoverlapping(file.data[offset..].as_ptr(), buf.stable_mut_ptr(), len);
        }

        // SAFETY: we just copied the bytes into the buffer from the file above,
        // so they may now be considered initialized
        unsafe {
            if len > buf.bytes_init() {
                buf.set_init(len);
            }
        }

        self.pending.push_back((handle, Ok(len as u32)));
        Ok(ReadHandle::new(buf))
    }

    fn write_at<B: crate::IoBuf>(
        &mut self,
        fd: Self::Fd,
        buf: B,
        offset: u64,
        handle: OpHandle,
    ) -> Result<WriteHandle<B>, (std::io::Error, B)> {
        let path = self.fds[fd as usize]
            .as_ref()
            .expect("write_at called on invalid fd");

        let file = self
            .files
            .get_mut(path)
            .expect("fd points to a file that no longer exists");

        let offset = offset as usize;
        let mut len = buf.bytes_init();

        // clamp to fault if injected
        if let Some(&Fault::PartialWrite { bytes }) = self.faults.front() {
            self.faults.pop_front();
            len = len.min(bytes)
        }

        let required = offset + len;

        if file.data.len() < required {
            file.data.resize(required, 0);
        }

        // SAFETY: the two buffers can not overlap, they are valid for creating pointers,
        // and file fits at least `len` bytes at the end, due to resizing above if required
        unsafe {
            std::ptr::copy_nonoverlapping(buf.stable_ptr(), file.data[offset..].as_mut_ptr(), len);
        }

        self.pending.push_back((handle, Ok(len as u32)));
        Ok(WriteHandle::new(buf))
    }

    fn fsync(&mut self, fd: Self::Fd, handle: OpHandle) -> io::Result<()> {
        let path = self.fds[fd as usize]
            .as_ref()
            .expect("fsync called on invalid fd");

        assert!(
            self.files.contains_key(path),
            "fd points to a file that no longer exists"
        );

        self.pending.push_back((handle, Ok(0)));
        Ok(())
    }

    fn rename(&mut self, from: &Path, to: &Path, handle: OpHandle) -> io::Result<()> {
        let file = self
            .files
            .remove(from)
            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "source file not found"))?;
        self.files.insert(to.to_owned(), file);

        for slot in self.fds.iter_mut().flatten() {
            if slot == from {
                *slot = to.to_owned();
            }
        }

        self.pending.push_back((handle, Ok(0)));
        Ok(())
    }

    fn remove(&mut self, path: &Path, handle: OpHandle) -> io::Result<()> {
        if self.files.remove(path).is_none() {
            self.pending.push_back((
                handle,
                Err(io::Error::new(io::ErrorKind::NotFound, "file not found")),
            ));
        } else {
            self.pending.push_back((handle, Ok(0)));
        }

        Ok(())
    }

    fn poll(&mut self) -> io::Result<()> {
        self.completions.extend(self.pending.drain(..));
        Ok(())
    }

    fn in_flight(&self) -> usize {
        self.pending.len()
    }

    fn park(&mut self) -> io::Result<()> {
        Ok(())
    }

    fn completions(&mut self) -> &mut Completions {
        &mut self.completions
    }

    fn list_dir(&self, path: &Path) -> io::Result<Vec<DirEntry>> {
        let mut prefix = path.to_string_lossy().into_owned();
        if !prefix.ends_with('/') {
            prefix.push('/');
        }

        let mut seen = HashMap::new();

        for file_path in self.files.keys() {
            let file_str = file_path.to_string_lossy();

            if !file_str.starts_with(&prefix) {
                continue;
            }

            let relative = &file_str[prefix.len()..];
            let (name, is_dir) = match relative.find('/') {
                Some(idx) => (&relative[..idx], true),
                None => (relative, false),
            };

            let entry_path = PathBuf::from(&prefix).join(name);
            seen.entry(entry_path.clone()).or_insert(DirEntry {
                path: entry_path,
                is_dir,
            });
        }

        Ok(seen.into_values().collect())
    }

    fn create_dir_all(&self, _path: &Path) -> io::Result<()> {
        Ok(())
    }
}

impl Drop for WasmIo {
    fn drop(&mut self) {
        if self.config.report_unclosed_fds {
            for (i, slot) in self.fds.iter().enumerate() {
                if let Some(path) = slot {
                    tracing::warn!("WasmIo dropped with open fd {} to {:?}", i, path)
                }
            }
        }
    }
}