//! 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;
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::p2::bindings::Command;
use wasmtime_wasi::p2::bindings::filesystem::{preopens, types};
use wasmtime_wasi::{FsPerms, I32Exit, WasiCtx, WasiCtxView, WasiView};
use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpCtxView, WasiHttpHooks, WasiHttpView};

use crate::fs::{ActorFs, ActorFsData, FsExt, Mount};

pub struct WasiPart {
    ctx: WasiCtx,
    table: ResourceTable,
}

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

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 {
    wasi: WasiPart,
    fsx: FsExt,
    http: WasiHttpCtx,
    hooks: HttpHooks,
}

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

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

fn actor_fs(s: &mut ActorState) -> ActorFs<'_> {
    ActorFs {
        inner: wasmtime_wasi::filesystem::WasiFilesystemView::filesystem(&mut s.wasi),
        ext: &mut s.fsx,
    }
}

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 Exit {
    pub fn is_clean(&self) -> bool {
        matches!(self, Exit::Clean)
    }
}

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>,
    pub allow_http: bool,
    pub stderr_log: PathBuf,
    pub env: Vec<(String, String)>,
}

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

impl Instance {
    /// Write one line to the guest's stdin.
    pub async fn send_line(&mut self, line: &str) -> std::io::Result<()> {
        self.stdin.write_all(line.as_bytes()).await?;
        self.stdin.write_all(b"\n").await?;
        self.stdin.flush().await
    }

    /// 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()
    }

    /// Send `shutdown`, close stdin, and wait. Kill the guest after `deadline`.
    pub async fn stop(mut self, shutdown_line: &str, deadline: Duration) -> Exit {
        let _ = self.send_line(shutdown_line).await;
        let _ = self.stdin.shutdown().await;
        drop(self.stdin);
        match tokio::time::timeout(deadline, &mut self.task).await {
            Ok(Ok(exit)) => exit,
            Ok(Err(e)) => Exit::Trap(format!("host task failed: {e}")),
            Err(_) => {
                self.task.abort();
                let _ = (&mut self.task).await;
                Exit::Trap("killed after shutdown deadline".into())
            }
        }
    }
}

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 the filesystem with the commit-on-close layer.
        linker.allow_shadowing(true);
        types::add_to_linker::<ActorState, ActorFsData>(&mut linker, actor_fs)?;
        preopens::add_to_linker::<ActorState, ActorFsData>(&mut linker, actor_fs)?;
        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();
        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 {
            wasi: WasiPart { ctx: b.build(), table: ResourceTable::new() },
            fsx: FsExt::new(cfg.mounts),
            http: WasiHttpCtx::new(),
            hooks: HttpHooks { allow: cfg.allow_http },
        };

        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 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;
            let exit = 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 => Exit::Trap(format!("{e:#}")),
                },
            };
            // Process exit closes every file, so a clean exit commits.
            if exit.is_clean() {
                store.data_mut().fsx.commit_all();
            }
            exit
        });

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