use crate::{
	config::rezon_config::RezonConfig,
	error::Error::{self, *},
	tui::{clean_console, indent},
	util::{get_and_display_model_names, terminal_cleanup, BASE_BFF_URL},
};
use async_trait::async_trait;
use clap::{Arg, ArgAction, ArgMatches, Command};
use colorful::{Color, Colorful};
use csv::Reader;
use regex::Regex;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use spinach::Spinach;
use std::{collections::HashMap, fs::File, time::Instant, vec::Vec};
use uuid::Uuid;
use reqwest_middleware::{ClientBuilder, Result as ReqwestMiddlewareResult};
use reqwest_retry::{RetryTransientMiddleware, policies::ExponentialBackoff, RetryableStrategy, Retryable};

use super::RezonCommand;

pub struct Run {}

#[derive(Serialize, Deserialize, Debug)]
pub struct RunBody {
	pub input: Value,
	pub include_explanations: bool,
	pub model_name: String,
	pub parameters: Vec<String>,
}

#[derive(Debug, Deserialize, Serialize)]
// TODO: once we have more defined structures for explanations we can use this...
pub struct Explanation {
	input_values: HashMap<String, Value>,
	input_contributions: Vec<Vec<f64>>,
	input_dependencies: Vec<Vec<f64>>,
	independent_term: Vec<Vec<f64>>,
	neural_network_output: Vec<f64>,
	linear_function_output: Vec<f64>,
	text_explanations: Vec<Vec<String>>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Model {
	pub id: i32,
	pub created_at: String,
	pub updated_at: String,
	pub build_id: Uuid,
	pub name: String,
	pub user_id: String,
	pub storage_id: Uuid,
	pub source_url: String,
	pub deploy_url: String,
	pub api_keys_id: Option<i32>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct ApiResponse {
	output: Value,
	explanation: Value,
}

#[derive(Debug, Deserialize, Serialize)]
struct ModelServerError {
	#[serde(rename = "type")]
	err_type: String,
	message: String,
	code: i32
}

#[derive(Debug, Deserialize, Serialize)]
struct ModelServerErrorBody {
	error: ModelServerError
}


#[async_trait]
impl RezonCommand for Run {
	fn cmd() -> clap::Command {
		Command::new("run")
			.about("Invoke machine learning models hosted on Rezon cloud")
			.arg(
				Arg::new("input_values")
					.required(true)
					.conflicts_with("input_file")
					.short('i')
					.long("input")
					.short_aliases(['d', 'v'])
					.num_args(0..=10000)
					.action(ArgAction::Set)
					.value_name("INPUT_VALUES")
					.help("Nested array of comma-separated input values"),
			)
			.arg(
				Arg::new("input_file")
					.required(true)
					.conflicts_with("input_values")
					.short('f')
					.long("input_file")
					.num_args(1)
					.action(ArgAction::Set)
					.value_name("INPUT_FILE_PATH")
					.help("Path to a CSV file of input values"),
			)
			.arg(
				Arg::new("model_name")
					.short('m')
					.long("model")
					.action(ArgAction::Set)
					.value_name("MODEL_NAME")
					.required(false)
					.help("Name of the model to invoke"),
			)
			.arg(
				Arg::new("no_explain")
					.long("no-explain")
					.required(false)
					.action(ArgAction::SetTrue)
					.value_name("NO_EXPLAIN")
					.help("Run your model without explanations"),
			)
	}

	async fn execute(args: &ArgMatches) -> Result<(), Error> {
		run_model(RezonConfig::read_config()?, args).await?;
		Ok(())
	}
}

struct ModelInvokeRetry;

impl RetryableStrategy for ModelInvokeRetry {
	fn handle(&self, res: &ReqwestMiddlewareResult<reqwest::Response>) -> Option<Retryable> {
		match res {
			// Retry only if we got back a 401 (this is the error that gets thrown from GCP when the server first starts up)
			Ok(success) if success.status() == 401 => Some(Retryable::Transient),
			// otherwise do not retry a successful request or any other error
			Ok(_) => None,
			Err(_) => None
	   }
   }
}

async fn run_model(config: RezonConfig, args: &ArgMatches) -> Result<(), Error> {
	let mut input_values_str: String;
	if args.contains_id("input_file") {
		input_values_str = get_file_input(&args)?
	} else {
		// Collect json string as user input
		// This unwrap is safe because arg is required.
		input_values_str = args.get_one::<String>("input_values").unwrap().to_owned();
	}

	// Check if input_values_str is at 2D json array of the form: [ any amount of whitespace [...
	let regex = Regex::new(r"\[\s*\[").unwrap();
	// If input_values_str is not a 2D array, make it a 2d array
	if !regex.is_match(input_values_str.as_str()) {
		input_values_str = format!("[{input_values_str}]");
	}

	let input_values = serde_json::from_str(&input_values_str)?;
	// This will always be Some() as due to clap default_value
	let model_name = match args.get_one::<String>("model_name") {
		Some(name) => name.to_owned(),
		None => {
			// No model name provided, fetch list of available models and return
			get_and_display_model_names(&config).await?
		}
	};

	let retry_policy = ExponentialBackoff::builder().build_with_total_retry_duration(std::time::Duration::from_secs(45));
	let retry_strategy = RetryTransientMiddleware::new_with_policy_and_strategy(
        retry_policy,
        ModelInvokeRetry,
    );

	let client = ClientBuilder::new(Client::new()).with(retry_strategy).build();

	// After doing any necessary model name selection, lets run the model
	clean_console();

	// Show a loading spinner while we wait for the model to run...
	let loading = Spinach::new("Invoking model...");
	let timer = Instant::now();

	let response = client
		.post(format!("{}/models/run", BASE_BFF_URL))
		.header("Authorization", config.config.api_key.clone())
		.header("Content-Type", "application/json")
		.json(&RunBody {
			input: input_values,
			include_explanations: !args.get_flag("no_explain"),
			model_name: model_name.clone(),
			parameters: config.model.parameters.clone(),
		})
		.send()
		.await?;


	if response.status() != StatusCode::OK {
		terminal_cleanup(loading);

		let json = response.json::<Value>().await?;

		// Check to see if we can parse the response and pull the `message` key out
		// The `message` key gives us the most descriptive error message from the server
		if let Some(val) = json.get("message") {
			return Err(CliError(val.to_string()));
		}

		return Err(NotOk {
			response: serde_json::to_string(&json)?,
		});
	}

	let json_response = response.json::<Value>().await?;

	// Try and deserialize the response body as a ServerError that gets returned from the model server
	match serde_json::from_value::<ModelServerErrorBody>(json_response.clone()) {
		Ok(body) => {
			return Err(ModelServerErrorResponse {
				response: body.error.message
			});
		},
		Err(_) => {} // Do nothing if we were not able to deserialize the response as a ServerError
	}

	let response_body: Vec<ApiResponse> = serde_json::from_value(json_response.clone())?;
	let formatted_response_body = serde_json::to_string_pretty(&response_body)?;

	terminal_cleanup(loading);

	println!("{formatted_response_body}");
	println!("{}", indent(1));
	println!(
		"{} Model result computed in {}",
		"➜".color(Color::Purple1a),
		format!("{:?}", timer.elapsed()).color(Color::Purple1a)
	);

	return Ok(());
}

/// Process the given csv file for serde to interpret
fn get_file_input(args: &ArgMatches) -> Result<String, Error> {
	// Open file
	// This unwrap is safe because arg is required
	let file_path = args.get_one::<String>("input_file").unwrap();
	let file = match File::open(file_path) {
		Ok(file) => file,
		Err(e) => {
			return Err(Io {
				path: std::path::PathBuf::from(file_path),
				source: e,
			})
		}
	};

	// Convert CSV input to JSON string
	let mut temp_values_str = String::from("[");
	for result in Reader::from_reader(file).records() {
		let record = result?;
		temp_values_str.push('[');
		for field in record.iter() {
			temp_values_str.push_str(field);
			temp_values_str.push_str(", ");
		}
		temp_values_str.pop();
		temp_values_str.pop();
		temp_values_str.push_str("], ");
	}
	temp_values_str.pop();
	temp_values_str.pop();
	temp_values_str.push(']');
	Ok(temp_values_str.to_string())
}
