nitro_plugin 0.30.0

use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Debug;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::Context;
use anyhow::bail;
use nitro_shared::output::NitroOutput;
use serde::Serialize;
use serde::{Deserialize, Deserializer};
use tokio::sync::Mutex;

use crate::PluginPaths;
use crate::hook::Hook;
use crate::hook::PLUGIN_DIR_TOKEN;
use crate::hook::WASM_FILE_NAME;
use crate::hook::call::HookCallArg;
use crate::hook::call::HookCallContext;
use crate::hook::call::HookHandle;
use crate::hook::hooks::StartWorker;
use crate::hook::wasm::call_wasm;
use crate::hook::wasm::loader::WASMLoader;
use crate::host::PluginContext;

/// The newest protocol version for plugin communication
pub const NEWEST_PROTOCOL_VERSION: u16 = 3;
/// The default protocol version used for compatability
pub const DEFAULT_PROTOCOL_VERSION: u16 = 1;
/// Token used for file replacement in hook handlers
pub static FILE_REPLACEMENT_TOKEN: &str = "$file:";

/// A plugin
pub struct Plugin {
	/// The plugin's ID
	id: String,
	/// The plugin's manifest
	pub manifest: PluginManifest,
	/// The custom config for the plugin, serialized from JSON
	custom_config: Option<String>,
	/// The working directory for the plugin
	working_dir: Option<PathBuf>,
	/// The persistent state of the plugin
	persistence: Arc<Mutex<PluginPersistence>>,
}

impl Plugin {
	/// Create a new plugin from an ID and manifest
	pub fn new(id: String, manifest: PluginManifest) -> Self {
		Self {
			id,
			manifest,
			custom_config: None,
			working_dir: None,
			persistence: Arc::new(Mutex::new(PluginPersistence::new())),
		}
	}

	/// Get the ID of the plugin
	pub fn get_id(&self) -> &String {
		&self.id
	}

	/// Get the manifest of the plugin
	pub fn get_manifest(&self) -> &PluginManifest {
		&self.manifest
	}

	/// Call a hook on the plugin
	pub async fn call_hook<H: Hook>(
		&self,
		hook: &H,
		arg: &H::Arg,
		paths: &PluginPaths,
		nitro_version: Option<&str>,
		plugin_list: &[String],
		wasm_loader: Arc<Mutex<WASMLoader>>,
		context: Option<&Arc<dyn PluginContext>>,
		o: &mut impl NitroOutput,
	) -> anyhow::Result<Option<HookHandle<H>>> {
		let Some(handler) = self.manifest.hooks.get(hook.get_name()) else {
			return Ok(None);
		};

		self.call_hook_handler(
			hook,
			handler,
			arg,
			paths,
			nitro_version,
			plugin_list,
			wasm_loader,
			context,
			o,
		)
		.await
	}

	/// Call a hook handler on the plugin
	async fn call_hook_handler<H: Hook>(
		&self,
		hook: &H,
		handler: &HookHandler,
		arg: &H::Arg,
		paths: &PluginPaths,
		nitro_version: Option<&str>,
		plugin_list: &[String],
		wasm_loader: Arc<Mutex<WASMLoader>>,
		context: Option<&Arc<dyn PluginContext>>,
		o: &mut impl NitroOutput,
	) -> anyhow::Result<Option<HookHandle<H>>> {
		match handler {
			HookHandler::Wasm { .. } => {
				let file = self
					.working_dir
					.as_ref()
					.context("WASM handler without working dir")?
					.join(WASM_FILE_NAME)
					.to_string_lossy()
					.to_string();

				let subscriptions = HashSet::new();
				let ctx = HookCallContext {
					subscriptions: &subscriptions,
					custom_config: self.custom_config.clone(),
					nitro_version,
					plugin_list,
					global_context: context,
				};

				let arg = HookCallArg {
					cmd: &file,
					arg,
					additional_args: &[],
					working_dir: self.working_dir.as_deref(),
					ctx,
					use_base64: !self.manifest.raw_transfer,
					persistence: self.persistence.clone(),
					paths,
					plugin_id: &self.id,
					protocol_version: self
						.manifest
						.protocol_version
						.unwrap_or(DEFAULT_PROTOCOL_VERSION),
					wasm_loader,
				};
				call_wasm(hook, arg, o).await.map(Some)
			}
			HookHandler::Execute {
				executable,
				args,
				priority: _,
				subscriptions,
			} => {
				let ctx = HookCallContext {
					subscriptions,
					custom_config: self.custom_config.clone(),
					nitro_version,
					plugin_list,
					global_context: context,
				};

				let arg = HookCallArg {
					cmd: executable,
					arg,
					additional_args: args,
					working_dir: self.working_dir.as_deref(),
					ctx,
					use_base64: !self.manifest.raw_transfer,
					persistence: self.persistence.clone(),
					paths,
					plugin_id: &self.id,
					protocol_version: self
						.manifest
						.protocol_version
						.unwrap_or(DEFAULT_PROTOCOL_VERSION),
					wasm_loader,
				};
				hook.call(arg, o).await.map(Some)
			}
			HookHandler::Constant {
				constant,
				priority: _,
			} => {
				// Replace file tokens
				let mut value = constant.clone();

				replace_file_tokens(&mut value, &self.working_dir, false)?;

				Ok(Some(HookHandle::constant(
					serde_json::from_value(value)?,
					self.id.clone(),
				)))
			}
			HookHandler::File { file, priority: _ } => {
				let Some(working_dir) = &self.working_dir else {
					bail!("Plugin does not have a directory for the file hook handler to look in");
				};

				let path = working_dir.join(file);
				let contents = std::fs::read_to_string(path)
					.context("Failed to read hook result from file")?;

				// Try to read the result with quotes wrapping it if it doesn't deserialize properly the first time
				let result = match serde_json::from_str(&contents) {
					Ok(result) => result,
					Err(_) => serde_json::from_value(serde_json::Value::String(contents))
						.context("Failed to deserialize hook result")?,
				};

				Ok(Some(HookHandle::constant(result, self.id.clone())))
			}
			HookHandler::Match {
				property,
				cases,
				priority: _,
			} => {
				let arg2 = serde_json::to_value(arg)?;
				let lhs = if let Some(property) = property {
					let arg2 = arg2.as_object().context(
						"Hook argument is not an object, so a property cannot be matched",
					)?;
					arg2.get(property)
						.context("Property does not exist on hook argument")
						.cloned()?
				} else {
					arg2
				};
				let lhs = serde_json::to_string(&lhs)?;

				for (case, handler) in cases.iter() {
					if &lhs == case {
						return Box::pin(self.call_hook_handler(
							hook,
							handler,
							arg,
							paths,
							nitro_version,
							plugin_list,
							wasm_loader,
							context,
							o,
						))
						.await;
					}
				}

				Ok(None)
			}
			HookHandler::Native {
				function,
				priority: _,
			} => {
				let arg = serde_json::to_value(arg)
					.context("Failed to serialize native hook argument")?;
				let result = function
					.call(arg)
					.await
					.context("Native hook handler failed")?;
				let result = serde_json::from_value(result)
					.context("Failed to deserialize native hook result")?;
				Ok(Some(HookHandle::constant(result, self.id.clone())))
			}
		}
	}

	/// Set the custom config of the plugin
	pub fn set_custom_config(&mut self, config: serde_json::Value) -> anyhow::Result<()> {
		let serialized =
			serde_json::to_string(&config).context("Failed to serialize custom plugin config")?;
		self.custom_config = Some(serialized);
		Ok(())
	}

	/// Set the working dir of the plugin
	pub fn set_working_dir(&mut self, dir: PathBuf) {
		self.working_dir = Some(dir);
	}

	/// Set the plugin's worker handle
	pub async fn set_worker(&mut self, worker: HookHandle<StartWorker>) -> anyhow::Result<()> {
		let mut lock = self.persistence.lock().await;
		lock.worker = Some(worker);

		Ok(())
	}

	/// Get the priority of the given hook
	pub fn get_hook_priority<H: Hook>(&self, hook: &H) -> HookPriority {
		let Some(handler) = self.manifest.hooks.get(hook.get_name()) else {
			return HookPriority::Any;
		};
		match handler {
			HookHandler::Wasm { priority, .. }
			| HookHandler::Execute { priority, .. }
			| HookHandler::Constant { priority, .. }
			| HookHandler::File { priority, .. }
			| HookHandler::Match { priority, .. }
			| HookHandler::Native { priority, .. } => *priority,
		}
	}
}

/// The manifest for a plugin that describes how it works
#[derive(Deserialize, Debug, Default)]
#[serde(default)]
pub struct PluginManifest {
	/// ID for the plugin
	pub id: Option<String>,
	/// Metadata for the plugin
	#[serde(flatten)]
	pub meta: PluginMetadata,
	/// The current version of the plugin
	pub version: Option<String>,
	/// The Nitrolaunch version this plugin is for
	#[serde(alias = "mcvm_version")]
	pub nitro_version: Option<String>,
	/// The hook handlers for the plugin
	pub hooks: HashMap<String, HookHandler>,
	/// Plugins that this plugin depends on
	pub dependencies: Vec<String>,
	/// Message to display when the plugin is installed
	pub install_message: Option<String>,
	/// The protocol version of the plugin
	pub protocol_version: Option<u16>,
	/// Whether to disable base64 encoding in the protocol
	pub raw_transfer: bool,
	/// Whether the plugin supports creating custom instances
	pub supports_instance_creation: bool,
	/// Whether the plugin supports creating custom templates
	pub supports_template_creation: bool,
	/// The subcommands the plugin provides
	pub subcommands: HashMap<String, PluginProvidedSubcommand>,
}

impl PluginManifest {
	/// Create a new PluginManifest
	pub fn new() -> Self {
		Self::default()
	}
}

/// Optional metadata for a plugin
#[derive(Serialize, Deserialize, Default, Debug)]
#[serde(default)]
pub struct PluginMetadata {
	/// The display name of the plugin
	pub name: Option<String>,
	/// The short description of the plugin
	pub description: Option<String>,
	/// URL for plugin documentation
	pub documentation: Option<String>,
}

/// A CLI subcommand provided by a plugin
#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum PluginProvidedSubcommand {
	/// A root-level subcommand, containing the description
	Global(String),
	/// A subsubcommand
	Specific {
		/// The command to be under
		supercommand: String,
		/// The description
		description: String,
	},
}

/// A handler for a single hook that a plugin uses
#[derive(Deserialize)]
#[serde(untagged)]
#[serde(rename_all = "snake_case")]
pub enum HookHandler {
	/// Handle this hook by running WASM
	Wasm {
		/// Marker thingy
		wasm: bool,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
	},
	/// Handle this hook by running an executable
	Execute {
		/// The executable to run
		executable: String,
		/// Arguments for the executable
		#[serde(default)]
		args: Vec<String>,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
		/// Data for the hook to subscribe to
		#[serde(default)]
		subscriptions: HashSet<HookSubscription>,
	},
	/// Handle this hook by returning a constant result
	Constant {
		/// The constant result
		constant: serde_json::Value,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
	},
	/// Handle this hook by getting the contents of a file
	File {
		/// The path to the file, relative to the plugin directory
		file: String,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
	},
	/// Match against the argument to handle the hook differently
	Match {
		/// The property to match against
		#[serde(default)]
		property: Option<String>,
		/// The cases of the match
		cases: HashMap<String, Box<HookHandler>>,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
	},
	/// Handle this hook with a native function call
	Native {
		/// The function to handle the hook
		#[serde(deserialize_with = "deserialize_native_function")]
		function: Arc<dyn NativeHookHandler>,
		/// The priority for the hook
		#[serde(default)]
		priority: HookPriority,
	},
}

impl Debug for HookHandler {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "HookHandler")
	}
}

/// Priority for a hook
#[derive(Deserialize, PartialEq, PartialOrd, Eq, Ord, Default, Clone, Copy)]
#[serde(rename_all = "snake_case")]
pub enum HookPriority {
	/// The plugin will try to run before other ones
	First,
	/// The plugin will run at any time, usually in the middle
	#[default]
	Any,
	/// The plugin will try to run after other ones
	Last,
}

/// Data that an executable handler can subscribe to receiving
#[derive(Deserialize, PartialEq, Eq, Clone, Copy, Hash)]
#[serde(rename_all = "snake_case")]
pub enum HookSubscription {
	/// List of instances
	Instances,
	/// List of templates
	Templates,
}

/// Deserialize function for the native hook. No plugin manifests should ever use this,
/// so just return a function that does nothing.
fn deserialize_native_function<'de, D>(_: D) -> Result<Arc<dyn NativeHookHandler>, D::Error>
where
	D: Deserializer<'de>,
{
	Ok(Arc::new(NoneHookHandler))
}

/// Trait for native plugin hook handlers
#[async_trait::async_trait]
pub trait NativeHookHandler: Send + Sync {
	/// Call the hook
	async fn call(&self, arg: serde_json::Value) -> anyhow::Result<serde_json::Value>;
}

/// Used for deserialization
struct NoneHookHandler;

#[async_trait::async_trait]
impl NativeHookHandler for NoneHookHandler {
	async fn call(&self, arg: serde_json::Value) -> anyhow::Result<serde_json::Value> {
		let _ = arg;
		Ok(serde_json::Value::Null)
	}
}

/// Persistent state for plugins
pub struct PluginPersistence {
	/// The persistent state of the plugin
	pub state: serde_json::Value,
	/// The long-running plugin worker
	pub worker: Option<HookHandle<StartWorker>>,
}

impl Default for PluginPersistence {
	fn default() -> Self {
		Self::new()
	}
}

impl PluginPersistence {
	/// Initalize the persistent plugin state
	pub fn new() -> Self {
		Self {
			state: serde_json::Value::Null,
			worker: None,
		}
	}
}

/// Replaces file tokens in the string values of JSON value
fn replace_file_tokens(
	value: &mut serde_json::Value,
	working_dir: &Option<PathBuf>,
	test: bool,
) -> anyhow::Result<()> {
	match value {
		serde_json::Value::Array(values) => {
			for value in values {
				replace_file_tokens(value, working_dir, test)?;
			}
		}
		serde_json::Value::Object(props) => {
			for prop in props.values_mut() {
				replace_file_tokens(prop, working_dir, test)?;
			}
		}
		serde_json::Value::String(value) => {
			if let Some(path) = value.strip_prefix(FILE_REPLACEMENT_TOKEN) {
				if test {
					*value = "test".into();
					return Ok(());
				}

				let Some(working_dir) = working_dir else {
					bail!("Plugin does not have a directory for the file hook handler to look in");
				};

				let path = working_dir.join(path);
				let contents = std::fs::read_to_string(path)
					.context("Failed to read hook result from file")?;

				*value = contents;
			}

			if let Some(working_dir) = working_dir {
				*value = value.replace(PLUGIN_DIR_TOKEN, &working_dir.to_string_lossy());
			}
		}
		_ => {}
	}

	Ok(())
}

#[cfg(test)]
mod tests {
	use super::*;
	use serde_json::json;

	#[test]
	fn test_file_token_replacement() {
		let mut json = json!([{
			"foo": "bar",
			"baz": format!("{FILE_REPLACEMENT_TOKEN}foobar")
		}]);

		replace_file_tokens(&mut json, &None, true).unwrap();

		let expected = json!([{
			"foo": "bar",
			"baz": "test"
		}]);
		assert_eq!(json, expected);
	}
}