#![allow(non_snake_case)]
use super::RezonCommand;
use crate::{
	cli::CliError,
	config::rezon_config::RezonConfig,
	tui::{clean_console, indent, SPINNER_COLOR},
	util::{is_build_completed, BASE_BFF_URL},
};
use async_trait::async_trait;
use clap::{ArgMatches, Command};
use colorful::{Color, Colorful};
use requestty::{prompt_one, Question};
use reqwest::{
	multipart::{Form, Part},
	Client,
};
use serde::{Deserialize, Serialize};
use spinoff::{spinners, Spinner};
use std::{error::Error as StdError, fmt, fs::File, io::Read, thread, time::Duration};
use uuid::Uuid;

/**
 * `rezon deploy`
 * > Deploys a model by searching for a valid `h5` and .json file in the current project
 */
pub struct Deploy {}

// TODO: We need a shared crate/package for some of these types to prevent duplication
#[derive(Serialize, Deserialize)]
pub struct UploadResponse {
	pub model_id: Uuid,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct StorageSource {
	pub bucket: String,
	pub object: String,
	pub generation: String,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct Source {
	pub storageSource: StorageSource,
}

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

#[derive(Serialize, Deserialize)]
struct DeployInput {
	pub model_id: uuid::Uuid,
	pub model_name: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct DeployResponse {
	pub status: String,
	pub build_id: Uuid,
	pub source: Source,
	pub model: InputModel,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct DeployStatusResponse {
	pub status: String,
	pub build_create_time: String,
}

#[derive(Debug, Clone)]
pub struct DeployCommandError {
	pub message: String,
	pub exit_code: i32,
}

impl StdError for DeployCommandError {}

impl fmt::Display for DeployCommandError {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		write!(f, "Error: could not upload model.zip to Google Cloud Storage!")
	}
}

#[async_trait]
impl RezonCommand for Deploy {
	fn cmd() -> clap::Command {
		Command::new("deploy").about("Deploy a valid model to Rezon cloud!")
	}

	async fn execute(config: Option<RezonConfig>, _: ArgMatches) -> Result<(), crate::cli::CliError<'static>> {
		deploy_model(config.unwrap()).await
	}
}

pub async fn deploy_model(config: RezonConfig) -> Result<(), CliError<'static>> {
	// (1) Grab references to our api_key and model paths, then open the files
	let api_key = config.config.api_key;
	let h5_file = config.model.h5_file;
	let json_path = config.model.json_path;

	// Open our files..
	let h5 = match File::open(h5_file) {
		Ok(file) => file,
		Err(_) => {
			return Err(CliError::new("Error: Failed to open the model's weights (.h5)", 64));
		}
	};

	let json = match File::open(json_path) {
		Ok(file) => file,
		Err(_) => {
			return Err(CliError::new("Error: Failed to open the model's configuration file (.json)", 64));
		}
	};

	let model_files = vec![h5, json];

	// (2) Iterate over the files and add them to the form data in the same request
	let upload_form_data = match chunk_model_files(model_files) {
		Ok(form_data) => form_data,
		Err(_) => {
			return Err(CliError::new("Error: Failed to chunk model files into form data for upload!", 64));
		}
	};

	// (3) Ask the user for the name of their model
	let model_name = prompt_one(
		Question::input("model_name")
			.message("Choose a name for your model:")
			.default("my-model")
			.build(),
	)
	.expect("Error: Failed to Deploy model to Rezon cloud. Please try again!");

	let model_name = model_name.as_string().unwrap();

	// Validate that the project name does not contain any invalid chars
	if !model_name.chars().all(|x| x.is_alphanumeric() || x == '-' || x == '_') {
		return Err(CliError::new(
			"Aborting...your model name may only contain alphanumeric characters along with '-' and '_'...",
			64,
		));
	}

	// (4) Hit the `/models/upload endpoint and upload the model files to Google Cloud storage
	let http_client = reqwest::Client::new();
	let upload_endpoint = format!("{BASE_BFF_URL}/models/upload");

	clean_console();

	// Show a loading progress bar while we wait for the upload, build, and deployment to complete
	let mut spinner = Spinner::new(spinners::Aesthetic, "Uploading model...", SPINNER_COLOR);

	let model_upload_req = match trigger_model_upload(http_client.clone(), upload_endpoint, upload_form_data, api_key.clone()).await {
		Ok(req) => req,
		Err(e) => {
			// Before outputting an error, lets stop the spinner and clear the console
			spinner.stop();
			clean_console();
			return Err(e);
		}
	};

	// (5) Using the `model_id` from the upload request, lets
	// hit the `/models/deploy` endpoint and send the model to Google Cloud Build
	let deploy_endpoint = format!("{BASE_BFF_URL}/models/deploy");

	let deploy_body = DeployInput {
		model_id: model_upload_req.model_id,
		model_name: model_name.to_string(),
	};

	// Update the spinner to have the deployment text relative to the current task (deploying the actual model)
	spinner.update_text(format!("Deploying model to {} Cloud...", "Rezon".color(Color::Purple1a)));

	let model_deploy_req = match trigger_model_deployment(http_client.clone(), deploy_endpoint, deploy_body, api_key).await {
		Ok(req) => req,
		Err(e) => {
			// Before outputting an error, lets stop the spinner and clear the console
			spinner.stop();
			clean_console();
			return Err(e);
		}
	};

	// (6) Load until the status is equal to "SUCCESS"
	while !is_build_completed(http_client.clone(), model_deploy_req.build_id).await {
		// Wait to request the build status every 45 seconds to avoid a rate limit by GCP
		thread::sleep(Duration::from_secs(45));
	}

	// Now that the model deployment is complete lets stop the loader
	spinner.stop();

	deployment_confirmation(model_deploy_req, model_name);

	Ok(())
}

fn deployment_confirmation(model_deploy_req: DeployResponse, model_name: &str) {
	println!("{}", indent(1));

	// Clear the console
	clean_console();

	// (6) Finally, show confirmation logs with the newly created model deploy Url and the user can run the model
	println!("{} Deployment Complete!\n", "✓".color(Color::Green));
	println!(
		"Your {} model is now live and ready to serve traffic!",
		model_name.color(Color::Purple1a).bold()
	);

	println!(
		"\n{} {}",
		format!("{}", "➜".color(Color::Purple1a)),
		format!(
			"{} {}",
			"Access your model here: ",
			model_deploy_req.model.deploy_url.color(Color::Purple1a)
		)
	);
	println!(
		"\n{} Start invoking your model using the {} command",
		"➜".color(Color::Purple1a),
		"rezon run".color(Color::Purple1a).bold()
	);
}

/// Trigger the model upload to GCS
async fn trigger_model_upload<'a>(
	http_client: Client,
	upload_endpoint: String,
	form_data: Form,
	api_key: String,
) -> Result<UploadResponse, CliError<'a>> {
	let err_template = CliError {
		error: Some("Error: Failed to Upload model to Rezon cloud. Please try again!"),
		exit_code: 64,
	};

	match http_client
		.post(upload_endpoint)
		.header("Authorization", api_key)
		.multipart(form_data)
		.send()
		.await
	{
		Ok(res) => {
			// Check to make sure that our request returned a 200
			if !res.status().is_success() {
				return Err(err_template);
			}

			match res.json::<UploadResponse>().await {
				Ok(res) => Ok(res),
				Err(_) => Err(err_template),
			}
		}
		Err(_) => Err(err_template),
	}
}

/// Trigger the model deployment to Google Cloud Build via a model_id
async fn trigger_model_deployment<'a>(
	http_client: Client,
	deploy_endpoint: String,
	deploy_body: DeployInput,
	api_key: String,
) -> Result<DeployResponse, CliError<'a>> {
	let err_template = CliError {
		error: Some("Error: Failed to Deploy model to Rezon cloud!"),
		exit_code: 64,
	};
	match http_client
		.post(deploy_endpoint)
		.header("Authorization", api_key)
		.json(&deploy_body)
		.send()
		.await
	{
		Ok(res) => {
			// Check to make sure that our request returned a 200
			if !res.status().is_success() {
				return Err(err_template);
			}

			match res.json::<DeployResponse>().await {
				Ok(res) => Ok(res),
				Err(_) => {
					return Err(err_template);
				}
			}
		}
		Err(_) => {
			return Err(err_template);
		}
	}
}

/// Takes in `model_files` and converts any file larger than 32MB into chunks of 32MB (this is the largest chunk size allowed by Cloud run)
fn chunk_model_files<'a>(model_files: Vec<File>) -> Result<Form, CliError<'a>> {
	let file_names = vec!["model.h5", "model.json"];
	let mut upload_form_data = Form::new();
	let err_template = CliError {
		error: Some("Error: Failed to Upload model to Rezon cloud. Please try again!"),
		exit_code: 64,
	};

	for (index, mut file) in model_files.iter().enumerate() {
		let file_size = match file.metadata() {
			Ok(meta) => meta.len(),
			Err(_) => {
				return Err(err_template);
			}
		};

		let file_name = file_names[index];

		// Check if the file size is greater than 32MB...
		if file_size > 32 * 1024 * 1024 {
			// Set a chunk size of 32 MB
			let mut buffer = vec![0; 32 * 1024 * 1024];
			let mut bytes_read = 0;

			while bytes_read < file_size {
				let bytes_to_read = std::cmp::min(buffer.len() as u64, file_size - bytes_read) as usize;

				if let Err(_) = file.read_exact(&mut buffer[..bytes_to_read]) {
					return Err(err_template);
				}

				let chunk = Part::bytes(buffer.clone());

				let part = chunk.file_name(file_name);

				// Now that we have all the chunks for this specifc file lets add it to the form data as an additional part
				upload_form_data = upload_form_data.part(file_name, part);

				// Increment the amount of bytes in the file that we have read
				bytes_read += bytes_to_read as u64;
			}
		} else {
			// If the file was not larger than 32mb then we can safely send it to Cloud run with no issues
			// Read the file to bytes
			let mut buffer = Vec::new();
			if let Err(_) = file.read_to_end(&mut buffer) {
				return Err(err_template);
			}

			// Create a new part from the bytes
			let part = Part::bytes(buffer.clone()).file_name(file_name);

			// Add the part to the form data
			upload_form_data = upload_form_data.part(file_name, part);
		}
	}

	Ok(upload_form_data)
}
