env_file.rs

use crate::parser::{file, FileLine};
use nom_locate::LocatedSpan;
use std::env;
use std::env::VarError;
use std::fmt::Display;
use std::str::FromStr;
use std::string::ParseError;

#[derive(Clone, Debug, Default)]
pub struct EnvironmentFile {
	lines: Vec<FileLine>,
}

#[derive(Debug, thiserror::Error)]
pub enum EnvironmentFileError {
	#[error("Unable to determine eof marker")]
	InvalidEof,
	#[error("Parse error at line {line}, column {column}: {kind:?}")]
	ParseError {
		line: usize,
		column: usize,
		kind: nom::error::ErrorKind,
	},
	#[error(transparent)]
	StdParseError(#[from] ParseError),
	#[error("Value could not be determined to be valid UTF-8")]
	InvalidValue,
}

impl FromStr for EnvironmentFile {
	type Err = EnvironmentFileError;

	fn from_str(s: &str) -> Result<Self, Self::Err> {
		let (_, lines) = file(LocatedSpan::new(s)).map_err(|err| match err {
			nom::Err::Incomplete(_) => EnvironmentFileError::InvalidEof,
			nom::Err::Error(err) | nom::Err::Failure(err) => EnvironmentFileError::ParseError {
				column: err.input.get_column(),
				line: err.input.location_line() as usize,
				kind: err.code,
			},
		})?;

		Ok(Self::new(lines))
	}
}

fn is_kv_line(line: &FileLine) -> bool {
	matches!(line, FileLine::KeyValue { .. })
}

fn is_comment_line(line: &FileLine) -> bool {
	matches!(line, FileLine::Comment(_))
}

impl EnvironmentFile {
	pub fn new(lines: Vec<FileLine>) -> Self {
		Self { lines }
	}

	pub fn parse(value: impl AsRef<str>) -> Result<Self, EnvironmentFileError> {
		value.as_ref().parse()
	}

	pub fn iter(&self) -> EnvFileIterator {
		EnvFileIterator {
			lines: &self.lines,
			current: 0,
		}
	}

	pub fn lines_kv(&self) -> SubTypeFileIterator {
		SubTypeFileIterator {
			lines: &self.lines,
			current: 0,
			predicate: is_kv_line,
		}
	}

	pub fn lines_comment(&self) -> SubTypeFileIterator {
		SubTypeFileIterator {
			lines: &self.lines,
			current: 0,
			predicate: is_comment_line,
		}
	}

	pub fn get_raw(&self, key: &str) -> Option<String> {
		for line in &self.lines {
			if let FileLine::KeyValue { key: k, value } = line {
				if k == key {
					return Some(value.iter().map(|part| part.to_string()).collect::<String>());
				}
			}
		}

		None
	}

	pub fn len(&self) -> usize {
		self.lines.len()
	}

	pub fn is_empty(&self) -> bool {
		self.lines.is_empty()
	}
}

#[derive(Clone, Debug, Default)]
pub struct ApplyOptions {
	pub prefix: Option<String>,
	pub overwrite: bool,
}

impl ApplyOptions {
	pub fn new(prefix: impl Display, overwrite: bool) -> Self {
		Self {
			prefix: Some(prefix.to_string()),
			overwrite,
		}
	}

	pub fn with_prefix(prefix: impl Display) -> Self {
		Self::new(prefix, false)
	}

	pub fn with_overwrite(overwrite: bool) -> Self {
		Self { prefix: None, overwrite }
	}
}

pub trait ApplyEnvironmentFile {
	fn apply(&self, options: ApplyOptions);
}

impl ApplyEnvironmentFile for EnvironmentFile {
	fn apply(&self, options: ApplyOptions) {
		let _ = set_from_file(self, options);
	}
}

impl<E> ApplyEnvironmentFile for Result<EnvironmentFile, E> {
	fn apply(&self, options: ApplyOptions) {
		if let Ok(file) = self {
			file.apply(options);
		}
	}
}

impl ApplyEnvironmentFile for Option<EnvironmentFile> {
	fn apply(&self, options: ApplyOptions) {
		if let Some(file) = self {
			file.apply(options);
		}
	}
}

pub struct EnvFileIterator<'a> {
	lines: &'a [FileLine],
	current: usize,
}

impl<'a> Iterator for EnvFileIterator<'a> {
	type Item = &'a FileLine;

	fn next(&mut self) -> Option<Self::Item> {
		if self.current >= self.lines.len() {
			None
		} else {
			let line = &self.lines[self.current];
			self.current += 1;
			Some(line)
		}
	}
}

impl ExactSizeIterator for EnvFileIterator<'_> {
	fn len(&self) -> usize {
		self.lines.len() - self.current
	}
}

pub struct SubTypeFileIterator<'a> {
	lines: &'a [FileLine],
	current: usize,
	predicate: fn(&FileLine) -> bool,
}

impl<'a> Iterator for SubTypeFileIterator<'a> {
	type Item = &'a FileLine;

	fn next(&mut self) -> Option<Self::Item> {
		let mut found = None;

		while self.current < self.lines.len() {
			let line = &self.lines[self.current];
			self.current += 1;
			if (self.predicate)(line) {
				found = Some(line);
				break;
			}
		}

		found
	}
}

fn is_missing_key(key: &str) -> Result<bool, EnvironmentFileError> {
	match env::var(key) {
		Ok(_) => Ok(false),
		Err(VarError::NotPresent) => Ok(true),
		Err(_) => Err(EnvironmentFileError::InvalidValue),
	}
}

fn set_from_file(file: &EnvironmentFile, options: ApplyOptions) -> Result<(), EnvironmentFileError> {
	let mut defferred = Vec::with_capacity(file.len());

	for line in file.lines_kv() {
		if let FileLine::KeyValue { key, .. } = &line {
			let key_name = if let Some(prefix) = &options.prefix {
				format!("{}{}", prefix, key)
			} else {
				key.to_string()
			};

			if options.overwrite || is_missing_key(&key_name)? {
				if line.is_complete() {
					eprintln!("Setting {} with value {}", key_name, line.assemble_value());
					env::set_var(key_name, line.assemble_value());
				} else {
					defferred.push(line);
				}
			}
		}
	}

	for line in defferred {
		if let FileLine::KeyValue { key, .. } = line {
			let key_name = if let Some(prefix) = &options.prefix {
				format!("{}{}", prefix, key)
			} else {
				key.to_string()
			};
			env::set_var(key_name, line.assemble_value());
		}
	}

	Ok(())
}

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

	#[test]
	fn test_lines_kv_iterator() {
		let env_file = EnvironmentFile::new(vec![
			FileLine::empty(),
			FileLine::comment("# This is a comment"),
			FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]),
			FileLine::empty(),
			FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]),
		]);

		let kv_lines: Vec<&FileLine> = env_file.lines_kv().collect();

		assert_eq!(kv_lines.len(), 2);
		assert!(
			matches!(kv_lines[0], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())])
		);
		assert!(
			matches!(kv_lines[1], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())])
		);
	}

	#[test]
	fn test_lines_comment_iterator() {
		let env_file = EnvironmentFile::new(vec![
			FileLine::empty(),
			FileLine::comment(" Comment 1"),
			FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]),
			FileLine::comment(" Comment 2"),
			FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]),
			FileLine::comment(" Comment 3"),
		]);

		let comment_lines: Vec<&FileLine> = env_file.lines_comment().collect();

		assert_eq!(comment_lines.len(), 3);
		assert!(matches!(comment_lines[0], FileLine::Comment(comment) if comment == " Comment 1"));
		assert!(matches!(comment_lines[1], FileLine::Comment(comment) if comment == " Comment 2"));
		assert!(matches!(comment_lines[2], FileLine::Comment(comment) if comment == " Comment 3"));
	}

	#[test]
	fn test_iter_returns_all_lines() {
		let env_file = EnvironmentFile::new(vec![
			FileLine::empty(),
			FileLine::comment("This is a comment"),
			FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]),
			FileLine::empty(),
			FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]),
		]);

		let all_lines: Vec<&FileLine> = env_file.iter().collect();

		assert_eq!(all_lines.len(), 5);
		assert!(matches!(all_lines[0], FileLine::Empty));
		assert!(matches!(all_lines[1], FileLine::Comment(comment) if comment == "This is a comment"));
		assert!(
			matches!(all_lines[2], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())])
		);
		assert!(matches!(all_lines[3], FileLine::Empty));
		assert!(
			matches!(all_lines[4], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())])
		);
	}

	#[test]
	fn test_lines_kv_maintains_order() {
		let env_file = EnvironmentFile::new(vec![
			FileLine::empty(),
			FileLine::key_value("KEY1".to_string(), vec![ValuePart::Static("VALUE1".to_string())]),
			FileLine::comment("Comment"),
			FileLine::key_value("KEY2".to_string(), vec![ValuePart::Static("VALUE2".to_string())]),
			FileLine::key_value("KEY3".to_string(), vec![ValuePart::Static("VALUE3".to_string())]),
			FileLine::empty(),
		]);

		let kv_lines: Vec<&FileLine> = env_file.lines_kv().collect();

		assert_eq!(kv_lines.len(), 3);
		assert!(
			matches!(kv_lines[0], FileLine::KeyValue { key, value } if key == "KEY1" && value == &vec![ValuePart::Static("VALUE1".to_string())])
		);
		assert!(
			matches!(kv_lines[1], FileLine::KeyValue { key, value } if key == "KEY2" && value == &vec![ValuePart::Static("VALUE2".to_string())])
		);
		assert!(
			matches!(kv_lines[2], FileLine::KeyValue { key, value } if key == "KEY3" && value == &vec![ValuePart::Static("VALUE3".to_string())])
		);
	}
}