// Copyright (C) 2024-2025  Johannes Sasongko <johannes sasongko org>
// SPDX-License-Identifier: GPL-3.0-or-later

#![doc = include_str!("../README.md")]

// env references:
// POSIX.1-2024: https://pubs.opengroup.org/onlinepubs/9799919799/utilities/env.html
// BusyBox: https://www.busybox.net/downloads/BusyBox.html#env
// FreeBSD: https://man.freebsd.org/env(1)
// GNU: https://www.gnu.org/software/coreutils/manual/html_node/env-invocation.html
// illumos: https://illumos.org/man/1/env
// NetBSD: https://man.netbsd.org/env.1
// OpenBSD: https://man.openbsd.org/env.1

mod cli;
mod windows;

use std::{
	ffi::{OsStr, OsString},
	io::Write as _,
	process::Command,
};

#[derive(Debug)]
struct Error {
	kind: ErrorKind,
	message: String,
}

#[derive(Debug, PartialEq)]
enum ErrorKind {
	CliError,
	ReadError,
	InvalidDotenv,
	CwdError,
	ProgramNotFound,
	ExecFailed,
	#[cfg(not(unix))]
	AbnormalExit,
}

// This matches GNU and NetBSD.
// POSIX only specifies 126 and 127.
// BusyBox, FreeBSD, and OpenBSD use 1 for generic errors.
enum ErrorExitCode {
	GenericError = 125,
	ExecFailed = 126,
	ProgramNotFound = 127,
}

const HELP: &str = include_str!("../generated/help.txt");
const _: () = assert!(!HELP.is_empty());

fn main() {
	std::process::exit(match main_(std::env::args_os().skip(1)) {
		Ok(exit_status) => exit_status,
		Err(err) => {
			let _ = std::io::stderr().write_all(concat!(env!("CARGO_BIN_NAME"), ": ").as_bytes());
			let _ = std::io::stderr().write_all(err.message.as_bytes());
			(match err.kind {
				ErrorKind::ProgramNotFound => ErrorExitCode::ProgramNotFound,
				ErrorKind::ExecFailed => ErrorExitCode::ExecFailed,
				_ => ErrorExitCode::GenericError,
			}) as _
		}
	})
}

fn main_(args: impl IntoIterator<Item = std::ffi::OsString>) -> Result<i32, Error> {
	let options = match cli::parse(args) {
		cli::ArgsResult::Options(options) => options,
		cli::ArgsResult::Help => {
			let _ = std::io::stdout().write_all(HELP.as_bytes());
			return Ok(0);
		}
		cli::ArgsResult::Version => {
			let _ = std::io::stdout().write_all(
				concat!(env!("CARGO_BIN_NAME"), " ", env!("CARGO_PKG_VERSION"), "\n").as_bytes(),
			);
			return Ok(0);
		}
		cli::ArgsResult::Err(message) => {
			return Err(Error {
				kind: ErrorKind::CliError,
				message,
			});
		}
	};
	let cli::Options {
		actions,
		directory,
		ignore_current_envs,
		program,
		program_args,
		zero_end,
	} = options;

	#[inline(always)]
	fn set_var(name: impl AsRef<OsStr>, value: impl AsRef<OsStr>) {
		// SAFETY: This program is single-threaded.
		unsafe { std::env::set_var(name, value) }
	}
	#[inline(always)]
	fn remove_var(name: impl AsRef<OsStr>) {
		// SAFETY: This program is single-threaded.
		unsafe { std::env::remove_var(name) }
	}

	if ignore_current_envs {
		for (name, _) in std::env::vars_os() {
			remove_var(name);
		}
	}
	for action in actions {
		match action {
			cli::Action::AddFile(path) => {
				let dotenv: Box<dyn Iterator<Item = _>> = if path == "-" {
					Box::new(dotenvy::from_read_iter(std::io::stdin()))
				} else {
					Box::new(dotenvy::from_path_iter(&path).map_err(|err| Error {
						kind: ErrorKind::ReadError,
						message: format!("{path}: {err}\n", path = path.to_string_lossy()),
					})?)
				};
				for env in dotenv {
					let (name, value) = env.map_err(|err| Error {
						kind: ErrorKind::InvalidDotenv,
						message: format!(
							"{path}: {err}\n",
							path = if path == "-" {
								"<stdin>".into()
							} else {
								path.to_string_lossy()
							},
						),
					})?;
					set_var(name, value);
				}
			}
			cli::Action::AddVar(name, value) => set_var(name, value),
			cli::Action::RemoveVar(name) => remove_var(name),
		}
	}

	let Some(program) = program else {
		let mut stdout = std::io::stdout();
		for (name, value) in std::env::vars_os() {
			fn maybe_quote(string: OsString) -> String {
				// TODO: into_string() doesn't work because then the string
				// becomes unavailable to the unwrap_or_else(); this will be
				// fixed in Rust 2024 (rustc 1.85.0), allowing us to avoid the
				// to_owned() call (which copies).
				string
					.to_str()
					.and_then(|str| {
						let mut chars = str.chars();
						let should_quote_first =
							chars.next().is_some_and(|ch| ch == '"' || ch.is_control());
						let should_quote = should_quote_first || chars.any(char::is_control);
						(!should_quote).then(|| str.to_owned())
					})
					.unwrap_or_else(|| format!("{string:?}"))
			}
			if zero_end {
				let _ = stdout.write_all(name.as_encoded_bytes());
				let _ = stdout.write_all(b"=");
				let _ = stdout.write_all(value.as_encoded_bytes());
				let _ = stdout.write_all(b"\0");
			} else {
				let _ = stdout.write_all(maybe_quote(name).as_bytes());
				let _ = stdout.write_all(b"=");
				let _ = stdout.write_all(maybe_quote(value).as_bytes());
				let _ = stdout.write_all(b"\n");
			}
		}
		return Ok(0);
	};

	if let Some(directory) = directory {
		std::env::set_current_dir(directory).map_err(|err| Error {
			kind: ErrorKind::CwdError,
			message: format!("{program}: {err}\n", program = program.to_string_lossy()),
		})?;
	}

	#[cfg(windows)]
	{
		let program_locs =
			windows::get_possible_program_locations(&program, || std::env::var_os("PATH"));
		let mut first_program_not_found = None;
		for program_loc in program_locs {
			match exec_or_spawn(&program, Some(program_loc), &program_args) {
				Err(err) if err.kind == ErrorKind::ProgramNotFound => {
					if first_program_not_found.is_none() {
						first_program_not_found = Some(err);
					}
					continue;
				}
				result => return result,
			}
		}
		Err(first_program_not_found.unwrap_or_else(|| Error {
			kind: ErrorKind::ProgramNotFound,
			message: format!(
				"{program}: PATH must not be empty if command does not specify a directory",
				program = program.to_string_lossy()
			),
		}))
	}
	#[cfg(not(windows))]
	exec_or_spawn(&program, None::<&OsStr>, program_args)
}

/// Runs a command using `exec` (on Unix) or `spawn` (elsewhere).
///
/// On Unix, if the `exec` succeeds, the current process is replaced and thus
/// `Ok(_)` is never returned.
fn exec_or_spawn(
	program: impl AsRef<OsStr>,
	program_loc: Option<impl AsRef<OsStr>>,
	args: impl IntoIterator<Item = impl AsRef<OsStr>>,
) -> Result<i32, Error> {
	let program = program.as_ref();
	let mut command = program_loc
		.map(Command::new)
		.unwrap_or_else(|| Command::new(program));
	command.args(args);
	#[cfg(unix)]
	{
		use std::os::unix::process::CommandExt as _;
		let err = command.exec();
		Err(transform_error(err, program))
	}
	#[cfg(not(unix))]
	{
		let result = command.spawn().and_then(|mut proc| proc.wait());
		let exit_status = result.map_err(|err| transform_error(err, program))?;
		exit_status.code().ok_or_else(|| Error {
			// This cannot happen on Windows because there is always an exit
			// code, but maybe it happens on some non-Unix non-Windows platform.
			kind: ErrorKind::AbnormalExit,
			message: format!(
				"{program}: abnormal exit\n",
				program = program.to_string_lossy(),
			),
		})
	}
}

/// Converts [`std::io::Error`] to [`Error`].
fn transform_error(err: std::io::Error, program: impl AsRef<OsStr>) -> Error {
	let program = program.as_ref();
	Error {
		// Translate std NotFound to our ProgramNotFound.
		// This does not behave consistently across platforms for edge cases
		// like empty programs or programs with ` ` as basename.
		// This is fine; we are just passing the Rust error through.
		kind: if err.kind() == std::io::ErrorKind::NotFound {
			ErrorKind::ProgramNotFound
		} else {
			ErrorKind::ExecFailed
		},
		message: format!("{program}: {err}\n", program = program.to_string_lossy(),),
	}
}
