//! FireCracker hypervisor backend for PlasmaVMC
//!
//! This crate provides the FireCracker backend implementation for the HypervisorBackend trait.
//! It uses FireCracker microVM to run lightweight virtual machines.

mod api;
mod env;

use async_trait::async_trait;
use env::{
    resolve_firecracker_path, resolve_jailer_path, resolve_kernel_path, resolve_rootfs_path,
    resolve_initrd_path, resolve_runtime_dir, resolve_socket_base_path, resolve_boot_args,
    resolve_use_jailer,
};
use api::FireCrackerClient;
use plasmavmc_hypervisor::{BackendCapabilities, HypervisorBackend, UnsupportedReason};
use plasmavmc_types::{
    AttachedDisk, DiskBus, Error, HypervisorType, NetworkSpec, NicModel, Result, VirtualMachine,
    VmHandle, VmSpec, VmStatus, VmState,
};
use plasmavmc_types::FireCrackerConfig;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio::process::Command;
use tokio::time::Instant;

/// FireCracker hypervisor backend
pub struct FireCrackerBackend {
    /// Path to FireCracker binary
    firecracker_path: PathBuf,
    /// Path to Jailer binary (optional)
    jailer_path: Option<PathBuf>,
    /// Runtime directory for VM state
    runtime_dir: PathBuf,
    /// Base path for FireCracker API sockets
    socket_base_path: PathBuf,
    /// Kernel image path
    kernel_path: PathBuf,
    /// Rootfs image path
    rootfs_path: PathBuf,
    /// Initrd image path (optional)
    initrd_path: Option<PathBuf>,
    /// Boot arguments
    boot_args: String,
    /// Use jailer for security
    use_jailer: bool,
}

impl FireCrackerBackend {
    /// Create a new FireCracker backend from config (with env fallbacks)
    pub fn from_config(config: &FireCrackerConfig) -> Result<Self> {
        let kernel_path = config
            .kernel_path
            .clone()
            .or_else(resolve_kernel_path)
            .ok_or_else(|| {
                Error::HypervisorError(
                    "PLASMAVMC_FIRECRACKER_KERNEL_PATH not set".to_string(),
                )
            })?;
        let rootfs_path = config
            .rootfs_path
            .clone()
            .or_else(resolve_rootfs_path)
            .ok_or_else(|| {
                Error::HypervisorError(
                    "PLASMAVMC_FIRECRACKER_ROOTFS_PATH not set".to_string(),
                )
            })?;

        let firecracker_path = config
            .firecracker_path
            .clone()
            .unwrap_or_else(resolve_firecracker_path);
        let jailer_path = config
            .jailer_path
            .clone()
            .or_else(resolve_jailer_path);
        let runtime_dir = config
            .runtime_dir
            .clone()
            .unwrap_or_else(resolve_runtime_dir);
        let socket_base_path = config
            .socket_base_path
            .clone()
            .unwrap_or_else(resolve_socket_base_path);
        let initrd_path = config
            .initrd_path
            .clone()
            .or_else(resolve_initrd_path);
        let boot_args = config
            .boot_args
            .clone()
            .unwrap_or_else(resolve_boot_args);
        let use_jailer = config
            .use_jailer
            .unwrap_or_else(|| config.jailer_path.is_some() || resolve_use_jailer());

        Ok(Self {
            firecracker_path,
            jailer_path,
            runtime_dir,
            socket_base_path,
            kernel_path,
            rootfs_path,
            initrd_path,
            boot_args,
            use_jailer,
        })
    }

    /// Create a new FireCracker backend from environment variables
    pub fn from_env() -> Result<Self> {
        Self::from_config(&FireCrackerConfig::default())
    }

    /// Create with default paths (for testing)
    pub fn with_defaults(
        kernel_path: impl Into<PathBuf>,
        rootfs_path: impl Into<PathBuf>,
    ) -> Self {
        Self {
            firecracker_path: PathBuf::from("/usr/bin/firecracker"),
            jailer_path: None,
            runtime_dir: PathBuf::from("/var/run/plasmavmc/firecracker"),
            socket_base_path: PathBuf::from("/tmp/firecracker"),
            kernel_path: kernel_path.into(),
            rootfs_path: rootfs_path.into(),
            initrd_path: None,
            boot_args: "console=ttyS0".to_string(),
            use_jailer: false,
        }
    }

    fn api_socket_path(&self, handle: &VmHandle) -> PathBuf {
        if let Some(path) = handle.backend_state.get("api_socket") {
            PathBuf::from(path)
        } else {
            self.socket_base_path.join(format!("{}.sock", handle.vm_id))
        }
    }

    fn vm_runtime_dir(&self, vm_id: &plasmavmc_types::VmId) -> PathBuf {
        self.runtime_dir.join(vm_id.to_string())
    }
}

/// Wait for FireCracker API socket to become available.
async fn wait_for_socket(socket_path: &Path, timeout: Duration) -> Result<()> {
    let start = Instant::now();
    loop {
        match tokio::net::UnixStream::connect(socket_path).await {
            Ok(_stream) => {
                return Ok(());
            }
            Err(_e) => {
                if start.elapsed() >= timeout {
                    return Err(Error::HypervisorError(format!(
                        "Timed out waiting for FireCracker socket {}",
                        socket_path.display()
                    )));
                }
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        }
    }
}

fn kill_pid(pid: u32) -> Result<()> {
    let status = std::process::Command::new("kill")
        .arg("-9")
        .arg(pid.to_string())
        .status()
        .map_err(|e| Error::HypervisorError(format!("Failed to invoke kill -9: {e}")))?;
    if status.success() {
        Ok(())
    } else {
        Err(Error::HypervisorError(format!(
            "kill -9 exited with status: {status}"
        )))
    }
}

#[async_trait]
impl HypervisorBackend for FireCrackerBackend {
    fn backend_type(&self) -> HypervisorType {
        HypervisorType::Firecracker
    }

    fn capabilities(&self) -> BackendCapabilities {
        BackendCapabilities {
            live_migration: false,
            hot_plug_cpu: false,
            hot_plug_memory: false,
            hot_plug_disk: false,
            hot_plug_nic: false,
            vnc_console: false,
            serial_console: true,
            nested_virtualization: false,
            gpu_passthrough: false,
            max_vcpus: 32,
            max_memory_gib: 1024,
            supported_disk_buses: vec![DiskBus::Virtio],
            supported_nic_models: vec![NicModel::VirtioNet],
        }
    }

    fn supports(&self, spec: &VmSpec) -> std::result::Result<(), UnsupportedReason> {
        // Check vCPU limit
        if spec.cpu.vcpus > 32 {
            return Err(UnsupportedReason::Feature(format!(
                "FireCracker supports max 32 vCPUs, requested {}",
                spec.cpu.vcpus
            )));
        }

        // Check disk bus types
        for disk in &spec.disks {
            if disk.bus != DiskBus::Virtio {
                return Err(UnsupportedReason::DiskBus(disk.bus));
            }
        }

        // Check NIC models
        for nic in &spec.network {
            if nic.model != NicModel::VirtioNet {
                return Err(UnsupportedReason::NicModel(nic.model));
            }
        }

        Ok(())
    }

    async fn create(&self, vm: &VirtualMachine, _disks: &[AttachedDisk]) -> Result<VmHandle> {
        tracing::info!(
            vm_id = %vm.id,
            name = %vm.name,
            "Creating FireCracker microVM"
        );

        // Validate spec
        self.supports(&vm.spec).map_err(|e| {
            Error::HypervisorError(format!("VM spec not supported by FireCracker: {e:?}"))
        })?;

        let runtime_dir = self.vm_runtime_dir(&vm.id);
        tokio::fs::create_dir_all(&runtime_dir)
            .await
            .map_err(|e| Error::HypervisorError(format!("Failed to create runtime dir: {e}")))?;

        let api_socket = self.socket_base_path.join(format!("{}.sock", vm.id));
        // Remove stale socket if it exists
        let _ = tokio::fs::remove_file(&api_socket).await;

        // Start FireCracker process
        let mut cmd = if self.use_jailer {
            let jailer_path = self.jailer_path.as_ref().ok_or_else(|| {
                Error::HypervisorError("Jailer requested but path not found".to_string())
            })?;
            let mut cmd = Command::new(jailer_path);
            cmd.args([
                "--id",
                &vm.id.to_string(),
                "--exec-file",
                self.firecracker_path.to_str().unwrap(),
                "--uid",
                "1000",
                "--gid",
                "1000",
                "--chroot-base-dir",
                runtime_dir.to_str().unwrap(),
                "--",
                "--api-sock",
                api_socket.to_str().unwrap(),
            ]);
            cmd
        } else {
            let mut cmd = Command::new(&self.firecracker_path);
            cmd.args(["--api-sock", api_socket.to_str().unwrap()]);
            cmd
        };

        tracing::debug!(
            vm_id = %vm.id,
            api_socket = %api_socket.display(),
            "Spawning FireCracker process"
        );

        let mut child = cmd
            .spawn()
            .map_err(|e| Error::HypervisorError(format!("Failed to spawn FireCracker: {e}")))?;
        let pid = child.id().map(|p| p);
        // Detach process
        tokio::spawn(async move {
            let _ = child.wait().await;
        });

        // Wait for API socket to be ready
        wait_for_socket(&api_socket, Duration::from_secs(5)).await?;

        // Configure VM via API
        let mut client = FireCrackerClient::connect(&api_socket).await?;

        // Set machine config
        client
            .put_machine_config(vm.spec.cpu.vcpus, vm.spec.memory.size_mib)
            .await?;

        // Set boot source
        let initrd_path = self.initrd_path.as_ref().map(|p| p.to_str().unwrap());
        client
            .put_boot_source(
                self.kernel_path.to_str().unwrap(),
                initrd_path,
                &self.boot_args,
            )
            .await?;

        // Set rootfs drive
        if !vm.spec.disks.is_empty() {
            client
                .put_drive(
                    "rootfs",
                    self.rootfs_path.to_str().unwrap(),
                    true,
                    false,
                )
                .await?;
        }

        // Set network interfaces
        for (idx, nic) in vm.spec.network.iter().enumerate() {
            let iface_id = format!("eth{}", idx);
            let mac = nic.mac_address.clone().unwrap_or_else(|| {
                format!("AA:FC:00:00:{:02X}:{:02X}", (vm.id.as_uuid().as_u128() >> 8) as u8, vm.id.as_uuid().as_u128() as u8)
            });
            // Note: host_dev_name should be set up externally (TAP interface)
            // For now, we'll use a placeholder
            client
                .put_network_interface(&iface_id, &mac, "tap0")
                .await?;
        }

        let mut handle = VmHandle::new(vm.id, runtime_dir.to_string_lossy().to_string());
        handle
            .backend_state
            .insert("api_socket".into(), api_socket.display().to_string());
        handle.pid = pid;

        Ok(handle)
    }

    async fn start(&self, handle: &VmHandle) -> Result<()> {
        let api_socket = self.api_socket_path(handle);
        wait_for_socket(&api_socket, Duration::from_secs(2)).await?;
        tracing::info!(
            vm_id = %handle.vm_id,
            api_socket = %api_socket.display(),
            "Starting FireCracker microVM"
        );
        let mut client = FireCrackerClient::connect(&api_socket).await?;
        client.instance_start().await?;
        Ok(())
    }

    async fn stop(&self, handle: &VmHandle, timeout: Duration) -> Result<()> {
        let api_socket = self.api_socket_path(handle);
        if let Err(e) = wait_for_socket(&api_socket, Duration::from_secs(2)).await {
            if let Some(pid) = handle.pid {
                tracing::warn!(vm_id = %handle.vm_id, pid, "API socket unavailable; sending SIGKILL");
                return kill_pid(pid);
            }
            return Err(e);
        }
        tracing::info!(
            vm_id = %handle.vm_id,
            timeout_secs = timeout.as_secs(),
            "Stopping FireCracker microVM"
        );
        // FireCracker doesn't support graceful shutdown via API
        // We'll send Ctrl+Alt+Del if ACPI is enabled, otherwise kill
        let mut client = FireCrackerClient::connect(&api_socket).await?;
        if let Err(_) = client.send_ctrl_alt_del().await {
            // Fallback to kill
            if let Some(pid) = handle.pid {
                tracing::warn!(vm_id = %handle.vm_id, pid, "Ctrl+Alt+Del failed; sending SIGKILL");
                return kill_pid(pid);
            }
        }

        // Wait for VM to stop
        let start = Instant::now();
        loop {
            let status = self.status(handle).await?;
            if matches!(status.actual_state, VmState::Stopped | VmState::Failed) {
                break;
            }
            if start.elapsed() >= timeout {
                if let Some(pid) = handle.pid {
                    tracing::warn!(vm_id = %handle.vm_id, pid, "Stop timed out; sending SIGKILL");
                    kill_pid(pid)?;
                    break;
                }
                return Err(Error::HypervisorError(format!(
                    "Timeout waiting for VM {} to stop",
                    handle.vm_id
                )));
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        Ok(())
    }

    async fn kill(&self, handle: &VmHandle) -> Result<()> {
        tracing::info!(vm_id = %handle.vm_id, "Force killing FireCracker microVM");
        if let Some(pid) = handle.pid {
            kill_pid(pid)
        } else {
            Err(Error::HypervisorError(
                "No PID available for VM".to_string(),
            ))
        }
    }

    async fn reboot(&self, handle: &VmHandle) -> Result<()> {
        tracing::info!(vm_id = %handle.vm_id, "Rebooting FireCracker microVM");
        let api_socket = self.api_socket_path(handle);
        wait_for_socket(&api_socket, Duration::from_secs(2)).await?;
        let mut client = FireCrackerClient::connect(&api_socket).await?;
        client.send_ctrl_alt_del().await?;
        Ok(())
    }

    async fn prepare_incoming(
        &self,
        _vm: &VirtualMachine,
        _listen_uri: &str,
        _disks: &[AttachedDisk],
    ) -> Result<VmHandle> {
        Err(Error::UnsupportedFeature(
            "FireCracker does not support live migration".into(),
        ))
    }

    async fn migrate(
        &self,
        _handle: &VmHandle,
        _destination_uri: &str,
        _timeout: Duration,
        _wait: bool,
    ) -> Result<()> {
        Err(Error::UnsupportedFeature(
            "FireCracker does not support live migration".into(),
        ))
    }

    async fn delete(&self, handle: &VmHandle) -> Result<()> {
        // Stop VM if running
        if let Ok(status) = self.status(handle).await {
            if matches!(status.actual_state, VmState::Running | VmState::Starting) {
                self.kill(handle).await?;
            }
        }

        // Clean up runtime directory
        let runtime_dir = PathBuf::from(&handle.runtime_dir);
        if runtime_dir.exists() {
            tokio::fs::remove_dir_all(&runtime_dir)
                .await
                .map_err(|e| Error::HypervisorError(format!("Failed to remove runtime dir: {e}")))?;
        }

        // Remove API socket
        let api_socket = self.api_socket_path(handle);
        let _ = tokio::fs::remove_file(&api_socket).await;

        tracing::info!(vm_id = %handle.vm_id, "Deleted FireCracker microVM");
        Ok(())
    }

    async fn status(&self, handle: &VmHandle) -> Result<VmStatus> {
        let api_socket = self.api_socket_path(handle);
        tracing::debug!(
            vm_id = %handle.vm_id,
            api_socket = %api_socket.display(),
            "Querying FireCracker microVM status"
        );

        // Check if process is still running
        let process_running = if let Some(pid) = handle.pid {
            std::process::Command::new("kill")
                .arg("-0")
                .arg(pid.to_string())
                .status()
                .map(|s| s.success())
                .unwrap_or(false)
        } else {
            false
        };

        if !process_running {
            return Ok(VmStatus {
                actual_state: VmState::Stopped,
                host_pid: handle.pid,
                ..VmStatus::default()
            });
        }

        // Try to get status from API
        match FireCrackerClient::connect(&api_socket).await {
            Ok(mut client) => {
                match client.get_vm_info().await {
                    Ok(_info) => {
                        // FireCracker API doesn't provide explicit running state
                        // If we can connect and get info, assume running
                        Ok(VmStatus {
                            actual_state: VmState::Running,
                            host_pid: handle.pid,
                            ..VmStatus::default()
                        })
                    }
                    Err(_) => {
                        // API error might mean VM is stopping
                        Ok(VmStatus {
                            actual_state: VmState::Stopped,
                            host_pid: handle.pid,
                            ..VmStatus::default()
                        })
                    }
                }
            }
            Err(_) => {
                // Can't connect - VM might be stopped or starting
                Ok(VmStatus {
                    actual_state: if process_running {
                        VmState::Starting
                    } else {
                        VmState::Stopped
                    },
                    host_pid: handle.pid,
                    ..VmStatus::default()
                })
            }
        }
    }

    async fn attach_disk(&self, _handle: &VmHandle, _disk: &AttachedDisk) -> Result<()> {
        Err(Error::HypervisorError(
            "FireCracker does not support hot-plugging disks".to_string(),
        ))
    }

    async fn detach_disk(&self, _handle: &VmHandle, _disk_id: &str) -> Result<()> {
        Err(Error::HypervisorError(
            "FireCracker does not support hot-unplugging disks".to_string(),
        ))
    }

    async fn attach_nic(&self, _handle: &VmHandle, _nic: &NetworkSpec) -> Result<()> {
        Err(Error::HypervisorError(
            "FireCracker does not support hot-plugging NICs".to_string(),
        ))
    }

    async fn detach_nic(&self, _handle: &VmHandle, _nic_id: &str) -> Result<()> {
        Err(Error::HypervisorError(
            "FireCracker does not support hot-unplugging NICs".to_string(),
        ))
    }
}

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

    #[test]
    fn test_firecracker_backend_creation() {
        let backend = FireCrackerBackend::with_defaults(
            "/tmp/kernel",
            "/tmp/rootfs",
        );
        assert_eq!(backend.backend_type(), HypervisorType::Firecracker);
    }

    #[test]
    fn test_firecracker_capabilities() {
        let backend = FireCrackerBackend::with_defaults(
            "/tmp/kernel",
            "/tmp/rootfs",
        );
        let caps = backend.capabilities();

        assert!(!caps.live_migration);
        assert!(!caps.vnc_console);
        assert!(caps.serial_console);
        assert_eq!(caps.max_vcpus, 32);
        assert_eq!(caps.supported_disk_buses.len(), 1);
        assert_eq!(caps.supported_disk_buses[0], DiskBus::Virtio);
    }

    #[test]
    fn test_firecracker_supports_valid_spec() {
        let backend = FireCrackerBackend::with_defaults(
            "/tmp/kernel",
            "/tmp/rootfs",
        );
        let mut spec = VmSpec::default();
        spec.cpu.vcpus = 2;
        spec.disks.push(DiskSpec {
            bus: DiskBus::Virtio,
            ..DiskSpec::default()
        });

        assert!(backend.supports(&spec).is_ok());
    }

    #[test]
    fn test_firecracker_rejects_too_many_vcpus() {
        let backend = FireCrackerBackend::with_defaults(
            "/tmp/kernel",
            "/tmp/rootfs",
        );
        let mut spec = VmSpec::default();
        spec.cpu.vcpus = 64; // Exceeds FireCracker limit of 32

        assert!(backend.supports(&spec).is_err());
    }

    #[test]
    fn test_firecracker_rejects_non_virtio_disk() {
        let backend = FireCrackerBackend::with_defaults(
            "/tmp/kernel",
            "/tmp/rootfs",
        );
        let mut spec = VmSpec::default();
        spec.disks.push(DiskSpec {
            bus: DiskBus::Scsi,
            ..DiskSpec::default()
        });

        assert!(backend.supports(&spec).is_err());
    }
}
