//! Runs WASI programs in independent stores over the world root. Owns process groups, stdio, and Exec.
use crate::{
    exec::{Command, Exec, Options, Stream},
    world::World,
};
use sausage_vfs::{
    Mounts,
    adapter::{self, Adapter},
};
use std::{
    collections::HashMap,
    path::Path,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicU32, Ordering},
    },
};
use tokio::{
    io::{AsyncReadExt, DuplexStream},
    sync::watch,
    task::AbortHandle,
};
use wasmtime::{
    Engine, Store,
    component::{Resource, ResourceTable},
};
use wasmtime_wasi::{
    WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView,
    cli::AsyncStdoutStream,
    p2::{
        DynInputStream, DynOutputStream,
        bindings::CommandPre,
        pipe::{AsyncReadStream, AsyncWriteStream},
    },
};

mod bindings {
    wasmtime::component::bindgen!({
        path: "wit", world: "program",
        imports: { default: async | trappable },
        with: {
            "wasi:io": wasmtime_wasi::p2::bindings::io,
            "sausage:process/host.child": super::Child,
            "sausage:process/host.watcher": super::Watcher,
        },
    });
}
use bindings::sausage::process::host::{Host, HostChild, Request, Stdio};
mod cache;
mod http;
pub use http::Router;
pub mod net;
/// Credential scripts use T28's bounded QuickJS component.
pub fn default_script() -> Arc<dyn crate::cred::script::Script> {
    Arc::new(crate::cred::quickjs::QuickJs::default())
}
pub use cache::configure as configure_cache;
mod services;
mod watcher;
pub use watcher::Watcher;
mod stdio;
pub use services::supervise;
use stdio::Io;

#[derive(Clone)]
enum Compiled {
    Component(CommandPre<State>),
    Module(wasmtime::InstancePre<State>),
}

/// A finite snapshot of process metadata exposed through /proc/processes.
#[derive(Clone, Debug)]
pub struct ProcessInfo {
    pub parent: u32,
    pub argv: Vec<String>,
    pub cwd: String,
    pub state: String,
    pub conversation: Option<String>,
}
pub type Processes = Arc<Mutex<std::collections::BTreeMap<u32, ProcessInfo>>>;

struct Entry {
    command: u32,
    group: u32,
    cancelled: Arc<AtomicBool>,
    abort: AbortHandle,
}
/// One runtime keeps compiled components warm and shares the world's namespace across stores.
pub struct Runtime {
    world: World,
    engine: Engine,
    compiler: Arc<cache::Compiler>,
    cache: Mutex<std::collections::HashSet<Vec<u8>>>,
    entries: Mutex<HashMap<u32, Entry>>,
    next: AtomicU32,
    alive: Arc<AtomicBool>,
    cred: Arc<crate::cred::Cred>,
    script: Arc<dyn crate::cred::script::Script>,
    router: http::Router,
    pub network: Arc<net::NetPolicy>,
}
impl Drop for Runtime {
    fn drop(&mut self) {
        self.alive.store(false, Ordering::Relaxed);
    }
}
struct State {
    wasi: WasiCtx,
    table: ResourceTable,
    fs: Adapter,
    runtime: Arc<Runtime>,
    pid: u32,
    command: u32,
    group: u32,
    conversations: Vec<String>,
    stdio: Io,
    preview1: Option<sausage_vfs::preview1::State>,
    http: wasmtime_wasi_http::WasiHttpCtx,
    hooks: http::Hooks,
}
impl wasmtime_wasi_http::WasiHttpView for State {
    fn http(&mut self) -> wasmtime_wasi_http::WasiHttpCtxView<'_> {
        wasmtime_wasi_http::WasiHttpCtxView {
            ctx: &mut self.http,
            table: &mut self.table,
            hooks: &mut self.hooks,
        }
    }
}
impl WasiView for State {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView {
            ctx: &mut self.wasi,
            table: &mut self.table,
        }
    }
}
pub struct Child {
    pid: u32,
    group: u32,
    code: watch::Receiver<Option<i32>>,
    input: Option<DynOutputStream>,
    output: Option<DynInputStream>,
    error: Option<DynInputStream>,
}
struct Started {
    pid: u32,
    group: u32,
    code: watch::Receiver<Option<i32>>,
    input: Option<DuplexStream>,
    output: Option<DuplexStream>,
    error: Option<DuplexStream>,
}
impl Runtime {
    pub fn open(path: &Path) -> anyhow::Result<Arc<Self>> {
        Self::new(World::open(path)?)
    }
    pub fn new(world: World) -> anyhow::Result<Arc<Self>> {
        Self::with_compiler(world, cache::shared()?)
    }
    /// Open an independent engine using an explicit shared disk cache.
    pub fn new_with_cache(world: World, directory: &Path) -> anyhow::Result<Arc<Self>> {
        Self::with_compiler(world, cache::Compiler::new(directory)?)
    }
    fn with_compiler(world: World, compiler: Arc<cache::Compiler>) -> anyhow::Result<Arc<Self>> {
        let cred = Arc::new(crate::cred::Cred::open(&world.host_dir)?);
        Self::with_hooks(world, compiler, cred, default_script(), http::whoami)
    }
    /// Share credentials and the host API router with the host's other subsystems.
    pub fn new_with_hooks(
        world: World,
        cred: Arc<crate::cred::Cred>,
        script: Arc<dyn crate::cred::script::Script>,
        router: Option<Router>,
    ) -> anyhow::Result<Arc<Self>> {
        Self::with_hooks(
            world,
            cache::shared()?,
            cred,
            script,
            router.unwrap_or(http::whoami),
        )
    }
    fn with_hooks(
        world: World,
        compiler: Arc<cache::Compiler>,
        cred: Arc<crate::cred::Cred>,
        script: Arc<dyn crate::cred::script::Script>,
        router: http::Router,
    ) -> anyhow::Result<Arc<Self>> {
        let network = Arc::new(net::NetPolicy::open(&world.host_dir)?);
        let engine = compiler.engine.clone();
        let alive = Arc::new(AtomicBool::new(true));
        let ticking = alive.clone();
        let clock = engine.clone();
        std::thread::spawn(move || {
            while ticking.load(Ordering::Relaxed) {
                std::thread::sleep(std::time::Duration::from_millis(10));
                clock.increment_epoch();
            }
        });
        let runtime = Arc::new(Self {
            world,
            engine,
            cred,
            script,
            router,
            network,
            alive,
            compiler,
            cache: Mutex::new(Default::default()),
            entries: Mutex::new(HashMap::new()),
            next: AtomicU32::new(1),
        });
        cache::precompile(runtime.compiler.clone(), runtime.world.root.clone());
        Ok(runtime)
    }
    /// Number of distinct program contents compiled by this runtime.
    pub fn compiled_programs(&self) -> usize {
        self.cache.lock().unwrap().len()
    }
    async fn resolve_program(
        &self,
        cwd: &str,
        name: &str,
        env: &[(String, String)],
    ) -> anyhow::Result<String> {
        let paths = if name.contains('/') {
            vec![sausage_vfs::resolve(cwd, name)]
        } else {
            let path = env
                .iter()
                .find(|(key, _)| key == "PATH")
                .map(|(_, value)| value.as_str())
                .unwrap_or("/usr/local/bin:/bin");
            path.split(':')
                .map(|folder| {
                    sausage_vfs::resolve(
                        cwd,
                        &std::path::Path::new(folder).join(name).to_string_lossy(),
                    )
                })
                .collect()
        };
        for path in paths {
            for candidate in [path.clone(), format!("{path}.wasm")] {
                match self.world.root.stat(&candidate, true).await {
                    Ok(stat) if !stat.directory => return self.canonical_program(candidate).await,
                    Ok(_) => {}
                    Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
                    Err(e) => return Err(e.into()),
                }
            }
        }
        anyhow::bail!("program not found: {name}")
    }
    async fn canonical_program(&self, mut path: String) -> anyhow::Result<String> {
        for _ in 0..40 {
            let mut prefix = String::new();
            let components: Vec<_> = path.trim_start_matches('/').split('/').collect();
            let mut target = None;
            for (index, component) in components.iter().enumerate() {
                prefix.push('/');
                prefix.push_str(component);
                if self.world.root.stat(&prefix, false).await?.symlink {
                    let link = self.world.root.readlink(&prefix).await?;
                    let parent = prefix.rsplit_once('/').unwrap().0;
                    let resolved = sausage_vfs::resolve(parent, &link);
                    target = Some(sausage_vfs::resolve(
                        &resolved,
                        &components[index + 1..].join("/"),
                    ));
                    break;
                }
            }
            match target {
                Some(target) => path = target,
                None => return Ok(path),
            }
        }
        anyhow::bail!("too many program symlinks")
    }
    fn kill_command(&self, command: u32) {
        for entry in self.entries.lock().unwrap().values() {
            if entry.command == command {
                entry.cancelled.store(true, Ordering::Relaxed);
                entry.abort.abort();
            }
        }
    }
    fn kill_group(&self, group: u32) {
        for entry in self.entries.lock().unwrap().values() {
            if entry.group == group {
                entry.abort.abort();
            }
        }
    }
    fn kill_pid(&self, pid: u32) -> anyhow::Result<()> {
        let group = self
            .entries
            .lock()
            .unwrap()
            .get(&pid)
            .map(|e| e.group)
            .ok_or_else(|| anyhow::anyhow!("no live process {pid}"))?;
        self.kill_group(group);
        Ok(())
    }
    async fn start(
        self: &Arc<Self>,
        mut request: Request,
        parent: u32,
        command: u32,
        inherited_group: u32,
        conversations: Vec<String>,
        inherited_io: Option<Io>,
    ) -> anyhow::Result<Started> {
        anyhow::ensure!(!request.argv.is_empty(), "empty argv");
        let cancelled = if parent == 0 {
            Arc::new(AtomicBool::new(false))
        } else {
            self.entries
                .lock()
                .unwrap()
                .get(&parent)
                .map(|entry| entry.cancelled.clone())
                .ok_or_else(|| anyhow::anyhow!("parent exited"))?
        };
        if !request.env.iter().any(|(k, _)| k == "PATH") {
            request
                .env
                .push(("PATH".into(), "/usr/local/bin:/bin".into()));
        }
        request.cwd = sausage_vfs::resolve("/", &request.cwd);
        request.env.retain(|(name, _)| name != "PWD");
        request.env.push(("PWD".into(), request.cwd.clone()));
        let executable = self
            .resolve_program(&request.cwd, &request.program, &request.env)
            .await?;
        let file = self
            .world
            .root
            .open(
                &executable,
                sausage_vfs::Open {
                    follow: true,
                    ..Default::default()
                },
            )
            .await?;
        let bytes = file
            .read(0, usize::try_from(file.stat().await?.size)?)
            .await?;
        let pre = self.compiler.get(bytes.clone()).await?;
        self.cache.lock().unwrap().insert(bytes);
        let pid = self.next.fetch_add(1, Ordering::Relaxed);
        let command = if command == 0 { pid } else { command };
        let group = match request.group {
            Some(0) => pid,
            Some(group) => group,
            None if inherited_group != 0 => inherited_group,
            None => pid,
        };
        request.env.retain(|(name, _)| name != "SAUSAGE_PID");
        request.env.push(("SAUSAGE_PID".into(), pid.to_string()));
        request.env.push(("PID".into(), pid.to_string()));
        request.env.push(("PPID".into(), parent.to_string()));
        let (input, child_input) = stdio::input(request.stdin, inherited_io.as_ref());
        let (output, child_output) =
            stdio::output(request.stdout, inherited_io.as_ref().map(|io| &io.output));
        let (error, child_error) =
            stdio::output(request.stderr, inherited_io.as_ref().map(|io| &io.error));
        let stdio = Io {
            input: child_input,
            output: child_output,
            error: child_error,
        };
        let preview1 = if matches!(pre, Compiled::Module(_)) {
            Some(sausage_vfs::preview1::State::new(
                self.world.root.clone(),
                &request.cwd,
                request.argv.clone(),
                request.env.clone(),
                stdio.input(),
                Box::new(stdio.output.clone()),
                Box::new(stdio.error.clone()),
            ))
        } else {
            None
        };
        let mut mounts = Mounts::default();
        mounts.mount("/", self.world.root.clone(), false);
        let mut wasi = WasiCtxBuilder::new();
        wasi.args(&request.argv)
            .envs(&request.env)
            .initial_cwd(&request.cwd)
            .stdin(stdio.input.clone())
            .stdout(AsyncStdoutStream::new(8192, stdio.output.clone()))
            .stderr(AsyncStdoutStream::new(8192, stdio.error.clone()))
            .allow_tcp(true)
            .allow_udp(false)
            .allow_ip_name_lookup(true);
        let policy = self.network.clone();
        wasi.socket_addr_check(move |address, _| {
            let allowed = policy.allows(address);
            Box::pin(async move { allowed })
        });
        let mut store = Store::new(
            &self.engine,
            State {
                wasi: wasi.build(),
                table: ResourceTable::new(),
                fs: Adapter::new(mounts),
                runtime: self.clone(),
                pid,
                command,
                group,
                conversations: conversations.clone(),
                stdio,
                preview1,
                http: wasmtime_wasi_http::WasiHttpCtx::new(),
                hooks: http::Hooks {
                    caller: crate::cred::Caller {
                        pid,
                        program: executable,
                        conversations: conversations.clone(),
                    },
                    cred: self.cred.clone(),
                    script: self.script.clone(),
                    policy: self.network.clone(),
                    router: self.router,
                },
            },
        );
        store.epoch_deadline_async_yield_and_update(1);
        self.world.processes.lock().unwrap().insert(
            pid,
            ProcessInfo {
                parent,
                argv: request.argv,
                cwd: request.cwd,
                state: "running".into(),
                conversation: conversations.first().cloned(),
            },
        );
        let (tx, code) = watch::channel(None);
        // Capture cleanup before the first poll, so an early abort still removes the entry.
        let cleanup = Cleanup {
            runtime: self.clone(),
            pid,
            tx,
        };
        // The start gate makes the entry visible before a guest can spawn children or exit.
        let (gate, ready) = tokio::sync::oneshot::channel();
        let task = tokio::spawn(async move {
            let _cleanup = cleanup;
            if ready.await.is_err() {
                return;
            }
            let result = async {
                let result = match pre {
                    Compiled::Component(pre) => {
                        let command = pre.instantiate_async(&mut store).await?;
                        command.wasi_cli_run().call_run(&mut store).await
                    }
                    Compiled::Module(pre) => {
                        let module = pre.instantiate_async(&mut store).await?;
                        module
                            .get_typed_func::<(), ()>(&mut store, "_start")?
                            .call_async(&mut store, ())
                            .await
                            .map(|()| Ok(()))
                    }
                };
                match result {
                    Ok(Ok(())) => Ok(0),
                    Ok(Err(())) => Ok(1),
                    Err(e) => match e.downcast_ref::<wasmtime_wasi::I32Exit>() {
                        Some(exit) => Ok(exit.0),
                        None => Err(e),
                    },
                }
            }
            .await;
            _cleanup.tx.send_replace(Some(result.unwrap_or(125)));
        });
        {
            let mut entries = self.entries.lock().unwrap();
            entries.insert(
                pid,
                Entry {
                    command,
                    group,
                    cancelled: cancelled.clone(),
                    abort: task.abort_handle(),
                },
            );
            // A spawn already in progress must observe a command-wide cancellation.
            if cancelled.load(Ordering::Relaxed) {
                task.abort();
            } else {
                let _ = gate.send(());
            }
        }
        Ok(Started {
            pid,
            group,
            code,
            input,
            output,
            error,
        })
    }
}
struct Cleanup {
    runtime: Arc<Runtime>,
    pid: u32,
    tx: watch::Sender<Option<i32>>,
}
impl Drop for Cleanup {
    fn drop(&mut self) {
        if self.tx.borrow().is_none() {
            self.tx.send_replace(Some(137));
        }
        self.runtime.entries.lock().unwrap().remove(&self.pid);
        self.runtime
            .world
            .processes
            .lock()
            .unwrap()
            .remove(&self.pid);
    }
}
impl Host for State {
    async fn spawn(
        &mut self,
        request: Request,
    ) -> wasmtime::Result<Result<Resource<Child>, String>> {
        let child = match self
            .runtime
            .start(
                request,
                self.pid,
                self.command,
                self.group,
                self.conversations.clone(),
                Some(self.stdio.clone()),
            )
            .await
        {
            Ok(child) => child,
            Err(e) => return Ok(Err(e.to_string())),
        };
        Ok(Ok(self.table.push(Child {
            pid: child.pid,
            group: child.group,
            code: child.code,
            input: child
                .input
                .map(|s| Box::new(AsyncWriteStream::new(8192, s)) as DynOutputStream),
            output: child
                .output
                .map(|s| Box::new(AsyncReadStream::new(s)) as DynInputStream),
            error: child
                .error
                .map(|s| Box::new(AsyncReadStream::new(s)) as DynInputStream),
        })?))
    }
    async fn kill(&mut self, pid: u32) -> wasmtime::Result<Result<(), String>> {
        Ok(self.runtime.kill_pid(pid).map_err(|e| e.to_string()))
    }
}
#[async_trait::async_trait]
impl wasmtime_wasi::p2::Pollable for Child {
    async fn ready(&mut self) {
        while self.code.borrow_and_update().is_none() {
            if self.code.changed().await.is_err() {
                break;
            }
        }
    }
}
impl HostChild for State {
    async fn subscribe(
        &mut self,
        child: Resource<Child>,
    ) -> wasmtime::Result<Resource<wasmtime_wasi::p2::DynPollable>> {
        wasmtime_wasi::p2::subscribe(&mut self.table, child)
    }
    async fn pid(&mut self, child: Resource<Child>) -> wasmtime::Result<u32> {
        Ok(self.table.get(&child)?.pid)
    }
    async fn group(&mut self, child: Resource<Child>) -> wasmtime::Result<u32> {
        Ok(self.table.get(&child)?.group)
    }
    async fn stdin(
        &mut self,
        child: Resource<Child>,
    ) -> wasmtime::Result<Resource<DynOutputStream>> {
        let stream = self
            .table
            .get_mut(&child)?
            .input
            .take()
            .ok_or_else(|| wasmtime::format_err!("stdin already taken"))?;
        Ok(self.table.push(stream)?)
    }
    async fn stdout(
        &mut self,
        child: Resource<Child>,
    ) -> wasmtime::Result<Resource<DynInputStream>> {
        let stream = self
            .table
            .get_mut(&child)?
            .output
            .take()
            .ok_or_else(|| wasmtime::format_err!("stdout already taken"))?;
        Ok(self.table.push(stream)?)
    }
    async fn stderr(
        &mut self,
        child: Resource<Child>,
    ) -> wasmtime::Result<Resource<DynInputStream>> {
        let stream = self
            .table
            .get_mut(&child)?
            .error
            .take()
            .ok_or_else(|| wasmtime::format_err!("stderr already taken"))?;
        Ok(self.table.push(stream)?)
    }
    async fn status(&mut self, child: Resource<Child>) -> wasmtime::Result<Option<i32>> {
        Ok(*self.table.get(&child)?.code.borrow())
    }
    async fn wait(&mut self, child: Resource<Child>) -> wasmtime::Result<i32> {
        let mut code = self.table.get(&child)?.code.clone();
        loop {
            if let Some(code) = *code.borrow_and_update() {
                return Ok(code);
            }
            code.changed().await?;
        }
    }
    async fn kill(&mut self, child: Resource<Child>) -> wasmtime::Result<()> {
        self.runtime
            .kill_pid(self.table.get(&child)?.pid)
            .map_err(|e| wasmtime::format_err!("{e}"))
    }
    async fn drop(&mut self, child: Resource<Child>) -> wasmtime::Result<()> {
        self.table.delete(child)?;
        Ok(())
    }
}
struct GroupGuard {
    runtime: Arc<Runtime>,
    group: u32,
}
impl Drop for GroupGuard {
    fn drop(&mut self) {
        self.runtime.kill_command(self.group);
    }
}
/// An Arc runtime implements Exec so the harness can retain one warm process table.
#[async_trait::async_trait]
impl Exec for Arc<Runtime> {
    async fn exec(
        &self,
        command: Command,
        options: Options,
        output: &mut (dyn for<'chunk> FnMut(Stream, &'chunk [u8]) + Send),
    ) -> anyhow::Result<i32> {
        let deadline = options
            .timeout
            .map(|duration| tokio::time::Instant::now() + duration);
        let argv = match command {
            Command::Shell(script) => vec![
                "/bin/sh".into(),
                "--noprofile".into(),
                "--norc".into(),
                "-c".into(),
                script,
            ],
            Command::Argv(argv) => argv,
        };
        let env = options.env;
        let start = self.start(
            Request {
                program: argv.first().cloned().unwrap_or_default(),
                argv,
                env,
                cwd: options.cwd,
                group: None,
                stdin: Stdio::Pipe,
                stdout: Stdio::Pipe,
                stderr: Stdio::Pipe,
            },
            0,
            0,
            0,
            options.conversations,
            None,
        );
        let mut child = if let Some(deadline) = deadline {
            tokio::time::timeout_at(deadline, start)
                .await
                .map_err(|_| anyhow::anyhow!("command timed out"))??
        } else {
            start.await?
        };
        let _guard = GroupGuard {
            runtime: self.clone(),
            group: child.pid,
        };
        drop(child.input);
        let mut stdout = child.output.take().unwrap();
        let mut stderr = child.error.take().unwrap();
        let run = async {
            let (mut out_open, mut err_open) = (true, true);
            let (mut out, mut err) = ([0; 8192], [0; 8192]);
            while out_open || err_open {
                tokio::select! {
                    n = stdout.read(&mut out), if out_open => { let n = n?; out_open = n != 0; if n != 0 { output(Stream::Stdout, &out[..n]); } }
                    n = stderr.read(&mut err), if err_open => { let n = n?; err_open = n != 0; if n != 0 { output(Stream::Stderr, &err[..n]); } }
                }
            }
            loop {
                if let Some(code) = *child.code.borrow_and_update() {
                    return Ok(code);
                }
                child.code.changed().await?;
            }
        };
        if let Some(deadline) = deadline {
            tokio::time::timeout_at(deadline, run)
                .await
                .map_err(|_| anyhow::anyhow!("command timed out"))?
        } else {
            run.await
        }
    }
}
/// Construct the harness's shared executor over the same open world.
pub fn exec(world: World) -> anyhow::Result<Arc<dyn Exec>> {
    Ok(Arc::new(Runtime::new(world)?))
}

/// Run the world's shell. Noninteractive shell scripts stream directly to host stdio.
pub async fn sh(world: &Path, command: Option<String>) -> anyhow::Result<i32> {
    let runtime = Runtime::open(world)?;
    let command = command.ok_or_else(|| anyhow::anyhow!("interactive shell is not implemented"))?;
    runtime
        .exec(
            Command::Shell(command),
            Options {
                cwd: "/".into(),
                ..Options::default()
            },
            &mut |stream, bytes| {
                use std::io::Write;
                match stream {
                    Stream::Stdout => {
                        let _ = std::io::stdout().write_all(bytes);
                    }
                    Stream::Stderr => {
                        let _ = std::io::stderr().write_all(bytes);
                    }
                }
            },
        )
        .await
}