use std::ffi::{OsStr, OsString};
use std::sync::OnceLock;

use tokio::sync::Mutex as AsyncMutex;

pub fn async_env_lock() -> &'static AsyncMutex<()> {
    static LOCK: OnceLock<AsyncMutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| AsyncMutex::new(()))
}

pub fn set_env(key: &str, value: impl AsRef<OsStr>) {
    // SAFETY: tests serialize env mutation with async_env_lock.
    unsafe { std::env::set_var(key, value.as_ref()) };
}

pub fn remove_env(key: &str) {
    // SAFETY: tests serialize env mutation with async_env_lock.
    unsafe { std::env::remove_var(key) };
}

#[must_use = "EnvVarRestore restores the original env var value when dropped"]
pub struct EnvVarRestore {
    key: String,
    original: Option<OsString>,
}

impl EnvVarRestore {
    pub fn capture(key: impl Into<String>) -> Self {
        let key = key.into();
        Self {
            original: std::env::var_os(&key),
            key,
        }
    }

    pub fn capture_many<const N: usize>(keys: [&str; N]) -> [Self; N] {
        keys.map(Self::capture)
    }
}

impl Drop for EnvVarRestore {
    fn drop(&mut self) {
        match &self.original {
            Some(value) => set_env(self.key.as_str(), value),
            None => remove_env(self.key.as_str()),
        }
    }
}