//! Shares compiled programs and single-flight compilation across worlds in one host process.
use super::{Compiled, State, adapter, bindings};
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::{Arc, Mutex, OnceLock},
};
use tokio::sync::watch;
type Compilation = watch::Receiver<Option<Result<Compiled, String>>>;
use wasmtime::{
    Engine,
    component::{Component, Linker},
};
use wasmtime_wasi::p2::bindings::CommandPre;

static DIRECTORY: OnceLock<PathBuf> = OnceLock::new();
static SHARED: Mutex<Option<Arc<Compiler>>> = Mutex::new(None);

/// Select the host user's cache before opening a world.
pub fn configure(directory: PathBuf) -> anyhow::Result<()> {
    DIRECTORY
        .set(directory)
        .map_err(|_| anyhow::anyhow!("cache directory already configured"))
}
fn directory() -> anyhow::Result<PathBuf> {
    if let Some(dir) = DIRECTORY.get() {
        return Ok(dir.clone());
    }
    if let Some(dir) = std::env::var_os("SAUSAGE_CACHE") {
        return Ok(dir.into());
    }
    if let Some(dir) = std::env::var_os("XDG_CACHE_HOME") {
        return Ok(PathBuf::from(dir).join("sausage"));
    }
    Ok(
        PathBuf::from(std::env::var_os("HOME").ok_or_else(|| anyhow::anyhow!("HOME is not set"))?)
            .join(".cache/sausage"),
    )
}
pub(super) fn shared() -> anyhow::Result<Arc<Compiler>> {
    let mut shared = SHARED.lock().unwrap();
    if shared.is_none() {
        *shared = Some(Compiler::new(&directory()?)?);
    }
    Ok(shared.as_ref().unwrap().clone())
}
pub(super) struct Compiler {
    pub engine: Engine,
    programs: Mutex<HashMap<Vec<u8>, Compilation>>,
    #[cfg(test)]
    before_compile: Mutex<Option<Box<dyn FnOnce() + Send>>>,
    #[cfg(test)]
    compilations: std::sync::atomic::AtomicUsize,
}
impl Compiler {
    pub fn new(directory: &Path) -> anyhow::Result<Arc<Self>> {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
        let mut config = wasmtime::Config::new();
        config.epoch_interruption(true);
        // Keep Cranelift's default optimization: compiles are cached per content, and guests run
        // many times. OptLevel::None saved only 12.7% of compile time (T15) but slows every run.
        let mut cache = wasmtime::CacheConfig::new();
        cache.with_directory(directory);
        config.cache(Some(wasmtime::Cache::new(cache)?));
        Ok(Arc::new(Self {
            engine: Engine::new(&config)?,
            programs: Mutex::new(HashMap::new()),
            #[cfg(test)]
            before_compile: Mutex::new(None),
            #[cfg(test)]
            compilations: std::sync::atomic::AtomicUsize::new(0),
        }))
    }
    pub async fn get(self: &Arc<Self>, bytes: Vec<u8>) -> anyhow::Result<Compiled> {
        let mut result = self
            .programs
            .lock()
            .unwrap()
            .entry(bytes.clone())
            .or_insert_with(|| {
                let (tx, rx) = watch::channel(None);
                let compiler = self.clone();
                // The job owns completion even if a command or a precompile waiter is cancelled.
                tokio::spawn(async move {
                    let result =
                        tokio::task::spawn_blocking(move || compiler.compile(&bytes)).await;
                    tx.send_replace(Some(
                        result
                            .map_err(anyhow::Error::from)
                            .and_then(|r| r)
                            .map_err(|e| format!("{e:#}")),
                    ));
                });
                rx
            })
            .clone();
        loop {
            if let Some(result) = result.borrow().clone() {
                return result.map_err(|e| anyhow::anyhow!(e));
            }
            result.changed().await?;
        }
    }
    fn compile(&self, bytes: &[u8]) -> anyhow::Result<Compiled> {
        #[cfg(test)]
        {
            self.compilations
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            if let Some(hook) = self.before_compile.lock().unwrap().take() {
                hook();
            }
        }
        if bytes.get(..8) == Some(b"\0asm\x01\0\0\0") {
            let module = wasmtime::Module::new(&self.engine, bytes)?;
            let mut linker = wasmtime::Linker::new(&self.engine);
            sausage_vfs::preview1::add_to_linker(&mut linker, |s: &mut State| {
                s.preview1.as_mut().unwrap()
            })?;
            Ok(Compiled::Module(linker.instantiate_pre(&module)?))
        } else {
            let component = Component::new(&self.engine, bytes)?;
            let mut linker = Linker::new(&self.engine);
            wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
            wasmtime_wasi_http::p2::add_only_http_to_linker_async(&mut linker)?;
            adapter::add_to_linker(&mut linker, |s: &mut State| adapter::View {
                adapter: &mut s.fs,
                table: &mut s.table,
            })?;
            bindings::sausage::process::host::add_to_linker::<_, wasmtime::component::HasSelf<_>>(
                &mut linker,
                |s| s,
            )?;
            Ok(Compiled::Component(CommandPre::new(
                linker.instantiate_pre(&component)?,
            )?))
        }
    }
}
/// Compile programs in both command folders without blocking world startup.
pub(super) fn precompile(
    compiler: Arc<Compiler>,
    root: Arc<dyn sausage_vfs::Layer>,
) -> tokio::task::JoinHandle<anyhow::Result<usize>> {
    tokio::spawn(async move {
        let mut tasks = tokio::task::JoinSet::new();
        // Leave CPU and file bandwidth for commands during startup.
        let slots = Arc::new(tokio::sync::Semaphore::new(2));
        for folder in ["/bin", "/usr/local/bin"] {
            let entries = match root.list(folder).await {
                Ok(entries) => entries,
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue,
                Err(e) => return Err(e.into()),
            };
            for entry in entries.into_iter().filter(|entry| !entry.directory) {
                let root = root.clone();
                let compiler = compiler.clone();
                let path = format!("{folder}/{}", entry.name);
                let slots = slots.clone();
                tasks.spawn(async move {
                    let _permit = slots.acquire().await?;
                    let file = root
                        .open(
                            &path,
                            sausage_vfs::Open {
                                follow: true,
                                ..Default::default()
                            },
                        )
                        .await?;
                    let bytes = file.read(0, file.stat().await?.size as usize).await?;
                    if bytes.get(..4) != Some(b"\0asm") {
                        return Ok::<_, anyhow::Error>(false);
                    }
                    compiler.get(bytes).await?;
                    Ok(true)
                });
            }
        }
        let mut count = 0;
        while let Some(result) = tasks.join_next().await {
            if result?? {
                count += 1;
            }
        }
        Ok(count)
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::exec::{Command, Exec, Options};
    use std::{future::Future, sync::atomic::Ordering};

    #[tokio::test(flavor = "multi_thread")]
    async fn command_joins_precompile_even_after_waiter_cancellation() {
        let temp = tempfile::tempdir().unwrap();
        let compiler = Compiler::new(&temp.path().join("cache")).unwrap();
        let world = crate::world::World::open(&temp.path().join("world")).unwrap();
        let bytes =
            wat::parse_str(r#"(module (memory (export "memory") 1) (func (export "_start")))"#)
                .unwrap();
        std::fs::create_dir_all(temp.path().join("world/world/bin")).unwrap();
        std::fs::write(temp.path().join("world/world/bin/program.wasm"), &bytes).unwrap();
        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
        let (finish_tx, finish_rx) = std::sync::mpsc::channel();
        *compiler.before_compile.lock().unwrap() = Some(Box::new(move || {
            started_tx.send(()).unwrap();
            finish_rx.recv().unwrap();
        }));
        // Startup returns while the compiler remains blocked behind the test gate.
        let runtime = super::super::Runtime::with_compiler(world, compiler.clone()).unwrap();
        started_rx.await.unwrap();
        let background = compiler.clone();
        let precompile = tokio::spawn(async move { background.get(bytes).await });
        let mut output = |_: crate::exec::Stream, _: &[u8]| {};
        let command = runtime.exec(
            Command::Argv(vec!["/bin/program".into()]),
            Options::default(),
            &mut output,
        );
        tokio::pin!(command);
        let pending = std::future::poll_fn(|cx| {
            std::task::Poll::Ready(command.as_mut().poll(cx).is_pending())
        })
        .await;
        assert!(pending);
        assert_eq!(compiler.compilations.load(Ordering::Relaxed), 1);
        precompile.abort();
        let _ = precompile.await;
        finish_tx.send(()).unwrap();
        assert_eq!(command.await.unwrap(), 0);
        assert_eq!(compiler.compilations.load(Ordering::Relaxed), 1);
        assert!(!temp.path().join("world/host/compile-cache").exists());
    }
}