//! Wasmtime engine and actor instances.
//!
//! An instance is one run of a `wasi:cli/command` component. The host writes
//! NDJSON lines to its stdin and reads NDJSON lines from its stdout.

use std::path::{Path, PathBuf};
use std::sync::{Arc, Weak};
use std::time::Duration;

use anyhow::{Context, Result};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use wasmtime::component::{Component, Linker, ResourceTable};
use wasmtime::{Config, Engine, Store};
use wasmtime_wasi::cli::{AsyncStdinStream, AsyncStdoutStream};
use wasmtime_wasi::filesystem::WasiFilesystemCtxView;
use wasmtime_wasi::p2::bindings::Command;
use wasmtime_wasi::p2::bindings::filesystem;
use wasmtime_wasi::{FsPerms, I32Exit, WasiCtx, WasiCtxView, WasiView};
use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpCtxView, WasiHttpHooks, WasiHttpView};

use crate::broker::Broker;
use crate::vfs::{Vfs, VfsState, VfsView};

/// A host folder that the guest sees at `guest`.
#[derive(Clone, Debug)]
pub struct Mount {
    pub guest: String,
    pub host: PathBuf,
    pub writable: bool,
}

pub struct HttpHooks {
    allow: bool,
}

impl WasiHttpHooks for HttpHooks {
    fn send_request(
        &mut self,
        request: http::Request<wasmtime_wasi_http::WasiBody>,
        options: Option<wasmtime_wasi_http::RequestOptions>,
        fut: Box<dyn Future<Output = Result<(), wasmtime_wasi_http::Error>> + Send>,
    ) -> Box<
        dyn Future<
                Output = wasmtime_wasi_http::Result<(
                    http::Response<wasmtime_wasi_http::WasiBody>,
                    Box<dyn Future<Output = Result<(), wasmtime_wasi_http::Error>> + Send>,
                )>,
            > + Send,
    > {
        _ = fut;
        let allow = self.allow;
        Box::new(async move {
            use http_body_util::BodyExt;
            if !allow {
                return Err(wasmtime_wasi_http::Error::HttpRequestDenied);
            }
            let (res, io) = wasmtime_wasi_http::default_send_request(request, options).await?;
            Ok((
                res.map(BodyExt::boxed_unsync),
                Box::new(io) as Box<dyn Future<Output = _> + Send>,
            ))
        })
    }
}

pub struct ActorState {
    ctx: WasiCtx,
    table: ResourceTable,
    http: WasiHttpCtx,
    hooks: HttpHooks,
    vfs: VfsState,
}

fn vfs_view(s: &mut ActorState) -> VfsView<'_> {
    VfsView { inner: WasiFilesystemCtxView { ctx: s.ctx.filesystem(), table: &mut s.table }, state: &mut s.vfs }
}

impl WasiView for ActorState {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView { ctx: &mut self.ctx, table: &mut self.table }
    }
}

impl WasiHttpView for ActorState {
    fn http(&mut self) -> WasiHttpCtxView<'_> {
        WasiHttpCtxView { ctx: &mut self.http, table: &mut self.table, hooks: &mut self.hooks }
    }
}

pub struct Runtime {
    engine: Engine,
    linker: Linker<ActorState>,
}

/// How an instance ended.
#[derive(Debug, Clone)]
pub enum Exit {
    /// `main` returned or the guest called `exit(0)`.
    Clean,
    /// The guest exited with a nonzero code.
    Code(i32),
    /// The guest trapped (panic, invalid memory access, ...).
    Trap(String),
}

impl std::fmt::Display for Exit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Exit::Clean => write!(f, "exit 0"),
            Exit::Code(c) => write!(f, "exit {c}"),
            Exit::Trap(t) => write!(f, "trap: {t}"),
        }
    }
}

pub struct InstanceConfig {
    pub mounts: Vec<Mount>,
    /// Outgoing HTTP, and TCP, UDP, and DNS through `wasi:sockets`.
    pub network: bool,
    pub stderr_log: PathBuf,
    pub env: Vec<(String, String)>,
    /// The actor that runs in the instance.
    pub actor: String,
    /// Asks the actors that serve `/actor/<id>/fs/` (see `vfs.rs`).
    pub broker: Weak<Broker>,
}

/// A running guest. Dropping it kills the guest.
pub struct Instance {
    stdin: Option<DuplexStream>,
    /// Lines that the guest writes to stdout. Closed when the guest ends.
    pub stdout: mpsc::Receiver<String>,
    task: JoinHandle<Exit>,
}

impl Drop for Instance {
    fn drop(&mut self) {
        self.task.abort();
    }
}

impl Instance {
    /// Write one line to the guest's stdin, in one write. Meanwhile, pass the
    /// lines that the guest writes to `route`: a guest that blocks on a full
    /// stdout would never read its stdin.
    pub async fn send_line(&mut self, line: &str, mut route: impl FnMut(String)) -> std::io::Result<()> {
        let stdin = self.stdin.as_mut().ok_or(std::io::ErrorKind::BrokenPipe)?;
        let bytes = format!("{line}\n");
        let write = async {
            stdin.write_all(bytes.as_bytes()).await?;
            stdin.flush().await
        };
        tokio::pin!(write);
        let mut open = true;
        loop {
            tokio::select! {
                done = &mut write => return done,
                out = self.stdout.recv(), if open => match out {
                    Some(out) => route(out),
                    None => open = false,
                },
            }
        }
    }

    /// Wait until the guest ends.
    pub async fn wait(&mut self) -> Exit {
        match (&mut self.task).await {
            Ok(exit) => exit,
            Err(e) if e.is_cancelled() => Exit::Trap("killed".into()),
            Err(e) => Exit::Trap(format!("host task failed: {e}")),
        }
    }

    pub fn is_finished(&self) -> bool {
        self.task.is_finished()
    }

    /// Kill the guest. Its stdout closes when it has ended.
    pub fn kill(&self) {
        self.task.abort();
    }

    /// Send `shutdown_line`, close stdin, and wait until the guest ends.
    /// Pass every line that the guest writes meanwhile to `route`. Kill the
    /// guest when `deadline` has passed.
    pub async fn stop(mut self, shutdown_line: &str, deadline: Duration, mut route: impl FnMut(String)) -> Exit {
        let stdin = self.stdin.take();
        let stdout = &mut self.stdout;
        let finished = tokio::time::timeout(deadline, async {
            let close = async {
                if let Some(mut stdin) = stdin {
                    let _ = stdin.write_all(format!("{shutdown_line}\n").as_bytes()).await;
                    let _ = stdin.shutdown().await;
                }
            };
            let drain = async {
                while let Some(line) = stdout.recv().await {
                    route(line);
                }
            };
            tokio::join!(close, drain);
        })
        .await;
        match finished {
            Ok(()) => self.wait().await,
            Err(_) => Exit::Trap("killed after shutdown deadline".into()),
        }
    }
}

fn log_note(path: &Path, text: &str) {
    use std::io::Write;
    if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(path) {
        let _ = writeln!(f, "[sausage] {text}");
    }
}

/// The most useful line of a trap: the `wasm trap: ...` reason.
fn short_trap(full: &str) -> String {
    let line = full
        .lines()
        .find_map(|l| l.find("wasm trap:").map(|i| l[i..].to_string()))
        .unwrap_or_else(|| full.lines().next().unwrap_or("trap").to_string());
    let mut line = line.trim().to_string();
    line.truncate(300);
    line
}

impl Runtime {
    pub fn new() -> Result<Arc<Self>> {
        let mut config = Config::new();
        config.epoch_interruption(true);
        let engine = Engine::new(&config)?;

        let mut linker: Linker<ActorState> = Linker::new(&engine);
        wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
        wasmtime_wasi_http::p2::add_only_http_to_linker_async(&mut linker)?;
        // Replace `wasi:filesystem` with the wrapper that serves `/actor/<id>/fs/`.
        linker.allow_shadowing(true);
        filesystem::types::add_to_linker::<ActorState, Vfs>(&mut linker, vfs_view)?;
        filesystem::preopens::add_to_linker::<ActorState, Vfs>(&mut linker, vfs_view)?;
        linker.allow_shadowing(false);

        // Time slices: every 10 ms each running guest yields to the scheduler.
        let ticker = engine.clone();
        std::thread::spawn(move || {
            loop {
                std::thread::sleep(Duration::from_millis(10));
                ticker.increment_epoch();
            }
        });

        Ok(Arc::new(Self { engine, linker }))
    }

    pub fn load(&self, wasm: &Path) -> Result<Component> {
        Component::from_file(&self.engine, wasm)
            .map_err(anyhow::Error::from)
            .with_context(|| format!("load {}", wasm.display()))
    }

    pub fn start(self: &Arc<Self>, component: &Component, cfg: InstanceConfig) -> Result<Instance> {
        let (host_in, guest_in) = tokio::io::duplex(64 * 1024);
        let (guest_out, host_out) = tokio::io::duplex(64 * 1024);

        let stderr = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&cfg.stderr_log)
            .with_context(|| format!("open {}", cfg.stderr_log.display()))?;

        let mut b = WasiCtx::builder();
        if cfg.network {
            b.inherit_network().allow_tcp(true).allow_udp(true).allow_ip_name_lookup(true);
        }
        b.stdin(AsyncStdinStream::new(guest_in));
        b.stdout(AsyncStdoutStream::new(8192, guest_out));
        b.stderr(AsyncStdoutStream::new(8192, tokio::fs::File::from_std(stderr)));
        for (k, v) in &cfg.env {
            b.env(k, v);
        }
        for m in &cfg.mounts {
            let perms = if m.writable { FsPerms::ReadWrite } else { FsPerms::ReadOnly };
            b.preopened_dir(&m.host, &m.guest, perms)
                .map_err(anyhow::Error::from)
                .with_context(|| format!("mount {} at {}", m.host.display(), m.guest))?;
        }
        let state = ActorState {
            ctx: b.build(),
            table: ResourceTable::new(),
            http: WasiHttpCtx::new(),
            hooks: HttpHooks { allow: cfg.network },
            vfs: VfsState::new(cfg.broker.clone(), cfg.actor.clone()),
        };

        let (line_tx, line_rx) = mpsc::channel(1024);
        tokio::spawn(async move {
            let mut lines = BufReader::new(host_out).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                if line_tx.send(line).await.is_err() {
                    break;
                }
            }
        });

        let rt = self.clone();
        let component = component.clone();
        let stderr_log = cfg.stderr_log.clone();
        let task = tokio::spawn(async move {
            let mut store = Store::new(&rt.engine, state);
            store.epoch_deadline_async_yield_and_update(1);
            let result = async {
                let cmd = Command::instantiate_async(&mut store, &component, &rt.linker).await?;
                cmd.wasi_cli_run().call_run(&mut store).await
            }
            .await;
            match result {
                Ok(Ok(())) => Exit::Clean,
                Ok(Err(())) => Exit::Code(1),
                Err(e) => match e.downcast_ref::<I32Exit>() {
                    Some(I32Exit(0)) => Exit::Clean,
                    Some(I32Exit(c)) => Exit::Code(*c),
                    None => {
                        // Keep the full backtrace in the log, and a short reason in the status.
                        let full = format!("{e:?}");
                        log_note(&stderr_log, &format!("trap details:\n{full}"));
                        Exit::Trap(short_trap(&full))
                    }
                },
            }
        });

        Ok(Instance { stdin: Some(host_in), stdout: line_rx, task })
    }
}