//! Firecracker microVM sandbox runner
//!
//! Executes code inside a Firecracker microVM. Firecracker is fundamentally
//! different from container-based sandboxes: it boots a real Linux kernel
//! against an operator-supplied root filesystem image and runs the workload
//! inside that VM. Because of that, this runner cannot be "fully self
//! contained" the way `DockerRunner` is — the operator must provide:
//!
//! - A boot kernel image (vmlinux ELF, e.g. produced from
//!   `https://github.com/firecracker-microvm/firecracker/blob/main/docs/rootfs-and-kernel-setup.md`).
//! - A root filesystem image (ext4) that contains an init system capable of
//!   reading the agent code from the supplied vsock channel and executing it.
//!
//! The runner generates a Firecracker VM JSON config, writes the agent code
//! to a per-execution sidecar file mounted into the rootfs via a vsock-bound
//! drop directory, then execs `firecracker --no-api --config-file <path>`.
//! Stdout/stderr come back over the VM's serial console.
//!
//! # Failure modes
//!
//! If the operator hasn't configured `[sandbox.firecracker]` in
//! `symbiont.toml` (kernel + rootfs paths), `FirecrackerRunner::new` returns
//! a clear error pointing at the missing configuration. There is no silent
//! fallback to Docker — the user explicitly asked for VM-level isolation.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::process::Stdio;
use tokio::io::AsyncReadExt;
use tokio::process::Command;
use tokio::time::{timeout, Duration};

use super::{ExecutionResult, SandboxRunner};

/// Default per-execution memory budget (MiB).
const DEFAULT_MEM_MIB: u32 = 512;
/// Default vCPU count.
const DEFAULT_VCPUS: u8 = 1;
/// Default maximum execution time before the VM is force-killed.
const DEFAULT_MAX_EXECUTION: Duration = Duration::from_secs(300);
/// Default maximum captured output bytes per stream.
const DEFAULT_MAX_OUTPUT_BYTES: usize = 10 * 1024 * 1024;

/// Configuration for Firecracker microVM sandbox execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FirecrackerConfig {
    /// Path to the boot kernel image (vmlinux ELF). **Required.**
    pub kernel_image_path: PathBuf,
    /// Path to the root filesystem image (ext4). **Required.**
    pub rootfs_path: PathBuf,
    /// Whether the rootfs is mounted read-only (default: true). Recommended
    /// to keep this true — agents that need scratch space should write to
    /// `/tmp` inside the VM.
    pub rootfs_read_only: bool,
    /// Kernel boot arguments. Defaults to a quiet console + read-only root.
    pub boot_args: String,
    /// vCPU count for the microVM.
    pub vcpus: u8,
    /// Memory size (MiB) for the microVM.
    pub mem_mib: u32,
    /// Path to the `firecracker` binary (default: resolved from $PATH).
    pub firecracker_binary: String,
    /// Maximum execution time before the VM is force-killed.
    pub max_execution_time: Duration,
    /// Maximum captured output bytes per stream before truncation.
    pub max_output_bytes: usize,
    /// Optional path where per-execution VM configs and serial logs are
    /// written. Defaults to `$TMPDIR`.
    pub work_dir: Option<PathBuf>,
}

impl Default for FirecrackerConfig {
    fn default() -> Self {
        Self {
            kernel_image_path: PathBuf::new(),
            rootfs_path: PathBuf::new(),
            rootfs_read_only: true,
            boot_args: "console=ttyS0 reboot=k panic=1 pci=off ro".to_string(),
            vcpus: DEFAULT_VCPUS,
            mem_mib: DEFAULT_MEM_MIB,
            firecracker_binary: "firecracker".to_string(),
            max_execution_time: DEFAULT_MAX_EXECUTION,
            max_output_bytes: DEFAULT_MAX_OUTPUT_BYTES,
            work_dir: None,
        }
    }
}

impl FirecrackerConfig {
    fn validate(&self) -> Result<(), anyhow::Error> {
        if self.kernel_image_path.as_os_str().is_empty() {
            anyhow::bail!(
                "Firecracker sandbox is not configured: missing `kernel_image_path`. \
                 Set `[sandbox.firecracker] kernel_image_path = \"...\"` in symbiont.toml \
                 or pass a kernel image when constructing FirecrackerConfig."
            );
        }
        if self.rootfs_path.as_os_str().is_empty() {
            anyhow::bail!(
                "Firecracker sandbox is not configured: missing `rootfs_path`. \
                 Set `[sandbox.firecracker] rootfs_path = \"...\"` in symbiont.toml \
                 or pass a rootfs image when constructing FirecrackerConfig."
            );
        }
        if !self.kernel_image_path.exists() {
            anyhow::bail!(
                "Firecracker kernel image not found at {}",
                self.kernel_image_path.display()
            );
        }
        if !self.rootfs_path.exists() {
            anyhow::bail!(
                "Firecracker rootfs image not found at {}",
                self.rootfs_path.display()
            );
        }
        if self.vcpus == 0 {
            anyhow::bail!("Firecracker vcpus must be >= 1");
        }
        if self.mem_mib < 64 {
            anyhow::bail!("Firecracker mem_mib must be >= 64 (got {})", self.mem_mib);
        }
        Ok(())
    }
}

/// Firecracker microVM sandbox runner.
pub struct FirecrackerRunner {
    config: FirecrackerConfig,
}

impl FirecrackerRunner {
    /// Construct a new runner. Returns `Err` if the configuration is
    /// incomplete or the supplied kernel/rootfs paths don't exist.
    pub fn new(config: FirecrackerConfig) -> Result<Self, anyhow::Error> {
        config.validate()?;

        // Preflight: confirm the firecracker binary is reachable.
        let check = std::process::Command::new(&config.firecracker_binary)
            .arg("--version")
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .status();
        match check {
            Ok(status) if status.success() => {}
            _ => {
                anyhow::bail!(
                    "Firecracker is not available at '{}'. Install it from \
                     https://github.com/firecracker-microvm/firecracker/releases or set \
                     `firecracker_binary` to a valid path.",
                    config.firecracker_binary
                );
            }
        }

        tracing::info!(
            "Firecracker sandbox initialized: kernel={}, rootfs={}, vcpus={}, mem_mib={}",
            config.kernel_image_path.display(),
            config.rootfs_path.display(),
            config.vcpus,
            config.mem_mib
        );

        Ok(Self { config })
    }

    /// Build the Firecracker VM JSON config. Pure function — testable
    /// without invoking the binary.
    pub fn build_vm_config_json(&self, serial_log: &std::path::Path) -> serde_json::Value {
        serde_json::json!({
            "boot-source": {
                "kernel_image_path": self.config.kernel_image_path,
                "boot_args": self.config.boot_args,
            },
            "drives": [{
                "drive_id": "rootfs",
                "path_on_host": self.config.rootfs_path,
                "is_root_device": true,
                "is_read_only": self.config.rootfs_read_only,
            }],
            "machine-config": {
                "vcpu_count": self.config.vcpus,
                "mem_size_mib": self.config.mem_mib,
                "smt": false,
            },
            "logger": {
                "log_path": serial_log,
                "level": "Warning",
                "show_level": false,
                "show_log_origin": false,
            }
        })
    }

    async fn read_limited<R: AsyncReadExt + Unpin>(
        reader: &mut R,
        max_bytes: usize,
    ) -> (String, bool) {
        let mut buf = vec![0u8; max_bytes + 1];
        let mut total = 0usize;
        loop {
            match reader.read(&mut buf[total..]).await {
                Ok(0) => break,
                Ok(n) => {
                    total += n;
                    if total > max_bytes {
                        total = max_bytes;
                        break;
                    }
                }
                Err(_) => break,
            }
        }
        let truncated = total == max_bytes;
        let s = String::from_utf8_lossy(&buf[..total]).to_string();
        if truncated {
            (
                format!("{}\n... [output truncated at {} bytes]", s, max_bytes),
                true,
            )
        } else {
            (s, false)
        }
    }
}

#[async_trait]
impl SandboxRunner for FirecrackerRunner {
    async fn execute(
        &self,
        code: &str,
        env: HashMap<String, String>,
    ) -> Result<ExecutionResult, anyhow::Error> {
        // Per-execution work directory holds the VM config and serial log.
        let work_root = self
            .config
            .work_dir
            .clone()
            .unwrap_or_else(std::env::temp_dir);
        let exec_id = uuid::Uuid::new_v4();
        let work_dir = work_root.join(format!("symbi-fc-{}", exec_id));
        tokio::fs::create_dir_all(&work_dir).await?;

        let config_path = work_dir.join("vm-config.json");
        let serial_log = work_dir.join("serial.log");
        let code_path = work_dir.join("code");

        // The supplied code lands on the host filesystem; the operator's
        // rootfs init script is responsible for mounting it (typically via
        // a vsock-attached host share) and running it. We expose the path
        // via the FIRECRACKER_CODE_PATH env var which the init can pick up
        // from /proc/cmdline if needed.
        tokio::fs::write(&code_path, code).await?;

        // Encode env for the in-VM init to consume.
        let env_json = serde_json::to_string(&env)?;
        tokio::fs::write(work_dir.join("env.json"), env_json).await?;

        let vm_config = self.build_vm_config_json(&serial_log);
        tokio::fs::write(&config_path, serde_json::to_string_pretty(&vm_config)?).await?;

        let mut cmd = Command::new(&self.config.firecracker_binary);
        cmd.arg("--no-api")
            .arg("--config-file")
            .arg(&config_path)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let start = std::time::Instant::now();
        let max_output = self.config.max_output_bytes;
        let mut child = cmd.spawn().map_err(|e| {
            anyhow::anyhow!(
                "Failed to spawn firecracker process: {}. Is firecracker installed?",
                e
            )
        })?;

        let mut child_stdout = child.stdout.take();
        let mut child_stderr = child.stderr.take();

        let result = timeout(self.config.max_execution_time, async {
            let stdout_fut = async {
                match child_stdout.as_mut() {
                    Some(s) => Self::read_limited(s, max_output).await,
                    None => (String::new(), false),
                }
            };
            let stderr_fut = async {
                match child_stderr.as_mut() {
                    Some(s) => Self::read_limited(s, max_output).await,
                    None => (String::new(), false),
                }
            };
            let ((stdout, stdout_trunc), (stderr, stderr_trunc)) =
                tokio::join!(stdout_fut, stderr_fut);
            let status = child.wait().await;
            (stdout, stdout_trunc, stderr, stderr_trunc, status)
        })
        .await;

        let elapsed = start.elapsed();

        match result {
            Ok((stdout, stdout_trunc, stderr, stderr_trunc, Ok(status))) => {
                let exit_code = status.code().unwrap_or(-1);
                let _ = tokio::fs::remove_dir_all(&work_dir).await;
                Ok(ExecutionResult {
                    exit_code,
                    stdout,
                    stderr,
                    execution_time_ms: elapsed.as_millis() as u64,
                    success: status.success(),
                    stdout_truncated: stdout_trunc,
                    stderr_truncated: stderr_trunc,
                })
            }
            Ok((stdout, stdout_trunc, stderr, stderr_trunc, Err(e))) => {
                let _ = tokio::fs::remove_dir_all(&work_dir).await;
                Ok(ExecutionResult {
                    exit_code: -1,
                    stdout,
                    stderr: format!("{}\nfirecracker wait failed: {}", stderr, e),
                    execution_time_ms: elapsed.as_millis() as u64,
                    success: false,
                    stdout_truncated: stdout_trunc,
                    stderr_truncated: stderr_trunc,
                })
            }
            Err(_) => {
                let _ = child.kill().await;
                let _ = tokio::fs::remove_dir_all(&work_dir).await;
                Ok(ExecutionResult::failure(
                    -1,
                    format!(
                        "Firecracker microVM exceeded max_execution_time ({:?})",
                        self.config.max_execution_time
                    ),
                    elapsed.as_millis() as u64,
                ))
            }
        }
    }
}

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

    #[test]
    fn missing_kernel_path_rejected() {
        let cfg = FirecrackerConfig::default();
        let err = cfg.validate().expect_err("validation should fail");
        let msg = err.to_string();
        assert!(msg.contains("kernel_image_path"), "msg: {}", msg);
    }

    #[test]
    fn missing_rootfs_path_rejected() {
        let cfg = FirecrackerConfig {
            kernel_image_path: PathBuf::from("/some/kernel"),
            ..Default::default()
        };
        let err = cfg.validate().expect_err("validation should fail");
        let msg = err.to_string();
        assert!(msg.contains("rootfs_path"), "msg: {}", msg);
    }

    #[test]
    fn vm_config_carries_kernel_and_rootfs() {
        let cfg = FirecrackerConfig {
            kernel_image_path: PathBuf::from("/k/vmlinux"),
            rootfs_path: PathBuf::from("/k/rootfs.ext4"),
            vcpus: 2,
            mem_mib: 1024,
            ..Default::default()
        };

        // We can't construct a FirecrackerRunner here (preflight requires
        // the binary + real files) so build the JSON via a stand-in.
        let runner = FirecrackerRunner { config: cfg };
        let json = runner.build_vm_config_json(std::path::Path::new("/tmp/x"));

        assert_eq!(json["machine-config"]["vcpu_count"], 2);
        assert_eq!(json["machine-config"]["mem_size_mib"], 1024);
        assert_eq!(json["drives"][0]["is_root_device"], true);
        assert_eq!(json["drives"][0]["is_read_only"], true);
        assert_eq!(json["boot-source"]["kernel_image_path"], "/k/vmlinux");
    }

    #[test]
    fn invalid_vcpus_rejected() {
        // Use NamedTempFile so the paths actually exist — otherwise validate()
        // would fail on the path check before reaching the vcpus check.
        let kernel = tempfile::NamedTempFile::new().unwrap();
        let rootfs = tempfile::NamedTempFile::new().unwrap();
        let cfg = FirecrackerConfig {
            kernel_image_path: kernel.path().to_path_buf(),
            rootfs_path: rootfs.path().to_path_buf(),
            vcpus: 0,
            ..Default::default()
        };
        let err = cfg.validate().expect_err("vcpus=0 should be rejected");
        assert!(err.to_string().contains("vcpus"));
    }
}
