From 7fcc0bcd3482bbc8771e330942ef6cd78cc8ec35 Mon Sep 17 00:00:00 2001
From: Amr Bashir <amr.bashir2015@gmail.com>
Date: Thu, 1 Feb 2024 13:05:45 +0200
Subject: [PATCH] feat(cli/add): add plugin init code to `main.rs` (#8490)

* feat(cli/add): add plugin to `main.rs`

closes #7696

* clippy

* fix change file

* adjust regex

* run cargo fmt

* check if already exists

* typo [skip ci]

---------

Co-authored-by: Lucas Nogueira <lucas@tauri.studio>
---
 .changes/cli-add-to-main-rs.md |  6 +++
 tooling/cli/src/add.rs         | 98 +++++++++++++++++++++-------------
 2 files changed, 68 insertions(+), 36 deletions(-)
 create mode 100644 .changes/cli-add-to-main-rs.md

diff --git a/.changes/cli-add-to-main-rs.md b/.changes/cli-add-to-main-rs.md
new file mode 100644
index 00000000000..91b83091f94
--- /dev/null
+++ b/.changes/cli-add-to-main-rs.md
@@ -0,0 +1,6 @@
+---
+'tauri-cli': 'patch:feat'
+'@tauri-apps/cli': 'patch:feat'
+---
+
+Add plugin initialization rust code when using `tauri add`
diff --git a/tooling/cli/src/add.rs b/tooling/cli/src/add.rs
index db535a03940..4a8736ff059 100644
--- a/tooling/cli/src/add.rs
+++ b/tooling/cli/src/add.rs
@@ -4,6 +4,8 @@
 
 use anyhow::Context;
 use clap::Parser;
+use colored::Colorize;
+use regex::Regex;
 
 use crate::{
   helpers::{
@@ -41,8 +43,10 @@ pub fn command(options: Options) -> Result<()> {
   let mut plugins = plugins();
   let metadata = plugins.remove(plugin.as_str()).unwrap_or_default();
 
+  let tauri_dir = tauri_dir();
+
   let mut cargo = Command::new("cargo");
-  cargo.current_dir(tauri_dir()).arg("add").arg(&crate_name);
+  cargo.current_dir(&tauri_dir).arg("add").arg(&crate_name);
 
   if options.tag.is_some() || options.rev.is_some() || options.branch.is_some() {
     cargo
@@ -106,45 +110,67 @@ pub fn command(options: Options) -> Result<()> {
     }
   }
 
-  let rust_code = if metadata.builder {
-    if metadata.desktop_only {
-      format!(
-        r#"tauri::Builder::default()
-    .setup(|app| {{
-        #[cfg(desktop)]
-        app.handle().plugin(tauri_plugin_{plugin_snake_case}::Builder::new().build());
-        Ok(())
-    }})
-    "#,
-      )
-    } else {
-      format!(
-        r#"tauri::Builder::default()
-    .setup(|app| {{
-        app.handle().plugin(tauri_plugin_{plugin_snake_case}::Builder::new().build());
-        Ok(())
-    }})
-    "#,
-      )
+  // add plugin init code to main.rs or lib.rs
+  let plugin_init_fn = if metadata.builder {
+    "Builder::new().build()"
+  } else {
+    "init()"
+  };
+  let plugin_init = format!(".plugin(tauri_plugin_{plugin_snake_case}::{plugin_init_fn})");
+  let re = Regex::new(r"(tauri\s*::\s*Builder\s*::\s*default\(\))(\s*)")?;
+  for file in [tauri_dir.join("src/main.rs"), tauri_dir.join("src/lib.rs")] {
+    let contents = std::fs::read_to_string(&file)?;
+
+    if contents.contains(&plugin_init) {
+      log::info!(
+        "Plugin initialization code already found on {}",
+        file.display()
+      );
+      return Ok(());
+    }
+
+    if re.is_match(&contents) {
+      let out = re.replace(&contents, format!("$1$2{plugin_init}$2"));
+
+      log::info!("Adding plugin to {}", file.display());
+      std::fs::write(file, out.as_bytes())?;
+
+      // run cargo fmt
+      log::info!("Running `cargo fmt`...");
+      let _ = Command::new("cargo")
+        .arg("fmt")
+        .current_dir(&tauri_dir)
+        .status();
+
+      return Ok(());
     }
-  } else if metadata.desktop_only {
-    format!(
-      r#"tauri::Builder::default()
-    .setup(|app| {{
-        #[cfg(desktop)]
-        app.handle().plugin(tauri_plugin_{plugin_snake_case}::init());
-        Ok(())
-    }})
-    "#,
-    )
+  }
+
+  let builder_code = if metadata.builder {
+    format!(r#"+    .plugin(tauri_plugin_{plugin_snake_case}::Builder::new().build())"#,)
   } else {
-    format!(
-      r#"tauri::Builder::default().plugin(tauri_plugin_{plugin_snake_case}::init())
-    "#,
-    )
+    format!(r#"+    .plugin(tauri_plugin_{plugin_snake_case}::init())"#)
   };
 
-  println!("You must enable the plugin in your Rust code:\n\n{rust_code}");
+  let rust_code = format!(
+    r#" {}
+{}
+     {}"#,
+    "tauri::Builder::default()".dimmed(),
+    builder_code.normal().green(),
+    r#".invoke_handler(tauri::generate_handler![])
+     .run(tauri::generate_context!())
+     .expect("error while running tauri application");"#
+      .dimmed(),
+  );
+
+  log::warn!(
+    "Couldn't find `{}` in `{}` or `{}`, you must enable the plugin in your Rust code manually:\n\n{}",
+    "tauri::Builder".cyan(),
+    "main.rs".cyan(),
+    "lib.rs".cyan(),
+    rust_code
+  );
 
   Ok(())
 }