quickenv 0.5.0

use std::collections::BTreeMap;
use std::ffi::{OsStr, OsString};
use std::io::{self, BufRead, BufReader, Read as _};
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};

use base64::prelude::*;
use flate2::read::ZlibDecoder;
use serde::Deserialize;

pub type Env = BTreeMap<OsString, OsString>;

pub struct EnvrcContext {
    pub envrc: std::fs::File,
    pub root: PathBuf,
    pub env_cache_path: PathBuf,
    pub env_cache_dir: PathBuf,
}

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("failed to find .envrc in current or any parent directory")]
    NoEnvrc,
    #[error(transparent)]
    Io(#[from] io::Error),
    #[error("failed to find QUICKENV_HOME or HOME")]
    NoQuickenvHome,
    #[error("failed to get current directory")]
    CurrentDir(#[source] io::Error),
}

pub fn resolve_envrc_context(quickenv_home: &Path) -> Result<EnvrcContext, Error> {
    let mut root = std::env::current_dir().map_err(Error::CurrentDir)?;

    let (envrc_path, envrc) = loop {
        let path = root.join(".envrc");
        if let Ok(f) = std::fs::File::open(&path) {
            log::debug!("loading {}", path.display());
            break (path, f);
        }

        if !root.pop() {
            return Err(Error::NoEnvrc);
        }
    };

    let env_cache_dir = quickenv_home.join("envs/");

    let mut env_hasher = blake3::Hasher::new();
    env_hasher.update(envrc_path.as_os_str().as_bytes());
    let env_cache_path = env_cache_dir.join(hex::encode(env_hasher.finalize().as_bytes()));

    Ok(EnvrcContext {
        root,
        env_cache_dir,
        envrc,
        env_cache_path,
    })
}

pub fn get_quickenv_home() -> Result<PathBuf, Error> {
    if let Ok(home) = std::env::var("QUICKENV_HOME") {
        Ok(Path::new(&home).to_owned())
    } else if let Ok(home) = std::env::var("HOME") {
        Ok(Path::new(&home).join(".quickenv/"))
    } else {
        Err(Error::NoQuickenvHome)
    }
}

pub fn parse_env_line(line: &[u8], env: &mut Env, prev_var_name: &mut Option<OsString>) {
    let mut split_iter = line.splitn(2, |&x| x == b'=');

    match split_iter
        .next()
        .and_then(|first| Some((first, split_iter.next()?)))
    {
        Some((var_name, value)) => {
            let var_name = OsString::from_vec(var_name.to_owned());
            let value = OsString::from_vec(value.to_owned());
            *prev_var_name = Some(var_name.clone());
            env.insert(var_name, value);
        }
        None => {
            let prev_value = env.get_mut(prev_var_name.as_ref().unwrap()).unwrap();
            prev_value.push(OsStr::new("\n"));
            prev_value.push(OsStr::from_bytes(line));
        }
    }
}

pub fn get_envvars(ctx: &EnvrcContext) -> Result<Option<Env>, Error> {
    if let Ok(file) = std::fs::File::open(&ctx.env_cache_path) {
        let mut loaded_env_cache = BTreeMap::new();
        let reader = BufReader::new(file);

        let mut prev_var_name = None;

        for line in reader.split(b'\n') {
            let raw_line = line?;
            let mut line = raw_line.as_slice();
            while let Some(b'\n') = line.last() {
                line = &line[..line.len()];
            }

            parse_env_line(line, &mut loaded_env_cache, &mut prev_var_name);
        }

        return Ok(Some(loaded_env_cache));
    }

    Ok(None)
}

#[derive(Debug, Deserialize)]
struct WatchedFile {
    path: PathBuf,
    modtime: i64,
}

pub fn is_cache_stale(env: &Env) -> bool {
    let watches = match env.get(OsStr::new("DIRENV_WATCHES")) {
        Some(w) => w,
        None => {
            log::debug!("no DIRENV_WATCHES in cached env, assuming not stale");
            return false;
        }
    };

    let watches_str = match watches.to_str() {
        Some(s) => s,
        None => {
            log::debug!("DIRENV_WATCHES is not valid UTF-8");
            return false;
        }
    };

    let watched_files = match decode_direnv_watches(watches_str) {
        Ok(files) => files,
        Err(e) => {
            log::debug!("failed to decode DIRENV_WATCHES: {e}");
            return false;
        }
    };

    for watch in watched_files {
        let current_mtime = match std::fs::metadata(&watch.path).and_then(|m| m.modified()) {
            Ok(mtime) => mtime
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs() as i64)
                .unwrap_or(0),
            Err(e) => {
                log::debug!("failed to stat {}: {e}", watch.path.display());
                return true;
            }
        };

        if current_mtime != watch.modtime {
            log::debug!(
                "{} changed: cached mtime={}, current mtime={}",
                watch.path.display(),
                watch.modtime,
                current_mtime
            );
            return true;
        }
    }

    false
}

fn decode_direnv_watches(encoded: &str) -> Result<Vec<WatchedFile>, Error> {
    use base64::engine::general_purpose::URL_SAFE;
    let bytes = URL_SAFE
        .decode(encoded)
        .or_else(|_| BASE64_URL_SAFE_NO_PAD.decode(encoded))
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

    let mut decoder = ZlibDecoder::new(&bytes[..]);
    let mut json_bytes = Vec::new();
    decoder
        .read_to_end(&mut json_bytes)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

    let files: Vec<WatchedFile> = serde_json::from_slice(&json_bytes)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

    Ok(files)
}

#[test]
fn test_decode_direnv_watches() {
    let encoded = "eJwEwEEKwjAQBdC7_HVxWqyIuYcrcVH0g1GahJkfUcS7911-aIseSLBz0MN6Udbyolvz-uRNYcEi_9qO5Y0Ba70rr0SajuNh3p_maRzATw4Fkrzzf90CAAD__wHyHR8";
    let files = decode_direnv_watches(encoded).unwrap();
    assert_eq!(files.len(), 1);
    assert_eq!(
        files[0].path.to_str().unwrap(),
        "/Users/untitaker/projects/sentry/.env"
    );
    assert_eq!(files[0].modtime, 1705439410);
}