/* 
 * Copyright (c) 2024 Emma Tebibyte <emma@tebibyte.media>
 *
 * Copying and distribution of this file, with or without modification, are
 * permitted in any medium without royalty provided the copyright notice and
 * this notice are preserved. This file is offered as-is, without any warranty.
 */

/* https://exercism.org/tracks/rust/exercises/protein-translation */

use std::{
	convert::TryFrom,
	env::args,
	fmt::{ Display, Error, Formatter },
	io,
	vec::IntoIter,
};

use Codon::*;

/* codon type */
#[derive(Clone, Debug)]
enum Codon {
	Methionine,
	Phenylalanine,
	Leucine,
	Serine,
	Tyrosine,
	Cysteine,
	Tryptophan,
	STOP,
}

/* impls ToString for Codon */
impl Display for Codon {
	fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
		match self {
			Methionine => write!(f, "Methionine"),
			Phenylalanine => write!(f, "Phenylalanine"),
			Leucine => write!(f, "Leucine"),
			Serine => write!(f, "Serine" ),
			Tyrosine => write!(f, "Tyrosine"),
			Cysteine => write!(f, "Cysteine"),
			Tryptophan => write!(f, "Tryptophan"),
			STOP => write!(f, "STOP"),
		}
	}
}

/* try to convert a String to a Codon and return error message otherwise */
impl TryFrom<String> for Codon {
	type Error = String;

	fn try_from(value: String) -> Result<Self, Self::Error> {
		match value.as_str() {
			"AUG" => Ok(Methionine),
			"UUU" | "UUC" => Ok(Phenylalanine),
			"UUA" | "UUG" => Ok(Leucine),
			"UCU" | "UCC" | "UCA" | "UCG" => Ok(Serine),
			"UAU" | "UAC" => Ok(Tyrosine),
			"UGU" | "UGC" => Ok(Cysteine),
			"UGG" => Ok(Tryptophan),
			"UAA" | "UAG" | "UGA" => Ok(STOP),
			_ => Err(format!("{}: unknown codon", value)),
		}
	}
}

/* protein type */
#[derive(Clone, Debug)]
struct Protein(Vec<Codon>);

/* convenience functions for constructing and appending to Proteins */
impl Protein {
	fn new() -> Self { Protein(Vec::new()) }
	fn push(&mut self, codon: Codon) { self.0.push(codon); }
}

/* implements to_string() for Protein */
impl Display for Protein {
	fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
		write!(f, "{}", self
			.clone()
			.into_iter()
			.map(|x| x.to_string())
			.collect::<Vec<_>>()
			.join(" "))
	}
}

/* iterator for Proteins that yields Codons */
impl IntoIterator for Protein {
	type Item = Codon;
	type IntoIter = IntoIter<Self::Item>;
	fn into_iter(self) -> Self::IntoIter { self.0.into_iter() }
}

/* try to convert a String to a Protein using Codon::try_from() */
impl TryFrom<String> for Protein {
	type Error = String;

	fn try_from(value: String) -> Result<Self, Self::Error> {
		let mut protein = Protein::new();
		let codons = value.chars().collect::<Vec<_>>();

		for codon in codons.chunks(3) {
			match Codon::try_from(codon.iter().collect::<String>())? {
				STOP => break,
				c => protein.push(c),
			};
		}

		Ok(protein)
	}
}

fn main() -> io::Result<()> {
	for arg in args().skip(1) {
		println!("{}", Protein::try_from(arg).map_err(|x| io::Error::other(x))?);
	}

	Ok(())
}

#[cfg(test)]
mod tests {
	use std::convert::TryFrom;
	use Protein;

	#[test]
	fn examples() {
		let sequence = String::from("AUGUUUUCUUAAAUG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Methionine Phenylalanine Serine");
	}

	#[test]
	fn methionine() {
		let sequence = String::from("AUG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Methionine");
	}

	#[test]
	fn phenylalanine() {
		let sequence = String::from("UUU");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Phenylalanine");
		
		let sequence = String::from("UUC");
		let protein = Protein::try_from(sequence).unwrap().to_string();
		assert_eq!(protein, "Phenylalanine");
	}

	#[test]
	fn leucine() {
		let sequence = String::from("UUA");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Leucine");

		let sequence = String::from("UUG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Leucine");
	}

	#[test]
	fn serine() {
		let sequence = String::from("UCU");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Serine");

		let sequence = String::from("UCC");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Serine");

		let sequence = String::from("UCA");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Serine");

		let sequence = String::from("UCG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Serine");
	}

	#[test]
	fn tyrosine() {
		let sequence = String::from("UAU");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Tyrosine");

		let sequence = String::from("UAC");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Tyrosine");
	}

	#[test]
	fn cysteine() {
		let sequence = String::from("UGU");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Cysteine");

		let sequence = String::from("UGC");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Cysteine");
	}

	#[test]
	fn tryptophan() {
		let sequence = String::from("UGG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "Tryptophan");
	}

	#[test]
	fn stop() {
		let sequence = String::from("UAA");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "");

		let sequence = String::from("UAG");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "");

		let sequence = String::from("UGA");
		let protein = Protein::try_from(sequence).unwrap().to_string();

		assert_eq!(protein, "");
	}
}