//! Filesystem layer with commit-on-close semantics.
//!
//! The layer wraps the wasmtime-wasi `wasi:filesystem` implementation. For
//! mounts with `commit_on_close`, an open for write goes to a hidden temporary
//! file next to the target. When the guest drops the descriptor (close), the
//! layer renames the temporary file into place. Until then, no other reader
//! sees the change. If the instance traps or is killed, the changes are lost.

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

use wasmtime::component::{HasData, Resource};
use wasmtime_wasi::filesystem::{Descriptor, WasiFilesystemCtxView};
use wasmtime_wasi::p2::bindings::filesystem::types::{
    self, ErrorCode, HostDescriptor, HostDirectoryEntryStream,
};
use wasmtime_wasi::p2::bindings::filesystem::preopens;
use wasmtime_wasi::p2::{DynInputStream, DynOutputStream, FsError, FsResult};

/// Name prefix of uncommitted files. The layer hides these names from guests.
pub const TMP_PREFIX: &str = ".sausage-tmp-";

#[derive(Clone, Debug)]
pub struct Mount {
    pub guest: String,
    pub host: PathBuf,
    pub writable: bool,
}

#[derive(Clone)]
struct DirInfo {
    host: PathBuf,
    root: PathBuf,
    commit: bool,
}

struct Pending {
    tmp: PathBuf,
    target: PathBuf,
}

/// Per-instance state of the layer.
pub struct FsExt {
    mounts: Vec<Mount>,
    dirs: HashMap<u32, DirInfo>,
    pending: HashMap<u32, Pending>,
    counter: u64,
}

impl FsExt {
    pub fn new(mounts: Vec<Mount>) -> Self {
        Self { mounts, dirs: HashMap::new(), pending: HashMap::new(), counter: 0 }
    }

    /// Commit every open write. Use this when the guest exits normally,
    /// because process exit closes all files.
    pub fn commit_all(&mut self) {
        for (_, p) in self.pending.drain() {
            commit(p);
        }
    }
}

impl Drop for FsExt {
    /// Discard writes that were not committed (trap, kill).
    fn drop(&mut self) {
        for (_, p) in self.pending.drain() {
            let _ = std::fs::remove_file(&p.tmp);
        }
    }
}

fn commit(p: Pending) {
    if let Err(e) = std::fs::rename(&p.tmp, &p.target) {
        eprintln!("sausage: commit of {} failed: {e}", p.target.display());
        let _ = std::fs::remove_file(&p.tmp);
    }
}

fn is_tmp_name(path: &str) -> bool {
    path.rsplit('/').next().is_some_and(|n| n.starts_with(TMP_PREFIX))
}

/// Join `rel` to the directory `dir` lexically. Return `None` if the result
/// leaves the mount root.
fn resolve(dir: &DirInfo, rel: &str) -> Option<PathBuf> {
    let mut out = dir.host.clone();
    for c in Path::new(rel).components() {
        match c {
            Component::Normal(n) => out.push(n),
            Component::CurDir => {}
            Component::ParentDir => {
                if !out.pop() {
                    return None;
                }
            }
            Component::RootDir | Component::Prefix(_) => return None,
        }
    }
    out.starts_with(&dir.root).then_some(out)
}

fn io_err(e: std::io::Error) -> FsError {
    ErrorCode::from(e).into()
}

pub struct ActorFs<'a> {
    pub inner: WasiFilesystemCtxView<'a>,
    pub ext: &'a mut FsExt,
}

pub struct ActorFsData;

impl HasData for ActorFsData {
    type Data<'a> = ActorFs<'a>;
}

impl preopens::Host for ActorFs<'_> {
    fn get_directories(&mut self) -> wasmtime::Result<Vec<(Resource<Descriptor>, String)>> {
        let dirs = preopens::Host::get_directories(&mut self.inner)?;
        for (r, name) in &dirs {
            if let Some(m) = self.ext.mounts.iter().find(|m| &m.guest == name) {
                self.ext.dirs.insert(
                    r.rep(),
                    DirInfo { host: m.host.clone(), root: m.host.clone(), commit: m.writable },
                );
            }
        }
        Ok(dirs)
    }
}

impl types::Host for ActorFs<'_> {
    fn convert_error_code(&mut self, err: FsError) -> wasmtime::Result<ErrorCode> {
        types::Host::convert_error_code(&mut self.inner, err)
    }

    fn filesystem_error_code(
        &mut self,
        err: Resource<wasmtime::Error>,
    ) -> wasmtime::Result<Option<ErrorCode>> {
        types::Host::filesystem_error_code(&mut self.inner, err)
    }
}

impl HostDescriptor for ActorFs<'_> {
    async fn advise(
        &mut self,
        fd: Resource<types::Descriptor>,
        offset: types::Filesize,
        len: types::Filesize,
        advice: types::Advice,
    ) -> FsResult<()> {
        HostDescriptor::advise(&mut self.inner, fd, offset, len, advice).await
    }

    async fn sync_data(&mut self, fd: Resource<types::Descriptor>) -> FsResult<()> {
        HostDescriptor::sync_data(&mut self.inner, fd).await
    }

    async fn get_flags(&mut self, fd: Resource<types::Descriptor>) -> FsResult<types::DescriptorFlags> {
        HostDescriptor::get_flags(&mut self.inner, fd).await
    }

    async fn get_type(&mut self, fd: Resource<types::Descriptor>) -> FsResult<types::DescriptorType> {
        HostDescriptor::get_type(&mut self.inner, fd).await
    }

    async fn set_size(&mut self, fd: Resource<types::Descriptor>, size: types::Filesize) -> FsResult<()> {
        HostDescriptor::set_size(&mut self.inner, fd, size).await
    }

    async fn set_times(
        &mut self,
        fd: Resource<types::Descriptor>,
        atim: types::NewTimestamp,
        mtim: types::NewTimestamp,
    ) -> FsResult<()> {
        HostDescriptor::set_times(&mut self.inner, fd, atim, mtim).await
    }

    async fn read(
        &mut self,
        fd: Resource<types::Descriptor>,
        len: types::Filesize,
        offset: types::Filesize,
    ) -> FsResult<(Vec<u8>, bool)> {
        HostDescriptor::read(&mut self.inner, fd, len, offset).await
    }

    async fn write(
        &mut self,
        fd: Resource<types::Descriptor>,
        buf: Vec<u8>,
        offset: types::Filesize,
    ) -> FsResult<types::Filesize> {
        HostDescriptor::write(&mut self.inner, fd, buf, offset).await
    }

    async fn read_directory(
        &mut self,
        fd: Resource<types::Descriptor>,
    ) -> FsResult<Resource<types::DirectoryEntryStream>> {
        HostDescriptor::read_directory(&mut self.inner, fd).await
    }

    async fn sync(&mut self, fd: Resource<types::Descriptor>) -> FsResult<()> {
        HostDescriptor::sync(&mut self.inner, fd).await
    }

    async fn create_directory_at(&mut self, fd: Resource<types::Descriptor>, path: String) -> FsResult<()> {
        HostDescriptor::create_directory_at(&mut self.inner, fd, path).await
    }

    async fn stat(&mut self, fd: Resource<types::Descriptor>) -> FsResult<types::DescriptorStat> {
        HostDescriptor::stat(&mut self.inner, fd).await
    }

    async fn stat_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        path_flags: types::PathFlags,
        path: String,
    ) -> FsResult<types::DescriptorStat> {
        if is_tmp_name(&path) {
            return Err(ErrorCode::NoEntry.into());
        }
        HostDescriptor::stat_at(&mut self.inner, fd, path_flags, path).await
    }

    async fn set_times_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        path_flags: types::PathFlags,
        path: String,
        atim: types::NewTimestamp,
        mtim: types::NewTimestamp,
    ) -> FsResult<()> {
        HostDescriptor::set_times_at(&mut self.inner, fd, path_flags, path, atim, mtim).await
    }

    async fn link_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        old_path_flags: types::PathFlags,
        old_path: String,
        new_descriptor: Resource<types::Descriptor>,
        new_path: String,
    ) -> FsResult<()> {
        HostDescriptor::link_at(&mut self.inner, fd, old_path_flags, old_path, new_descriptor, new_path).await
    }

    async fn open_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        path_flags: types::PathFlags,
        path: String,
        oflags: types::OpenFlags,
        flags: types::DescriptorFlags,
    ) -> FsResult<Resource<types::Descriptor>> {
        if is_tmp_name(&path) {
            return Err(ErrorCode::NoEntry.into());
        }
        let parent = self.ext.dirs.get(&fd.rep()).cloned();

        // Commit-on-close: redirect writes to a hidden temporary file.
        if let Some(dir) = parent.as_ref().filter(|d| d.commit)
            && flags.contains(types::DescriptorFlags::WRITE)
            && !oflags.contains(types::OpenFlags::DIRECTORY)
        {
            let target = resolve(dir, &path).ok_or(FsError::from(ErrorCode::NotPermitted))?;
            let meta = std::fs::metadata(&target).ok();
            if meta.as_ref().is_some_and(|m| m.is_dir()) {
                return Err(ErrorCode::IsDirectory.into());
            }
            let exists = meta.is_some();
            if exists && oflags.contains(types::OpenFlags::EXCLUSIVE) {
                return Err(ErrorCode::Exist.into());
            }
            if !exists && !oflags.contains(types::OpenFlags::CREATE) {
                return Err(ErrorCode::NoEntry.into());
            }
            let target_dir = target.parent().ok_or(FsError::from(ErrorCode::NotPermitted))?;
            if !target_dir.is_dir() {
                return Err(ErrorCode::NoEntry.into());
            }
            self.ext.counter += 1;
            let tmp_name = format!("{TMP_PREFIX}{}-{}", std::process::id(), self.ext.counter);
            let tmp_rel = match path.rsplit_once('/') {
                Some((d, _)) => format!("{d}/{tmp_name}"),
                None => tmp_name.clone(),
            };
            let tmp = target_dir.join(&tmp_name);
            let truncate = oflags.contains(types::OpenFlags::TRUNCATE);
            if exists && !truncate {
                std::fs::copy(&target, &tmp).map_err(io_err)?;
            }
            let oflags = (oflags | types::OpenFlags::CREATE)
                & !(types::OpenFlags::EXCLUSIVE | types::OpenFlags::TRUNCATE);
            return match HostDescriptor::open_at(&mut self.inner, fd, path_flags, tmp_rel, oflags, flags).await {
                Ok(r) => {
                    self.ext.pending.insert(r.rep(), Pending { tmp, target });
                    Ok(r)
                }
                Err(e) => {
                    let _ = std::fs::remove_file(&tmp);
                    Err(e)
                }
            };
        }

        let r = HostDescriptor::open_at(&mut self.inner, fd, path_flags, path.clone(), oflags, flags).await?;
        if let Some(dir) = parent
            && matches!(self.inner.table.get(&r)?, Descriptor::Dir(_))
            && let Some(host) = resolve(&dir, &path)
        {
            self.ext.dirs.insert(r.rep(), DirInfo { host, root: dir.root, commit: dir.commit });
        }
        Ok(r)
    }

    fn drop(&mut self, fd: Resource<types::Descriptor>) -> wasmtime::Result<()> {
        let rep = fd.rep();
        HostDescriptor::drop(&mut self.inner, fd)?;
        self.ext.dirs.remove(&rep);
        if let Some(p) = self.ext.pending.remove(&rep) {
            commit(p);
        }
        Ok(())
    }

    async fn readlink_at(&mut self, fd: Resource<types::Descriptor>, path: String) -> FsResult<String> {
        HostDescriptor::readlink_at(&mut self.inner, fd, path).await
    }

    async fn remove_directory_at(&mut self, fd: Resource<types::Descriptor>, path: String) -> FsResult<()> {
        HostDescriptor::remove_directory_at(&mut self.inner, fd, path).await
    }

    async fn rename_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        old_path: String,
        new_fd: Resource<types::Descriptor>,
        new_path: String,
    ) -> FsResult<()> {
        if is_tmp_name(&old_path) || is_tmp_name(&new_path) {
            return Err(ErrorCode::NoEntry.into());
        }
        HostDescriptor::rename_at(&mut self.inner, fd, old_path, new_fd, new_path).await
    }

    async fn symlink_at(
        &mut self,
        _fd: Resource<types::Descriptor>,
        _src_path: String,
        _dest_path: String,
    ) -> FsResult<()> {
        // Symlinks would break the lexical path tracking of commit-on-close.
        Err(ErrorCode::NotPermitted.into())
    }

    async fn unlink_file_at(&mut self, fd: Resource<types::Descriptor>, path: String) -> FsResult<()> {
        if is_tmp_name(&path) {
            return Err(ErrorCode::NoEntry.into());
        }
        HostDescriptor::unlink_file_at(&mut self.inner, fd, path).await
    }

    fn read_via_stream(
        &mut self,
        fd: Resource<types::Descriptor>,
        offset: types::Filesize,
    ) -> FsResult<Resource<DynInputStream>> {
        HostDescriptor::read_via_stream(&mut self.inner, fd, offset)
    }

    fn write_via_stream(
        &mut self,
        fd: Resource<types::Descriptor>,
        offset: types::Filesize,
    ) -> FsResult<Resource<DynOutputStream>> {
        HostDescriptor::write_via_stream(&mut self.inner, fd, offset)
    }

    fn append_via_stream(&mut self, fd: Resource<types::Descriptor>) -> FsResult<Resource<DynOutputStream>> {
        HostDescriptor::append_via_stream(&mut self.inner, fd)
    }

    async fn is_same_object(
        &mut self,
        a: Resource<types::Descriptor>,
        b: Resource<types::Descriptor>,
    ) -> wasmtime::Result<bool> {
        HostDescriptor::is_same_object(&mut self.inner, a, b).await
    }

    async fn metadata_hash(&mut self, fd: Resource<types::Descriptor>) -> FsResult<types::MetadataHashValue> {
        HostDescriptor::metadata_hash(&mut self.inner, fd).await
    }

    async fn metadata_hash_at(
        &mut self,
        fd: Resource<types::Descriptor>,
        path_flags: types::PathFlags,
        path: String,
    ) -> FsResult<types::MetadataHashValue> {
        HostDescriptor::metadata_hash_at(&mut self.inner, fd, path_flags, path).await
    }
}

impl HostDirectoryEntryStream for ActorFs<'_> {
    async fn read_directory_entry(
        &mut self,
        stream: Resource<types::DirectoryEntryStream>,
    ) -> FsResult<Option<types::DirectoryEntry>> {
        let rep = stream.rep();
        loop {
            let entry = HostDirectoryEntryStream::read_directory_entry(
                &mut self.inner,
                Resource::new_borrow(rep),
            )
            .await?;
            match entry {
                Some(e) if e.name.starts_with(TMP_PREFIX) => continue,
                other => return Ok(other),
            }
        }
    }

    fn drop(&mut self, stream: Resource<types::DirectoryEntryStream>) -> wasmtime::Result<()> {
        HostDirectoryEntryStream::drop(&mut self.inner, stream)
    }
}