use std::{
    collections::HashMap,
    fs,
    io::Read,
    path::{Path, PathBuf},
    str::FromStr,
    thread,
    time::{Duration, SystemTime},
};

use dir_lock::DirLock;
use eyre::{Context, Result, bail, eyre};
use nanoid::nanoid;
use serde::{Deserialize, Serialize};
use strum::Display;
use sysinfo::{ProcessRefreshKind, ProcessesToUpdate};
use tracing::{debug, info, instrument};

use crate::{
    error::Error,
    qemu::QemuLaunchOpts,
    ssh::SshLaunchOpts,
    utils::{HEX_ALPHABET, VmexecDirs, lock_state_dir},
};

/// A virtual machine identifier
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct VmId(String);

impl Default for VmId {
    fn default() -> Self {
        VmId(nanoid!(12, &HEX_ALPHABET))
    }
}

impl VmId {
    pub fn new(s: String) -> Result<Self, String> {
        if s.is_empty()
            || !s
                .chars()
                .all(|c| c.is_alphanumeric() || c == '-' || c == '_')
        {
            return Err(format!(
                "Invalid vmid '{}': must contain only alphanumeric characters, hyphens, or underscores",
                s
            ));
        }
        Ok(Self(s))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }

    pub fn into_inner(self) -> String {
        self.0
    }
}

impl std::fmt::Display for VmId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl FromStr for VmId {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::new(s.to_string())
    }
}

/// File containing VM config in the VM's state dir
pub const VM_CONFIG_FILENAME: &str = "vm.json";

/// File with vmexec's pid in the VM's state dir
pub const VMEXEC_PID_FILENAME: &str = "vmexec.pid";

/// File with QEMU's pid in the VM's state dir
pub const QEMU_PID_FILENAME: &str = "qemu.pid";

/// File with the VM's vsock CID in the VM's state dir
pub const CID_FILENAME: &str = "cid";

/// File with the daemonized vmexec's stdout and stderr in the VM's state dir
pub const DAEMON_LOG_FILENAME: &str = "daemon.log";

/// File with the VM's terminal output in the VM's state dir
pub const SESSION_LOG_FILENAME: &str = "session.log";

/// Holds some metainfo about a particular VM as well as SSH and QEMU-specific info
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VmConfig {
    pub vmid: VmId,
    pub created_at: SystemTime,
    pub qemu_launch_opts: QemuLaunchOpts,
    pub ssh_launch_opts: SshLaunchOpts,
    /// Whether to delete the VM data after running.
    pub rm: bool,
}
impl VmConfig {
    pub fn new(
        vmid: VmId,
        qemu_launch_opts: QemuLaunchOpts,
        ssh_launch_opts: SshLaunchOpts,
        rm: bool,
    ) -> Self {
        Self {
            vmid,
            created_at: SystemTime::now(),
            qemu_launch_opts,
            ssh_launch_opts,
            rm,
        }
    }
}

/// Check whether the given `pid` is currently an active process
fn pid_exists(pid: usize) -> bool {
    let mut system = sysinfo::System::default();
    system.refresh_processes_specifics(
        ProcessesToUpdate::Some(&[pid.into()]),
        true,
        ProcessRefreshKind::default(),
    ) > 0
}

/// Info for a vm
#[derive(Debug)]
pub struct Vm {
    pub vm_config: Option<VmConfig>,
    pub vm_dir: PathBuf,
    pub cid_pids: CidPids,
}

#[derive(Debug, PartialEq, Display)]
pub enum VmState {
    Running,
    Exited,
    Starting,
}

impl Vm {
    pub fn state(&self) -> VmState {
        if self.cid_pids.cid.is_some()
            && self.cid_pids.qemu_pid.is_some()
            && self.cid_pids.vmexec_pid.is_some()
        {
            // If everything is set, we're happily running.
            VmState::Running
        } else if self.cid_pids.vmexec_pid.is_some() {
            // If only the vmexec pid is currently set then we're likely in startup.
            VmState::Starting
        } else {
            // For any other case, we'll assume we exited.
            VmState::Exited
        }
    }
}

/// Contains potentially live cid and pids
#[derive(Debug)]
pub struct CidPids {
    pub cid: Option<u32>,
    pub qemu_pid: Option<usize>,
    pub vmexec_pid: Option<usize>,
}

/// Get a list of [`VmConfig`]s and their paths
///
/// They may be live or have already exited.
#[instrument]
pub fn get_vms(vms_dir: &Path) -> Result<HashMap<VmId, Vm>> {
    let entries = fs::read_dir(vms_dir)?;

    let mut vms = HashMap::new();
    for entry in entries {
        let vm_dir = entry?.path();

        if vm_dir.is_dir() {
            // If the vm_config file exists, we'll try to read it. If it doesn't exist, either the
            // VM isn't well-formed or it's still starting up.
            let vm_config_file = vm_dir.join(VM_CONFIG_FILENAME);
            let vm_config: Option<VmConfig> =
                if let Ok(vm_config_str) = fs::read_to_string(&vm_config_file) {
                    Some(
                        serde_json::from_str(&vm_config_str)
                            .context(format!("Error reading file {:?}", vm_config_file))?,
                    )
                } else {
                    None
                };

            // The cid should exist very soon after launching the VM but not quite immediately.
            let cid = if let Ok(cid_str) = fs::read_to_string(vm_dir.join(CID_FILENAME)) {
                let cid = cid_str.trim().parse::<u32>()?;
                Some(cid)
            } else {
                None
            };

            // The qemu.pid file could exist but it might also not yet or not anymore.
            let qemu_pid =
                if let Ok(qemu_pid_str) = fs::read_to_string(vm_dir.join(QEMU_PID_FILENAME)) {
                    let qemu_pid = qemu_pid_str.trim().parse::<usize>()?;
                    if pid_exists(qemu_pid) {
                        Some(qemu_pid)
                    } else {
                        None
                    }
                } else {
                    None
                };

            // The vmexec.pid exists for as long as vmexec is running. If it doesn't exist anymore,
            // it means the VM has exited.
            let vmexec_pid =
                if let Ok(vmexec_pid_str) = fs::read_to_string(vm_dir.join(VMEXEC_PID_FILENAME)) {
                    let vmexec_pid = vmexec_pid_str.trim().parse::<usize>()?;
                    if pid_exists(vmexec_pid) {
                        Some(vmexec_pid)
                    } else {
                        None
                    }
                } else {
                    None
                };

            let cid_pids = CidPids {
                cid,
                qemu_pid,
                vmexec_pid,
            };

            // The VM dir name is guaranteed to be the vmid.
            let name = vm_dir.file_name().unwrap().to_string_lossy().to_string();
            let vmid = VmId::new(name).map_err(|e| eyre!(e))?;
            let vm = Vm {
                vm_config,
                vm_dir,
                cid_pids,
            };
            vms.insert(vmid.clone(), vm);
        }
    }
    Ok(vms)
}

/// Look up a single VM by id
#[instrument]
pub fn get_vm(vms_dir: &Path, vmid: &VmId) -> Result<Vm, Error> {
    get_vms(vms_dir)?
        .remove(vmid)
        .ok_or_else(|| Error::VmNotFound(vmid.clone()))
}

/// Stop a VM
///
/// This stops a VM by sending SIGTERM to its vmexec and QEMU processes.
///
/// If the VM was started with `--rm`, this waits briefly for both processes to exit and then
/// removes its state dir. The wait might timeout and cleanup proceeds regardless once it elapses.
#[instrument]
pub fn stop_vm(vms_dir: &Path, vmid: &VmId) -> Result<(), Error> {
    let vm = get_vm(vms_dir, vmid)?;

    // A VM that is still starting up or already half-exited may be missing their vmexec or qemu
    // pids. Kill whichever ones are actually live.
    for (program, program_pid) in [
        ("vmexec", vm.cid_pids.vmexec_pid),
        ("qemu", vm.cid_pids.qemu_pid),
    ] {
        let Some(pid) = program_pid else {
            debug!("No live {program} process for this VM, nothing to signal");
            continue;
        };

        debug!("Found {program} process {pid} to be live, sending kill signal");
        nix::sys::signal::kill(
            nix::unistd::Pid::from_raw(pid as i32),
            nix::sys::signal::SIGTERM,
        )
        .map_err(|e| Error::Other(e.into()))?;
    }

    // Figure out whether we're done or whether --rm was passed.
    // If --rm was not passed, we can just return right away.
    let rm_flag = vm.vm_config.as_ref().map(|c| c.rm).unwrap_or(false);
    if !rm_flag {
        return Ok(());
    }

    // Clean up after running since this has been requested.
    debug!("Waiting for processes to exit before cleanup...");
    let mut retries = 0;
    loop {
        let current_vm = get_vm(vms_dir, vmid)?;
        if current_vm.cid_pids.qemu_pid.is_none() && current_vm.cid_pids.vmexec_pid.is_none() {
            break;
        }
        if retries > 20 {
            debug!("Processes still running after waiting, proceeding with cleanup anyway");
            break;
        }
        thread::sleep(Duration::from_millis(100));
        retries += 1;
    }

    debug!("Removing VM dir at {:?}", &vm.vm_dir);
    fs::remove_dir_all(&vm.vm_dir).map_err(|e| Error::Other(e.into()))?;

    Ok(())
}

/// A freshly reserved state dir for a new VM
///
/// We need this so we have a guaranteed `vmid` and `vm_dir` while we do some other plumbing in the meantime.
/// It also holds the state dir lock. As such, make sure to only have this struct around for as long
/// as needed.
pub struct VmReservation {
    pub vmid: VmId,
    pub vm_dir: PathBuf,
    pub(crate) state_dir_lock: DirLock,
}

/// Reserve a state dir for a new VM
///
/// Picks a random vmid if none is given and errors with [`Error::VmIdExists`] on a collision.
///
/// NOTE: Locks the state dir.
#[instrument(skip(dirs))]
pub fn reserve_vm(dirs: &VmexecDirs, vmid: Option<VmId>) -> Result<VmReservation, Error> {
    let state_dir_lock = lock_state_dir(dirs)?;

    let vmid = vmid.unwrap_or_default();

    // Check for collision with existing VMs
    let existing_vms = get_vms(&dirs.vms_dir)?;
    if existing_vms.contains_key(&vmid) {
        return Err(Error::VmIdExists(vmid));
    }

    let vm_dir = dirs.vms_dir.join(vmid.as_str());
    fs::create_dir(&vm_dir).wrap_err(format!("Couldn't make temp dir in {:?}", dirs.vms_dir))?;
    debug!("vm dir is: {:?}", vm_dir);

    Ok(VmReservation {
        vmid,
        vm_dir,
        state_dir_lock,
    })
}

/// Reap dead VMs that didn't get cleaned up
///
/// This can happen is vmexec is killed at a bad moment and doesn't have the opportunity to clean
/// up stuff.
///
/// NOTE: Locks the state dir.
#[instrument]
pub fn reap_dead_vm_dirs(dirs: &VmexecDirs) -> Result<Vec<PathBuf>> {
    // Hold the lock while reaping. A VM that is still being launched has a dir but no qemu.pid
    // yet which is exactly what this function treats as dead. `run` holds the lock across that
    // window so taking it here is what stops us from deleting a VM dir out from under it.
    let _lock = lock_state_dir(dirs)?;

    let entries = fs::read_dir(&dirs.vms_dir)?;

    let mut cleaned_vm_dirs = vec![];
    for entry in entries {
        let dir = entry?.path();

        if dir.is_dir() {
            debug!("Found existing VM dir {dir:?}, seeing if we need to clean up");
            let pid_path = dir.join(QEMU_PID_FILENAME);

            // Determine if `qemu.pid` exists and try to read its content.
            match fs::File::open(&pid_path) {
                Ok(mut file) => {
                    let mut pid_str = String::new();
                    file.read_to_string(&mut pid_str)
                        .wrap_err(format!("Failed to read {pid_path:?}"))?;

                    // Try to parse the pid into usize.
                    if let Ok(pid) = pid_str.trim().parse::<usize>() {
                        if !pid_exists(pid) {
                            // If no process with that pid exists, remove the directory
                            info!("Found dead VM dir {dir:?}, cleaning up");
                            fs::remove_dir_all(&dir).wrap_err("Failed to remove dead VM dir")?;
                            cleaned_vm_dirs.push(dir);
                        }
                    } else {
                        // Remove the directory if `qemu.pid` is present but invalid.
                        info!("Found VM dir {dir:?} with invalid qemu.pid file, cleaning up");
                        fs::remove_dir_all(&dir).wrap_err("Failed to remove dead VM dir")?;
                        cleaned_vm_dirs.push(dir);
                    }
                }
                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
                    // If `qemu.pid` does not exist, remove the directory.
                    info!("Found VM dir {dir:?} without qemu.pid file, cleaning up");
                    fs::remove_dir_all(&dir).wrap_err("Failed to remove dead VM dir")?;
                    cleaned_vm_dirs.push(dir);
                }
                Err(e) => {
                    bail!(e);
                }
            }
        }
    }

    Ok(cleaned_vm_dirs)
}

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

    #[rstest]
    #[case("myvm")]
    #[case("my-vm")]
    #[case("my_vm")]
    fn test_vmid_from_str_valid(#[case] input: &str) {
        let actual: VmId = input.parse().unwrap();
        assert_eq!(actual.as_str(), input);
    }

    #[rstest]
    #[case("", "Invalid vmid")]
    #[case(" ", "Invalid vmid")]
    #[case("my vm", "Invalid vmid")]
    #[case("my.vm", "Invalid vmid")]
    #[case("my@vm", "Invalid vmid")]
    fn test_vmid_from_str_invalid(#[case] input: &str, #[case] expected_error: &str) {
        let actual: Result<VmId, _> = input.parse();
        assert!(actual.is_err());
        assert!(actual.unwrap_err().contains(expected_error));
    }
}
