//! Wasmtime wrapper for executing sandboxed Wasm tools.

use std::path::Path;

use thiserror::Error;
use wasmbox_shared::Capability;
use wasmtime::Engine;
use wasmtime_wasi::{DirPerms, FilePerms, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
use wasmtime_wasi_http::body::HyperOutgoingBody;
use wasmtime_wasi_http::types::{
    default_send_request, HostFutureIncomingResponse, OutgoingRequestConfig,
};
use wasmtime_wasi_http::{HttpResult, WasiHttpCtx, WasiHttpView};

#[derive(Debug, Error)]
pub enum RuntimeError {
    #[error("wasmtime error: {0}")]
    Wasmtime(#[from] wasmtime::Error),

    #[error("wasm file not found: {0}")]
    FileNotFound(String),

    #[error("io error: {0}")]
    Io(#[from] std::io::Error),

    /// Tool called proc_exit with a non-zero code. Not a crash — forward to host.
    #[error("tool exited with non-zero status")]
    ToolExited,
}

/// Configuration for running a Wasm tool.
pub struct RunConfig {
    pub wasm_path: std::path::PathBuf,
    pub capabilities: Vec<Capability>,
    pub fuel_limit: Option<u64>,
    pub args: Vec<String>,
}

struct HostState {
    ctx: WasiCtx,
    http_ctx: WasiHttpCtx,
    table: wasmtime::component::ResourceTable,
    allowed_hosts: Vec<String>,
}

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

impl WasiHttpView for HostState {
    fn ctx(&mut self) -> &mut WasiHttpCtx {
        &mut self.http_ctx
    }

    fn table(&mut self) -> &mut wasmtime::component::ResourceTable {
        &mut self.table
    }

    fn send_request(
        &mut self,
        request: hyper::Request<HyperOutgoingBody>,
        config: OutgoingRequestConfig,
    ) -> HttpResult<HostFutureIncomingResponse> {
        // Enforce host allowlist from Network capability
        if self.allowed_hosts.is_empty() {
            return Err(wasmtime_wasi_http::HttpError::trap(
                wasmtime::Error::msg("network access denied: no Network capability granted"),
            ));
        }

        if let Some(authority) = request.uri().authority() {
            let host = authority.host();
            let allowed = self.allowed_hosts.iter().any(|h| {
                // Match "host:port" or just "host"
                h == authority.as_str() || h == host
            });
            if !allowed {
                return Err(wasmtime_wasi_http::HttpError::trap(wasmtime::Error::msg(
                    format!(
                        "network access denied: host '{}' not in allowed list",
                        authority
                    ),
                )));
            }
        }

        Ok(default_send_request(request, config))
    }
}

/// Execute a Wasm component tool with the given configuration.
pub fn run_tool(config: &RunConfig) -> Result<(), RuntimeError> {
    if !config.wasm_path.exists() {
        return Err(RuntimeError::FileNotFound(
            config.wasm_path.display().to_string(),
        ));
    }

    let mut engine_config = wasmtime::Config::new();
    engine_config.wasm_component_model(true);

    if config.fuel_limit.is_some() {
        engine_config.consume_fuel(true);
    }

    let engine = Engine::new(&engine_config)?;
    let wasm_bytes = std::fs::read(&config.wasm_path)?;
    let component = wasmtime::component::Component::new(&engine, &wasm_bytes)?;

    let (wasi_ctx, allowed_hosts) = build_wasi_context(&config.capabilities, &config.args)?;

    let host_state = HostState {
        ctx: wasi_ctx,
        http_ctx: WasiHttpCtx::new(),
        table: wasmtime::component::ResourceTable::new(),
        allowed_hosts,
    };

    let mut store = wasmtime::Store::new(&engine, host_state);

    if let Some(fuel) = config.fuel_limit {
        store.set_fuel(fuel)?;
    }

    let mut linker = wasmtime::component::Linker::new(&engine);

    // Full WASI CLI interfaces (stdin, stdout, filesystem, env, etc.)
    wasmtime_wasi::p2::add_to_linker_sync(&mut linker)?;
    // Add wasi:http on top for tools that need outbound HTTP
    wasmtime_wasi_http::add_only_http_to_linker_sync(&mut linker)?;

    let command = wasmtime_wasi::p2::bindings::sync::Command::instantiate(
        &mut store,
        &component,
        &linker,
    )?;

    let result = command.wasi_cli_run().call_run(&mut store);

    match result {
        Ok(Ok(())) => Ok(()),
        Ok(Err(())) => Err(RuntimeError::ToolExited),
        Err(e) => {
            // proc_exit(non_zero) surfaces as an I32Exit trap, not as Err(()).
            // Treat it the same as a clean non-zero return.
            if let Some(exit) = e.downcast_ref::<wasmtime_wasi::I32Exit>() {
                if exit.0 == 0 { Ok(()) } else { Err(RuntimeError::ToolExited) }
            } else {
                Err(RuntimeError::Wasmtime(e))
            }
        }
    }
}

fn build_wasi_context(
    capabilities: &[Capability],
    args: &[String],
) -> Result<(WasiCtx, Vec<String>), RuntimeError> {
    let mut builder = WasiCtxBuilder::new();
    let mut allowed_hosts: Vec<String> = Vec::new();

    // WASI expects argv[0] to be the program name, followed by actual arguments.
    // We always set args so the tool gets a proper argv even if no extra args given.
    let mut full_args: Vec<&str> = vec!["tool"];
    full_args.extend(args.iter().map(|s| s.as_str()));
    builder.args(&full_args);

    for cap in capabilities {
        match cap {
            Capability::Stdin => {
                builder.inherit_stdin();
            }
            Capability::Stdout => {
                builder.inherit_stdout();
                builder.inherit_stderr();
            }
            Capability::Env(vars) => {
                for var in vars {
                    if let Ok(val) = std::env::var(var) {
                        builder.env(var, &val);
                    }
                }
            }
            Capability::Filesystem(grant) => {
                configure_filesystem(&mut builder, &grant.path, grant.read, grant.write)?;
            }
            Capability::Network(hosts) => {
                allowed_hosts.extend(hosts.clone());
            }
            Capability::Clipboard => {
                // Clipboard is a host-specific capability, not part of WASI.
                // Would need a custom host function.
            }
        }
    }

    Ok((builder.build(), allowed_hosts))
}

fn configure_filesystem(
    builder: &mut WasiCtxBuilder,
    path: &str,
    read: bool,
    write: bool,
) -> Result<(), RuntimeError> {
    let expanded = expand_path(path);
    let host_path = Path::new(&expanded);

    if !host_path.exists() {
        return Err(RuntimeError::Io(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!("preopened path not found: {path}"),
        )));
    }

    let dir_perms = match (read, write) {
        (true, true) => DirPerms::all(),
        (true, false) => DirPerms::READ,
        (false, true) => DirPerms::MUTATE,
        (false, false) => DirPerms::empty(),
    };

    let file_perms = match (read, write) {
        (true, true) => FilePerms::all(),
        (true, false) => FilePerms::READ,
        (false, true) => FilePerms::WRITE,
        (false, false) => FilePerms::empty(),
    };

    builder.preopened_dir(host_path, path, dir_perms, file_perms)?;

    Ok(())
}

fn expand_path(path: &str) -> String {
    if let Some(rest) = path.strip_prefix("~/")
        && let Some(home) = dirs::home_dir()
    {
        return home.join(rest).to_string_lossy().to_string();
    }
    path.to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn missing_wasm_file_returns_error() {
        let config = RunConfig {
            wasm_path: std::path::PathBuf::from("/nonexistent/tool.wasm"),
            capabilities: vec![],
            fuel_limit: None,
            args: vec![],
        };
        assert!(run_tool(&config).is_err());
    }

    #[test]
    fn expand_tilde_path() {
        let expanded = expand_path("~/Documents");
        assert!(!expanded.starts_with('~'));
        assert!(expanded.contains("Documents"));
    }

    #[test]
    fn expand_absolute_path_unchanged() {
        let expanded = expand_path("/tmp/data");
        assert_eq!(expanded, "/tmp/data");
    }

    #[test]
    fn allowed_hosts_from_network_capability() {
        let caps = vec![
            Capability::Stdin,
            Capability::Network(vec![
                "api.example.com".to_string(),
                "cdn.example.com".to_string(),
            ]),
        ];
        let (_ctx, hosts) = build_wasi_context(&caps, &[]).unwrap();
        assert_eq!(hosts, vec!["api.example.com", "cdn.example.com"]);
    }

    #[test]
    fn no_network_capability_means_empty_hosts() {
        let caps = vec![Capability::Stdin, Capability::Stdout];
        let (_ctx, hosts) = build_wasi_context(&caps, &[]).unwrap();
        assert!(hosts.is_empty());
    }
}