use super::RezonCommand;
use crate::{cli::CliError, config::rezon_config::RezonConfig, tui::indent};
use async_trait::async_trait;
use chrono::NaiveDateTime;
use clap::{arg, Arg, ArgAction, ArgMatches, Command};
use colorful::{Color, Colorful};
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
	collections::HashMap,
	fs::{write, File},
	process::exit,
	time::Instant,
};

/**
 * rezon run --input "[[0.569, 0.175, 0.175, 0.175, 0.175, 0.37], [0.569, 0.98, 0.2345, 0.5674, 0.11, 0.37],[0.569, 0.889, 0.111, 0.56, 0.17, 0.42]]"
 * Note: you can include explainations with the `--explain` flag/arg
 * Note: the output can be formatted as json with the `--json` flag/arg
 */
pub struct Run {}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RezonInput<T> {
	pub input: T,
	pub include_explanations: bool,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RezonOutput {
	output: Value,
	explanation: Option<RezonExplanation>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RezonOutputWithDate {
	created_at: NaiveDateTime,
	input: Value,
	output: Value,
	explanation: Option<RezonExplanation>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RezonExplanation {
	input_values: Value,
	input_contributions: Value,
	input_dependencies: Value,
	independent_term: Value,
	neural_network_output: Value,
	linear_function_output: Value,
	text_explanations: Value,
	reference_id: String,
}

#[async_trait]
impl RezonCommand for Run {
	fn cmd() -> clap::Command {
		Command::new("run")
			.about("Invoke machine learning models hosted via Rezon cloud")
			.arg(
				Arg::new("input_values")
					.short('i')
					.long("input")
					.short_aliases(['d', 'v'])
					.num_args(0..=10000)
					.action(clap::ArgAction::Set)
					.help("Comma-separated input values")
					.required(true),
			)
			.arg(
				Arg::new("explain")
					.short('e')
					.long("explain")
					.help("Include explanations for your model's output")
					.required(false)
					.action(ArgAction::SetTrue),
			)
			.arg(
				Arg::new("json")
					.short('j')
					.long("json")
					.help("Run your model and receive the result as json")
					.required(false)
					.action(ArgAction::SetTrue),
			)
	}

	async fn execute(rezon_config: Option<RezonConfig>, args: ArgMatches) -> Result<(), CliError<'static>> {
		// if let Some(rezon_config) = rezon_config {
		// 	let model_params = rezon_config.model.parameters;
		// } else {
		// 	println!("💡 Name your model's input and output parameters in your Rezon.toml file for better explainability!")
		// }

		// TODO: instead of grabbing one we want to instead grab all of the inputs and loop through them (we should be able to just put everything inside of a for loop)
		// Use this snippet: input_values.into_iter().collect::<Vec<&String>>() -- and also change `get_one` to `get_many`
		let input_values = args
			.get_many::<String>("input_values")
			.expect("Input values are required")
			.into_iter()
			.collect::<Vec<&String>>();

		for (index, input) in input_values.iter().enumerate() {
			let sanitized_input: Value = match serde_json::from_str(input) {
				Ok(val) => val,
				Err(_) => {
					println!(
						"{}",
						"Invalid input. Make sure the input sent to the rezen CLI is valid JSON.".color(Color::Red)
					);
					exit(64);
				}
			};

			let should_explain = args.get_one::<bool>("explain").expect("Could not parse `explain` input arg!");
			let is_json = args.get_one::<bool>("json").expect("Could not parse `json` input arg!");

			// TODO: `should_explain` should not be used in an if statement here -- we should just directly pass it to the reqwest method
			// ^^ not an issue for now -- just pointing it out
			if !should_explain {
				let client = Client::new();

				// Sample input vals for the yacht example
				let _ = vec![
					[0.569, 0.175, 0.175, 0.175, 0.175, 0.37],
					[0.569, 0.98, 0.2345, 0.5674, 0.11, 0.37],
					[0.569, 0.889, 0.111, 0.56, 0.17, 0.42],
				];

				let start = Instant::now();

				let response = match client
					.post("http://localhost:8000") // TODO: this URL will get passed in
					.json(&RezonInput {
						input: sanitized_input,
						include_explanations: false,
					})
					.send()
					.await
				{
					Ok(res) => {
						if !(res.status() == StatusCode::OK) {
							println!(
								"{}",
								"Error: could not invoke model. Are your inputs formatted correctly?".color(Color::Red)
							);
							exit(64);
						}
						res.json::<Value>()
					}
					Err(_) => {
						println!(
							"{}",
							"Error: could not invoke model. Are your inputs formatted correctly?".color(Color::Red)
						);
						exit(64);
					}
				};

				// --------- WE NEED TO CACHE ALL OF OUR RESPONSES INTO A JSON FILE ---------
				// Note: in the real thing this will actually save to a database etc
				let curr_date = chrono::Local::now().naive_local();
				let file_path = "/Users/JRDarrenbaldwin/rezon/reports.json";

				// Check if a file exists at /rezon/reports.json
				match File::open(&file_path) {
					Ok(f) => f,
					Err(_) => {
						// If it does not exist we want to create one
						let f = File::create(&file_path).unwrap();
						f
					}
				};

				// Read the file...
				let contents = std::fs::read_to_string(&file_path).unwrap();

				// Output the final response
				let parsed_resp = response.await.unwrap();

				let sanitized_resp: RezonOutput = serde_json::from_value(parsed_resp.clone()).unwrap();

				// Update the file with the new data
				let new_entry = RezonOutputWithDate {
					created_at: curr_date,
					input: serde_json::from_str(&input).unwrap(),
					output: sanitized_resp.output.clone(),
					explanation: None,
				};

				// Check if there is nothing in the file...
				if contents == String::from("") {
					let entry = vec![new_entry];
					write(file_path, serde_json::to_string_pretty(&entry).unwrap()).unwrap();
				} else {
					// Since there was not nothing in the file we want to simply append to the array and re-write to the file
					// Parse the contents of the file to a type
					let mut sanitized_contents: Vec<RezonOutputWithDate> = serde_json::from_str(&contents).unwrap();
					sanitized_contents.push(new_entry);

					write(file_path, serde_json::to_string_pretty(&sanitized_contents).unwrap()).unwrap();
				}

				let duration = start.elapsed().as_millis();

				println!("{}", indent(1));

				// If the user passed in more than 1 input we want to print a divider (the terminal gets confusing if we dont)
				if &input_values.len() > &1 {
					println!(
						"{} {}",
						format!("{}", "➜".color(Color::Purple1a)).bold(),
						format!(
							"{}",
							format!(
								"───────────── Input {}: Model result computed in: {} ─────────────",
								format!("#{}", index + 1).color(Color::Purple1a),
								format!("{}ms", duration).color(Color::Purple1a).bold()
							)
						)
						.bold()
					);
				} else {
					println!(
						"{} {} {}:",
						format!("{}", "➜".color(Color::Purple1a)).bold(),
						format!("{}", "Model result computed in".bold()),
						format!("{}ms", duration).color(Color::Purple1a).bold()
					);
				}

				println!("{}", indent(1));

				if is_json == &true {
					println!("{}", serde_json::to_string_pretty(&parsed_resp).unwrap());
				} else {
					println!("{}", serde_json::to_string_pretty(&sanitized_resp.output).unwrap());
				}

				continue;
			} else {
				// TODO: this is duplicated code but no need to refactor for demo
				let client = Client::new();

				// Sample input vals for the yacht example
				let _ = vec![
					[0.569, 0.175, 0.175, 0.175, 0.175, 0.37],
					[0.569, 0.98, 0.2345, 0.5674, 0.11, 0.37],
					[0.569, 0.889, 0.111, 0.56, 0.17, 0.42],
				];

				let start = Instant::now();

				let response = match client
					.post("http://localhost:8000") // TODO: this URL will get passed in
					.json(&RezonInput {
						input: sanitized_input,
						include_explanations: true,
					})
					.send()
					.await
				{
					Ok(res) => {
						if !(res.status() == StatusCode::OK) {
							println!(
								"{}",
								"Error: could not invoke model. Are your inputs formatted correctly?".color(Color::Red)
							);
							exit(64);
						}
						res.json::<Value>()
					}
					Err(_) => {
						println!(
							"{}",
							"Error: could not invoke model. Are your inputs formatted correctly?".color(Color::Red)
						);
						exit(64);
					}
				};

				// --------- WE NEED TO CACHE ALL OF OUR RESPONSES INTO A JSON FILE ---------
				// Note: in the real thing this will actually save to a database etc
				let curr_date = chrono::Local::now().naive_local();
				let file_path = "/Users/JRDarrenbaldwin/rezon/reports.json";

				// Check if a file exists at /rezon/reports.json
				match File::open(&file_path) {
					Ok(f) => f,
					Err(_) => {
						// If it does not exist we want to create one
						let f = File::create(&file_path).unwrap();
						f
					}
				};

				// Read from the reports.json file...
				let contents = std::fs::read_to_string(&file_path).unwrap();

				// Output the final response
				let parsed_resp = response.await.unwrap();

				let sanitized_resp: RezonOutput = serde_json::from_value(parsed_resp.clone()).unwrap();

				// Update each text explanation with the param alias from the rezon.toml
				// We want to update the param names with the ones specified inside of the rezon.toml
				let mut explanations: HashMap<String, Vec<String>> =
					serde_json::from_value(sanitized_resp.explanation.clone().unwrap().text_explanations).unwrap();

				if let Some(rezon_config) = rezon_config.clone() {
					let model_params = rezon_config.model.parameters;
					if !model_params.is_empty() {
						for key in explanations.clone().keys() {
							let mut new_name = explanations.get(key).unwrap().to_owned();
							let mut name_index = 1;
							for param_name in model_params.clone() {
								let replace_string = format!("Variable {}", name_index);
								new_name = vec![new_name[0].replace(&replace_string, &param_name)];
								name_index += 1;
							}
							explanations.insert(key.to_string(), new_name);
						}
					}
				}

				let new_resp = sanitized_resp.clone();

				// Update the file with the new data
				let final_resp = RezonOutputWithDate {
					created_at: curr_date,
					input: serde_json::from_str(&input).unwrap(),
					output: sanitized_resp.output.clone(),
					explanation: Some(RezonExplanation {
						text_explanations: serde_json::to_value(explanations).unwrap(),
						..new_resp.explanation.unwrap()
					}),
				};

				// Check if there is nothing in the file...
				if contents == String::from("") {
					let entry = vec![final_resp.clone()];
					write(file_path, serde_json::to_string_pretty(&entry).unwrap()).unwrap();
				} else {
					// Since there was not nothing in the file we want to simply append to the array and re-write to the file
					// Parse the contents of the file to a type
					let mut sanitized_contents: Vec<RezonOutputWithDate> = serde_json::from_str(&contents).unwrap();
					sanitized_contents.push(final_resp.clone());

					write(file_path, serde_json::to_string_pretty(&sanitized_contents).unwrap()).unwrap();
				}

				// stop our benchmark
				let duration = start.elapsed().as_millis();

				println!("{}", indent(1));

				if &input_values.len() > &1 {
					println!(
						"{} {}",
						format!("{}", "➜".color(Color::Purple1a)).bold(),
						format!(
							"{}",
							format!(
								"───────────── Input {}: Model result computed in: {} ─────────────",
								format!("#{}", index + 1).color(Color::Purple1a),
								format!("{}ms", duration).color(Color::Purple1a).bold()
							)
						)
						.bold()
					);
				} else {
					println!(
						"{} {} {}:",
						format!("{}", "➜".color(Color::Purple1a)).bold(),
						format!("{}", "Model result computed in".bold()),
						format!("{}ms", duration).color(Color::Purple1a).bold()
					);
				}

				println!("{}", indent(1));

				if is_json == &true {
					println!("{}", serde_json::to_string_pretty(&final_resp).unwrap());
				} else {
					// We want to update the param names with the ones specified inside of the rezon.toml
					let mut explanations: HashMap<String, Vec<String>> =
						serde_json::from_value(sanitized_resp.explanation.unwrap().text_explanations).unwrap();

					if let Some(rezon_config) = rezon_config.clone() {
						let model_params = rezon_config.model.parameters;
						if !model_params.is_empty() {
							for key in explanations.clone().keys() {
								let mut new_name = explanations.get(key).unwrap().to_owned();
								let mut name_index = 1;
								for param_name in model_params.clone() {
									let replace_string = format!("Variable {}", name_index);
									new_name = vec![new_name[0].replace(&replace_string, &param_name)];
									name_index += 1;
								}
								explanations.insert(key.to_string(), new_name);
							}
						}
					}

					println!("{}", serde_json::to_string_pretty(&sanitized_resp.output).unwrap());

					println!("{}", indent(1));

					println!("{}", "Explanations:".color(Color::Purple1a));
					println!("{}", indent(1));

					println!("{} {}", "➜".color(Color::Purple1a).bold(), explanations.get("0").unwrap()[0]);
				}
			}
		}
		Ok(())
	}
}
