#![allow(non_snake_case)]
use super::RezonCommand;
use crate::error::Error::{self, *};
use crate::{
	config::rezon_config::RezonConfig,
	tui::{clean_console, indent, SPINNER_COLOR},
	util::{get_model_file_names, get_source_type, wait_for_deployment, BASE_BFF_URL},
};
use async_trait::async_trait;
use clap::{ArgMatches, Command};
use colorful::{Color, Colorful};
use log::info;
use requestty::{prompt_one, Question};
use reqwest::StatusCode;
use reqwest::{
	multipart::{Form, Part},
	Client,
};
use serde::{Deserialize, Serialize};
use spinoff::{spinners, Spinner};
use std::{fs::File, io::Read};
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,
	pub model_type: ModelType,
}

#[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,
}

pub struct ModelSource {
	pub files: Vec<File>,
	pub model_type: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct ModelType {
	pub framework: String,
	pub source_type: String,
}

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

	// TODO: replace dummy output with actual functionality
	async fn execute(_args: &ArgMatches) -> Result<(), Error> {
		// unwrapping config is safe because we check for it already
		deploy_model(RezonConfig::read_config()?).await?;
		Ok(())
	}
}

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

	let model_files = get_model_files(config)?;
	let is_h5_and_json = model_files.files.len() > 1;
	let model_source_type = get_source_type(is_h5_and_json, &model_type);

	// (2) Iterate over the files and add them to the form data in the same request
	let upload_form_data = chunk_model_files(model_files.files, is_h5_and_json, &model_type)?;

	// (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(),
	)?;

	// Sanitize the model name and always convert it to a lowercase string
	let model_name = model_name.as_string().unwrap().to_lowercase();

	// TODO: this needs moved to our deploy endpoint - we should not be validating schemas on clients
	// Validate that the project name does not contain any invalid chars
	if !model_name.chars().all(|x| x.is_alphanumeric() || x == '-') {
		return Err(CliError(format!(
			"Aborting...your model name may only contain alphanumeric characters along with '-'..."
		)));
	}

	// (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_with_stream(spinners::Aesthetic, "Uploading model...", SPINNER_COLOR, spinoff::Streams::Stderr);

	let model_upload_req = match trigger_model_upload(&http_client, &upload_endpoint, upload_form_data, &api_key).await {
		Ok(req) => req,
		Err(e) => {
			spinner.stop();
			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(),
		model_type: ModelType {
			framework: model_type[0..1].to_uppercase() + &model_type[1..],
			source_type: model_source_type.to_owned(),
		},
	};

	println!(
		"{} Current deployment status: {}\n",
		"➜".color(Color::Purple1a).bold(),
		"INITIALIZING".color(Color::Purple1a).bold()
	);
	clean_console();

	// Update the spinner to have the deployment text relative to the current task (deploying the actual model)
	spinner.update_text("Deploying model to Rezon Cloud...");

	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 before throwing the error
			spinner.stop();
			return Err(e);
		}
	};

	// (6) Load until the status is equal to "SUCCESS" or throw an error if the status is "FAILURE" (or anything else)
	if let Err(e) = wait_for_deployment(http_client, model_deploy_req.build_id).await {
		// since we got an error, lets stop the spinner before throwing the error
		spinner.stop();
		return Err(e);
	};

	// 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
	info!("{} 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(http_client: &Client, upload_endpoint: &String, form_data: Form, api_key: &String) -> Result<UploadResponse, Error> {
	let response = http_client
		.post(upload_endpoint)
		.header("Authorization", api_key)
		.multipart(form_data)
		.send()
		.await?;
	// Check to make sure that our request returned a 200
	if response.status() != StatusCode::OK {
		return Err(NotOk {
			response: response.text().await?,
		});
	}

	Ok(response.json::<UploadResponse>().await?)
}

/// Trigger the model deployment to Google Cloud Build via a model_id
async fn trigger_model_deployment(
	http_client: Client,
	deploy_endpoint: String,
	deploy_body: DeployInput,
	api_key: String,
) -> Result<DeployResponse, Error> {
	let response = http_client
		.post(deploy_endpoint)
		.header("Authorization", api_key)
		.json(&deploy_body)
		.send()
		.await?;

	// Check to make sure that our request returned a 200
	if response.status() != StatusCode::OK {
		return Err(NotOk {
			response: response.text().await?,
		});
	}

	Ok(response.json::<DeployResponse>().await?)
}

/// 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(model_files: Vec<File>, is_json_and_h5: bool, model_type: &str) -> Result<Form, Error> {
	let file_names = get_model_file_names(is_json_and_h5, model_type);
	let mut upload_form_data = Form::new();
	for (index, mut file) in model_files.iter().enumerate() {
		let file_name = file_names[index];

		let file_size = match file.metadata() {
			Ok(meta) => meta.len(),
			Err(e) => {
				return Err(Io {
					path: std::path::PathBuf::from(file_name),
					source: e,
				});
			}
		};

		// 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(e) = file.read_exact(&mut buffer[..bytes_to_read]) {
					return Err(Io {
						path: std::path::PathBuf::from(file_name),
						source: e,
					});
				}

				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(e) = file.read_to_end(&mut buffer) {
				return Err(Io {
					path: std::path::PathBuf::from(file_name),
					source: e,
				});
			}

			// 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)
}

pub fn get_model_files(config: RezonConfig) -> Result<ModelSource, Error> {
	// (1) check for the .json and .h5 keys
	let model_json = config.model.json_config_path;
	let model_h5 = config.model.h5_without_config_path;
	let model_path = config.model.model_path;
	let model_type = config.model.model_type;

	if let Some(path) = model_path {
		return check_for_model_path(path, model_type);
	}

	// if the .json and .h5 keys are not both present then check for the model_path key (return an error if either are None)
	let (Some(model_json), Some(model_h5 )) = (&model_json, &model_h5) else {
		return Err(CliError(format!("could not find model file(s):\nmodel_json: {}\nmodel_h5: {}", model_json.is_some(), model_h5.is_some())));
	};

	// Fallback to using the json/h5 file approach if the `model_path` key is not present
	// Note we can safely unwrap these because we already checked for their presence above (they also get validated via the RezonConfig validation)
	let json = match File::open(model_json) {
		Ok(file) => file,
		Err(e) => {
			return Err(Io {
				path: std::path::PathBuf::from(model_json),
				source: e,
			});
		}
	};

	let h5 = match File::open(model_h5) {
		Ok(file) => file,
		Err(e) => {
			return Err(Io {
				path: std::path::PathBuf::from(model_h5),
				source: e,
			});
		}
	};

	Ok(ModelSource {
		files: vec![h5, json],
		model_type,
	})
}

pub fn check_for_model_path(path: String, model_type: String) -> Result<ModelSource, Error> {
	// Since the user did specify a model file path lets open the file and send the source back up
	let model_file = match File::open(&path) {
		Ok(f) => f,
		Err(e) => {
			return Err(Io {
				path: std::path::PathBuf::from(path),
				source: e,
			})
		}
	};

	return Ok(ModelSource {
		files: vec![model_file],
		model_type,
	});
}
