pipeline.rs

use clap::Parser;
use std::collections::HashMap;
use std::path::{Path, PathBuf};

use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::cli_args::Args;
use crate::commands::{Palette, Remap};
use crate::format::make_paths;
use crate::load_image;
use crate::utils::normalize_path;

#[derive(Error, Debug)]
pub enum PipelineError {
	#[error("Use a file ending with '.toml' or '.json' to configure your pipeline")]
	FormatDetection,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PipelineType {
	Pipeline {
		input_path: String,
		output_path: String,
		actions: Vec<Args>,
	},
	Ref {
		input_path: String,
		output_path: String,
		reference: String,
	},
	GlobRef {
		pattern: String,
		output_dir: String,
		reference: String,
	},
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PipelineRef {
	pub actions: Vec<Args>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PipelineFile {
	pub refs: HashMap<String, PipelineRef>,
	pub pipelines: Vec<PipelineType>,
}

/// Execute a predefined pipeline
#[derive(Debug, Clone, Parser, Serialize, Deserialize)]
#[clap(author, version = "0.3.0")]
pub struct Pipeline {
	/// The path to the pipeline definition file
	#[serde(default)]
	pub config: String,
}

macro_rules! result {
	($value: expr) => {
		match $value {
			Ok(val) => val,
			Err(e) => {
				log::error!("{}", e);
				return;
			}
		}
	};
}

impl Pipeline {
	pub fn execute(&self) -> anyhow::Result<()> {
		let path = std::env::current_dir().map(|path| path.join(&self.config))?;
		let path_string = format!("{}", &path.display());

		log::debug!("Trying to read from input file: {}", &path.display());

		if !&path_string.ends_with(".toml") && !&path_string.ends_with(".json") {
			Err(PipelineError::FormatDetection)?;
		}

		let file_contents = std::fs::read(&path)?;

		log::debug!("Found correct file type and read bytes, trying to parse");
		let pipeline_data: PipelineFile = if path_string.ends_with(".toml") {
			toml::from_str(String::from_utf8(file_contents)?.as_str())?
		} else {
			serde_json::from_slice(&file_contents)?
		};

		log::debug!("Expanding pipeline file into targets");
		let base_path = PathBuf::from(path.parent().unwrap());
		get_targets(base_path.clone(), &pipeline_data).for_each(
			|(input_path, output_path, actions)| {
				match make_paths(&output_path) {
					Ok(_) => {}
					Err(e) => {
						log::error!("Failed to create target directory {}; {}", &output_path, e);
						return;
					}
				}

				if actions.is_empty() {
					match std::fs::copy(&input_path, &output_path) {
						Ok(_) => {}
						Err(e) => {
							log::error!("Failed to copy {} to {}; {}", input_path, output_path, e);
						}
					};
					return;
				}

				let mut file = result!(load_image(&input_path, None));

				log::debug!(
					"Loaded {}, Executing {} actions",
					&input_path,
					actions.len()
				);

				let mut count = 1;
				for step in actions {
					match step {
						Args::Rotate(rotate) => {
							file = result!(rotate.run(&file));
						}
						Args::Extrude(extrude) => {
							file = result!(extrude.run(&file));
						}
						Args::Scale(scale) => {
							file = result!(scale.run(&file));
						}
						Args::Flip(flip) => {
							file = result!(flip.run(&file));
						}
						Args::Remap(remap) => {
							let palette = result!(load_image(&remap.palette, None));
							let image_palette = result!(Palette::extract_from(&file));
							let target_palette = result!(Palette::extract_from(&palette));

							let mappings =
								Palette::calculate_mapping(&image_palette, &target_palette);
							file = result!(Remap::remap_image(file, mappings));
						}
						_ => {}
					}

					count += 1;
				}

				let mut outer_target_path = PathBuf::from(&output_path);
				outer_target_path.pop();

				if let Err(e) = std::fs::create_dir(&outer_target_path) {
					match e.kind() {
						std::io::ErrorKind::AlreadyExists => { /* This is fine */ }
						_ => log::error!(
							"Failed to create containing directory {}; {}",
							outer_target_path.to_string_lossy(),
							e
						),
					}
				}

				match file.save(&output_path) {
					Ok(_) => {}
					Err(e) => {
						log::error!("Failed to save to {}; {}", output_path, e);
					}
				}
			},
		);

		Ok(())
	}
}

fn join<T: AsRef<Path>>(root: &Path, rest: &T) -> String {
	let path = normalize_path(root.join(rest));
	format!("{}", path.display())
}

fn get_targets(
	base_path: PathBuf,
	pipeline_data: &PipelineFile,
) -> impl ParallelIterator<Item = (String, String, Vec<Args>)> + '_ {
	pipeline_data
		.pipelines
		.par_iter()
		.flat_map(move |pipe| match pipe {
			PipelineType::Pipeline {
				input_path,
				output_path,
				actions,
			} => vec![(
				join(&base_path, &input_path),
				join(&base_path, &output_path),
				actions.clone(),
			)],
			PipelineType::Ref {
				input_path,
				output_path,
				reference,
			} => pipeline_data
				.refs
				.get(reference.as_str())
				.iter()
				.map(|value| {
					(
						join(&base_path, &input_path),
						join(&base_path, &output_path),
						(*value).actions.clone(),
					)
				})
				.collect(),
			PipelineType::GlobRef {
				pattern,
				output_dir,
				reference,
			} => pipeline_data
				.refs
				.get(reference.as_str())
				.iter()
				.map(|value| (*value).actions.clone())
				.flat_map(|actions| {
					let mut paths = Vec::new();
					let target_path = join(&base_path, pattern);

					log::debug!("Mapping glob paths for '{}'", &target_path);
					for entry in glob::glob(target_path.as_str()).unwrap() {
						log::debug!("Found a glob match: [{:?}]", entry);
						paths.push((actions.clone(), entry));
					}
					paths
				})
				.filter_map(|(actions, inner)| inner.ok().map(|p| (actions, p)))
				.filter_map(|(actions, path)| {
					if let Some(filename) = path.file_name().and_then(|osstr| osstr.to_str()) {
						let output_path = Path::new(output_dir.as_str());
						let output_path = output_path.join(filename);

						Some((
							join(&base_path, &path),
							join(&base_path, &output_path),
							actions,
						))
					} else {
						None
					}
				})
				.collect(),
		})
}