use alloc::{collections::BTreeMap, vec::Vec};

use crate::sync::Mutex;

static ENV: Mutex<BTreeMap<Vec<u8>, Vec<u8>>> = Mutex::new(BTreeMap::new());

/// Initialize the environment from a C-style null-terminated `envp` array.
///
/// Each entry must be a null-terminated `KEY=VALUE` byte string. If `envp` is
/// null the environment is left empty. This function is intended to be called
/// once from `_start` before any other code runs.
///
/// # Safety
/// `envp` must either be null or point to a null-terminated array of pointers,
/// each of which must point to a null-terminated byte string. All memory must
/// remain valid for the duration of this call.
pub unsafe fn init_env(envp: *const *const u8) {
    let mut guard = ENV.lock();
    guard.clear();

    if envp.is_null() {
        return;
    }

    let mut i = 0usize;
    loop {
        let entry_ptr = unsafe { *envp.add(i) };
        if entry_ptr.is_null() {
            break;
        }

        // Read the null-terminated byte string.
        let mut len = 0usize;
        while unsafe { *entry_ptr.add(len) } != 0 {
            len += 1;
        }
        let bytes = unsafe { core::slice::from_raw_parts(entry_ptr, len) };

        // Split on the first '=' to get key and value.
        if let Some(eq_pos) = bytes.iter().position(|&b| b == b'=') {
            let key = bytes[..eq_pos].to_vec();
            let value = bytes[eq_pos + 1..].to_vec();
            guard.insert(key, value);
        }
        // Entries without '=' are silently ignored (not a valid KEY=VALUE pair).

        i += 1;
    }
}

/// Look up an environment variable by key.
///
/// Returns `Some(value)` if the key exists, `None` otherwise.
pub fn get_env(key: &[u8]) -> Option<Vec<u8>> {
    let guard = ENV.lock();
    guard.get(key).cloned()
}

/// Set or overwrite an environment variable.
pub fn set_env(key: &[u8], value: &[u8]) {
    let mut guard = ENV.lock();
    guard.insert(key.to_vec(), value.to_vec());
}

/// Remove an environment variable. Does nothing if the key is not present.
pub fn unset_env(key: &[u8]) {
    let mut guard = ENV.lock();
    guard.remove(key);
}

/// Return a snapshot of all environment variables as `(key, value)` pairs.
pub fn env_iter() -> Vec<(Vec<u8>, Vec<u8>)> {
    let guard = ENV.lock();
    guard.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
}

/// Return all environment variables formatted as null-terminated `KEY=VALUE\0`
/// byte vectors, suitable for passing to `execve`-style syscalls.
pub fn env_as_c_strings() -> Vec<Vec<u8>> {
    let guard = ENV.lock();
    guard
        .iter()
        .map(|(k, v)| {
            let mut s = Vec::with_capacity(k.len() + 1 + v.len() + 1);
            s.extend_from_slice(k);
            s.push(b'=');
            s.extend_from_slice(v);
            s.push(0); // null terminator
            s
        })
        .collect()
}
