use std::io::Write;
use std::path::{Path, PathBuf};

macro_rules! error {
	($($args:tt)*) => {
		{ writeln!(std::io::stderr().lock(), "{}: {}", yansi::Paint::red("error").bright(), format_args!($($args)*)).ok(); }
	};
}

const CARGO_TARGET_TMPDIR: &str = env!("CARGO_TARGET_TMPDIR");
const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");

#[test]
fn ui_tests() -> std::process::ExitCode {
	if let Err(()) = do_main() {
		std::process::ExitCode::FAILURE
	} else {
		std::process::ExitCode::SUCCESS
	}
}

fn do_main() -> Result<(), ()> {
	let dir = Path::new(CARGO_MANIFEST_DIR).join("tests/ui-tests");
	let cases = std::fs::read_dir(&dir)
		.map_err(|e| error!("Failed to open directory \"tests/ui-tests\": {e}"))?;

	let bless = std::env::var_os("UI_TESTS").is_some_and(|x| x == "bless");
	let toolchain = toolchain();

	let mut failed = 0;
	for entry in cases {
		let entry = entry.map_err(|e| error!("Failed to read dir entry from \"tests/ui-tests\": {e}"))?;
		let file_type = entry.file_type()
			.map_err(|e| error!("Failed to stat \"tests/ui-tests/{}\": {e}", entry.file_name().to_string_lossy()))?;
		if file_type.is_file() {
			let path = entry.path();
			if path.extension().map(|x| x == "rs") != Some(true) {
				continue;
			}
			let Some(name) = path.file_stem() else {
				continue;
			};
			let name = name.to_str()
				.ok_or_else(|| error!("File name contains invalid UTF-8: {name:?}"))?;

			let test = TestDefinition::new(name, &dir, &toolchain)?;
			write!(std::io::stderr().lock(), "UI test {} ... ", yansi::Paint::bold(&test.name).bright()).ok();
			let result = test.run(bless)?;
			let mut stderr = std::io::stderr().lock();
			if result.diagnostics.has_error {
				writeln!(stderr, "{}", yansi::Paint::red("fail")).ok();
				failed += 1;
			} else {
				writeln!(stderr, "{}", yansi::Paint::green("ok")).ok();
			}
			result.diagnostics.print_all();
			if result.stdout_mismatch {
				if let Some(output) = test.expected_stdout.as_deref().and_then(|x| std::str::from_utf8(x).ok()) {
					writeln!(stderr, "\nExpected stdout:\n==========\n{output}==========").ok();
				}
				if let Some(output) = result.output.as_ref().and_then(|x| std::str::from_utf8(&x.stdout).ok()) {
					writeln!(stderr, "\nActual stdout:\n==========\n{output}==========").ok();
				}
			}
			if result.stderr_mismatch {
				if let Some(output) = test.expected_stderr.as_deref().and_then(|x| std::str::from_utf8(x).ok()) {
					writeln!(stderr, "\nExpected stderr:\n==========\n{output}==========").ok();
				}
				if let Some(output) = result.output.as_ref().and_then(|x| std::str::from_utf8(&x.stderr).ok()) {
					writeln!(stderr, "\nActual stderr:\n==========\n{output}==========").ok();
				}
			}
			drop(stderr);
			test.cleanup().ok();
		}
	}

	if failed == 0 {
		Ok(())
	} else {
		Err(())
	}
}

fn toolchain() -> String {
	match std::env::var("RUSTUP_TOOLCHAIN") {
		Err(_) => "unknown".into(),
		Ok(toolchain) => match toolchain.split_once("-") {
			Some((left, _right)) => left.into(),
			None => toolchain,
		}
	}
}

struct TestDefinition<'a> {
	name: &'a str,
	expected_stdout: Option<Vec<u8>>,
	expected_stderr: Option<Vec<u8>>,
	expected_stdout_path: PathBuf,
	expected_stderr_path: PathBuf,
	actual_stdout_path: PathBuf,
	actual_stderr_path: PathBuf,
	working_dir: tempfile::TempDir,
	target_dir: PathBuf,
}

struct TestResults {
	diagnostics: Diagnostics,
	output: Option<std::process::Output>,
	stdout_mismatch: bool,
	stderr_mismatch: bool,
}

struct Diagnostics {
	data: Vec<(Level, String)>,
	has_error: bool,
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
enum Level {
	Hint,
	Error,
}

impl Diagnostics {
	fn new() -> Self {
		Self {
			data: Vec::new(),
			has_error: false,
		}
	}

	fn add_hint(&mut self, hint: impl std::fmt::Display) {
		self.data.push((Level::Hint, hint.to_string()));
	}

	fn add_error(&mut self, error: impl std::fmt::Display) {
		self.data.push((Level::Error, error.to_string()));
		self.has_error = true;
	}

	fn print_all(&self) {
		let mut stderr = std::io::stderr().lock();
		for (level, message) in &self.data {
			let prefix = match level {
				Level::Error => Some(yansi::Paint::new("error").red().bright().bold()),
				Level::Hint => None,
			};
			if let Some(prefix) = prefix {
				writeln!(&mut stderr, "{prefix}: {message}").ok();
			} else {
				writeln!(&mut stderr, "{message}").ok();
			}
		}
	}
}

impl<'a> TestDefinition<'a> {
	pub fn new(name: &'a str, dir: &'a Path, toolchain: &'a str) -> Result<Self, ()> {
		let expected_stdout_path = dir.join(format!("{name}/stdout-expected-{toolchain}"));
		let expected_stderr_path = dir.join(format!("{name}/stderr-expected-{toolchain}"));
		let actual_stdout_path = dir.join(format!("{name}/stdout-actual-{toolchain}"));
		let actual_stderr_path = dir.join(format!("{name}/stderr-actual-{toolchain}"));
		let expected_stdout = read_file(&expected_stdout_path)?;
		let expected_stderr = read_file(&expected_stderr_path)?;

		let working_dir = tempfile::TempDir::new_in(CARGO_TARGET_TMPDIR)
			.map_err(|e| error!("Failed to create temporary directory for test {name:?}: {e}"))?;

		setup_test_crate(name, dir, working_dir.path())?;

		Ok(Self {
			name,
			expected_stdout,
			expected_stderr,
			expected_stdout_path,
			expected_stderr_path,
			actual_stdout_path,
			actual_stderr_path,
			working_dir,
			target_dir: Path::new(CARGO_TARGET_TMPDIR).join("target"),
		})
	}

	pub fn run(&self, bless: bool) -> Result<TestResults, ()> {
		let mut results = TestResults {
			diagnostics: Diagnostics::new(),
			output: None,
			stdout_mismatch: false,
			stderr_mismatch: false,
		};

		let build_status = match self.build_crate() {
			Ok(status) => status,
			Err(e) => {
				results.diagnostics.add_error(format!("cargo build: failed to spawn process: {e}"));
				return Ok(results);
			},
		};

		if !build_status.success() {
			results.diagnostics.add_error(format!("cargo build: {build_status}"));
			return Ok(results);
		}

		let output = match self.run_test_binary() {
			Ok(output) => results.output.insert(output),
			Err(e) => {
				results.diagnostics.add_error(format!("cargo run: failed to spawn process: {e}"));
				return Ok(results);
			},
		};

		match output.status.code() {
			None => results.diagnostics.add_error(format!("test exit status: {}", output.status)),
			Some(0) => results.diagnostics.add_error("test did not panic"),
			Some(_) => (),
		}

		results.stdout_mismatch = !check_output(
			&mut results.diagnostics,
			"stdout",
			self.expected_stdout.as_deref(),
			&output.stdout,
			bless,
			&self.expected_stdout_path,
			&self.actual_stdout_path,
		)?;
		results.stderr_mismatch = !check_output(
			&mut results.diagnostics,
			"stderr",
			self.expected_stderr.as_deref(),
			&output.stderr,
			bless,
			&self.expected_stderr_path,
			&self.actual_stderr_path,
		)?;

		Ok(results)
	}

	fn build_crate(&self) -> Result<std::process::ExitStatus, std::io::Error> {
		std::process::Command::new("cargo")
			.current_dir(self.working_dir.path())
			.args(["build", "--quiet"])
			.arg("--target-dir")
			.arg(&self.target_dir)
			.status()
	}

	fn run_test_binary(&self) -> Result<std::process::Output, std::io::Error> {
		std::process::Command::new("cargo")
			.current_dir(self.working_dir.path())
			.args(["run", "--quiet"])
			.arg("--target-dir")
			.arg(&self.target_dir)
			.env("ASSERT2", "color")
			.output()
	}

	pub fn cleanup(self) -> Result<(), ()> {
		self.working_dir.close()
			.map_err(|e| error!("Failed to clean up temporary directory: {e}"))
	}
}

fn check_output(
	diagnostics: &mut Diagnostics,
	stream_name: &str,
	expected: Option<&[u8]>,
	actual: &[u8],
	bless: bool,
	expected_path: &Path,
	actual_path: &Path,
) -> Result<bool, ()> {
	match (bless, expected) {
		(true, _) => {
			write_file(expected_path, actual)?;
			Ok(true)
		}
		(false, None) => {
			diagnostics.add_error(format!("no expected {stream_name} available, output saved to: {}", actual_path.display()));
			diagnostics.add_hint("Run the tests with UI_TESTS=bless to store the current output as expected");
			write_file(actual_path, actual)?;
			Ok(false)
		},
		(false, Some(expected)) => {
			if actual == expected {
				Ok(true)
			} else {
				diagnostics.add_error(format!("{stream_name} does not match, output saved to: {}", actual_path.display()));
				diagnostics.add_hint("Run the tests with UI_TESTS=bless to store the current output as expected");
				write_file(actual_path, actual)?;
				Ok(false)
			}
		}
	}
}

fn read_file(path: &Path) -> Result<Option<Vec<u8>>, ()> {
	let mut file = match std::fs::File::open(path) {
		Ok(x) => x,
		Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
		Err(e) => {
			error!("Failed to open {} for reading: {e}", path.display());
			return Err(());
		},
	};
	let mut buffer = Vec::new();
	std::io::Read::read_to_end(&mut file, &mut buffer)
		.map_err(|e| error!("Failed to read from {}: {e}", path.display()))?;
	Ok(Some(buffer))
}

fn write_file(path: &Path, data: &[u8]) -> Result<(), ()> {
	if let Some(parent) = path.parent() {
		std::fs::create_dir_all(parent)
		.map_err(|e| error!("Failed to create parent directory of {}: {e}", path.display()))?;
	}
	let mut file = std::fs::File::create(path)
		.map_err(|e| error!("Failed to open {} for writing: {e}", path.display()))?;
	std::io::Write::write_all(&mut file, data)
		.map_err(|e| error!("Failed to write to {}: {e}", path.display()))?;
	Ok(())
}

fn setup_test_crate(name: &str, dir: &Path, working_dir: &Path) -> Result<(), ()> {
	let assert2_path = env!("CARGO_MANIFEST_DIR");

	std::fs::copy(dir.join(format!("{name}.rs")), working_dir.join("main.rs"))
		.map_err(|e| error!("Failed to copy {}/{name}.rs to {}: {e}", dir.display(), working_dir.display()))?;
	let manifest = std::fs::OpenOptions::new()
		.write(true)
		.create_new(true)
		.open(working_dir.join("Cargo.toml"))
		.map_err(|e| error!("Failed to create manifest at {}/Cargo.toml: {e}", working_dir.display()))?;
	write_manifest(manifest, name, assert2_path)
		.map_err(|e| error!("Failed to write manifest to {}/Cargo.toml: {e}", working_dir.display()))?;
	Ok(())
}

fn write_manifest<W: std::io::Write>(mut write: W, name: &str, assert2_path: &str) -> std::io::Result<()> {
	writeln!(&mut write, "[package]")?;
	writeln!(&mut write, "name = {name:?}")?;
	writeln!(&mut write, "version = \"0.0.0\"")?;
	writeln!(&mut write, "edition = \"2021\"")?;
	writeln!(&mut write, "publish = false")?;
	writeln!(&mut write, "[[bin]]")?;
	writeln!(&mut write, "name = {name:?}")?;
	writeln!(&mut write, "path = \"main.rs\"")?;
	writeln!(&mut write, "[dependencies]")?;
	writeln!(&mut write, "assert2 = {{ path = {assert2_path:?} }}")?;
	writeln!(&mut write, "reproducible-panic = \"0.1.2\"")?;
	writeln!(&mut write, "[workspace]")?;

	Ok(())
}
