feat Add PowerShell support (#2543)

This adds PowerShell support 🎉 

I built this script around @lzybkr's
[prototype](https://github.com/atuinsh/atuin/issues/84#issuecomment-1689168533),
so I added him as co-author (I hope that's ok). I wouldn't know where to
start without his contribution.

I'm not a PowerShell expert, so this was a nice opportunity to learn
some stuff. I think it's ok, but I would appreciate if someone more
knowledgeable in the matter could review this though.

It would be nice if other PowerShell users could test this with their
configs and report any issues. I wouldn't be surprised if there are some
remaining bugs or missing features.

Fixes #84

## Installation

If you'd like to test this, you can install the `atuin` from this PR by
running:

```powershell
cargo install --git https://github.com/ltrzesniewski/atuin.git --branch powershell-pr
```

Then, add the following to your PowerShell profile file (whose path is
in `$PROFILE`) and restart the shell:

```powershell
atuin init powershell | Out-String | Invoke-Expression
```

This requires `atuin` to be in the path and the
[PSReadLine](https://github.com/PowerShell/PSReadLine) module to be
installed, which is the case by default.

## Tests

I tested this on the following:
- PowerShell 7.4.6 / PSReadLine 2.3.5 / Windows (the latest one)
- PowerShell 7.5.1 / PSReadLine 2.3.6 / Windows (the latest one)
- PowerShell 5.1.22621.4391 / PSReadLine 2.0.0 / Windows (the one
shipped with Windows)
- PowerShell 7.4.6 / PSReadLine 2.3.5 / Ubuntu WSL (strangely, it didn't
behave exactly like the Windows version)
- PowerShell 7.5.1 / PSReadLine 2.3.6 / Ubuntu WSL

I also tested this with and without my custom [Oh My
Posh](https://ohmyposh.dev/) prompt. It works fine in both cases,
~except that since my OMP config contains `"newline": true`, my prompt
is multiline and shifts downwards by a single line on each `atuin search
-i` invocation. This can be adjusted with the
`$env:ATUIN_POWERSHELL_PROMPT_OFFSET` environment variable (e.g. I set
mine to `-1` to account for the additional prompt line).~ (this variable
is now auto-initialized).

## Checks
- [x] I am happy for maintainers to push small adjustments to this PR,
to speed up the review cycle
- [x] I have checked that there are no existing pull requests for the
same thing

---------

Co-authored-by: Jason Shirk <jasonsh@microsoft.com>

11 files changed, 528 insertions, 47 deletions

diff --git a/crates/atuin-dotfiles/src/shell.rs b/crates/atuin-dotfiles/src/shell.rs
index bd61aafa..73a9ce8c 100644
--- a/crates/atuin-dotfiles/src/shell.rs
+++ b/crates/atuin-dotfiles/src/shell.rs
@@ -8,6 +8,7 @@ use crate::store::AliasStore;
 
 pub mod bash;
 pub mod fish;
+pub mod powershell;
 pub mod xonsh;
 pub mod zsh;
 
diff --git a/crates/atuin-dotfiles/src/shell/powershell.rs b/crates/atuin-dotfiles/src/shell/powershell.rs
new file mode 100644
index 00000000..1daee28b
--- /dev/null
+++ b/crates/atuin-dotfiles/src/shell/powershell.rs
@@ -0,0 +1,169 @@
+use crate::shell::{Alias, Var};
+use crate::store::{AliasStore, var::VarStore};
+use std::path::PathBuf;
+
+async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String {
+    match tokio::fs::read_to_string(path).await {
+        Ok(aliases) => aliases,
+        Err(r) => {
+            // we failed to read the file for some reason, but the file does exist
+            // fallback to generating new aliases on the fly
+
+            store.powershell().await.unwrap_or_else(|e| {
+                format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",)
+            })
+        }
+    }
+}
+
+async fn cached_vars(path: PathBuf, store: &VarStore) -> String {
+    match tokio::fs::read_to_string(path).await {
+        Ok(vars) => vars,
+        Err(r) => {
+            // we failed to read the file for some reason, but the file does exist
+            // fallback to generating new vars on the fly
+
+            store.powershell().await.unwrap_or_else(|e| {
+                format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",)
+            })
+        }
+    }
+}
+
+/// Return powershell dotfile config
+///
+/// Do not return an error. We should not prevent the shell from starting.
+///
+/// In the worst case, Atuin should not function but the shell should start correctly.
+///
+/// While currently this only returns aliases, it will be extended to also return other synced dotfiles
+pub async fn alias_config(store: &AliasStore) -> String {
+    // First try to read the cached config
+    let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.ps1");
+
+    if aliases.exists() {
+        return cached_aliases(aliases, store).await;
+    }
+
+    if let Err(e) = store.build().await {
+        return format!("echo 'Atuin: failed to generate aliases: {e}'");
+    }
+
+    cached_aliases(aliases, store).await
+}
+
+pub async fn var_config(store: &VarStore) -> String {
+    // First try to read the cached config
+    let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.ps1");
+
+    if vars.exists() {
+        return cached_vars(vars, store).await;
+    }
+
+    if let Err(e) = store.build().await {
+        return format!("echo 'Atuin: failed to generate vars: {e}'");
+    }
+
+    cached_vars(vars, store).await
+}
+
+pub fn format_alias(alias: &Alias) -> String {
+    // Set-Alias doesn't support adding implicit arguments, so use a function.
+    // See https://github.com/PowerShell/PowerShell/issues/12962
+
+    let mut result = secure_command(&format!(
+        "function {} {{\n    {}{} @args\n}}",
+        alias.name,
+        if alias.value.starts_with(['"', '\'']) {
+            "& "
+        } else {
+            ""
+        },
+        alias.value
+    ));
+
+    // This makes the file layout prettier
+    result.insert(0, '\n');
+    result
+}
+
+pub fn format_var(var: &Var) -> String {
+    secure_command(&format!(
+        "${}{} = '{}'",
+        if var.export { "env:" } else { "" },
+        var.name,
+        var.value.replace("'", "''")
+    ))
+}
+
+/// Wraps the given command in an Invoke-Expression to ensure the outer script is not halted
+/// if the inner command contains a syntax error.
+fn secure_command(command: &str) -> String {
+    format!(
+        "Invoke-Expression -ErrorAction Continue -Command '{}'\n",
+        command.replace("'", "''")
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn aliases() {
+        assert_eq!(
+            format_alias(&Alias {
+                name: "gp".to_string(),
+                value: "git push".to_string(),
+            }),
+            "\n".to_string()
+                + &secure_command(
+                    "function gp {
+    git push @args
+}"
+                )
+        );
+
+        assert_eq!(
+            format_alias(&Alias {
+                name: "spc".to_string(),
+                value: "\"path with spaces\" arg".to_string(),
+            }),
+            "\n".to_string()
+                + &secure_command(
+                    "function spc {
+    & \"path with spaces\" arg @args
+}"
+                )
+        );
+    }
+
+    #[test]
+    fn vars() {
+        assert_eq!(
+            format_var(&Var {
+                name: "FOO".to_owned(),
+                value: "bar 'baz'".to_owned(),
+                export: true,
+            }),
+            secure_command("$env:FOO = 'bar ''baz'''")
+        );
+
+        assert_eq!(
+            format_var(&Var {
+                name: "TEST".to_owned(),
+                value: "1".to_owned(),
+                export: false,
+            }),
+            secure_command("$TEST = '1'")
+        );
+    }
+
+    #[test]
+    fn invoke_expression() {
+        assert_eq!(
+            secure_command("echo 'foo'"),
+            "Invoke-Expression -ErrorAction Continue -Command 'echo ''foo'''\n"
+        )
+    }
+}
diff --git a/crates/atuin-dotfiles/src/store.rs b/crates/atuin-dotfiles/src/store.rs
index 01316b4e..17597065 100644
--- a/crates/atuin-dotfiles/src/store.rs
+++ b/crates/atuin-dotfiles/src/store.rs
@@ -142,7 +142,20 @@ impl AliasStore {
 
     pub async fn posix(&self) -> Result<String> {
         let aliases = self.aliases().await?;
+        Ok(Self::format_posix(&aliases))
+    }
+
+    pub async fn xonsh(&self) -> Result<String> {
+        let aliases = self.aliases().await?;
+        Ok(Self::format_xonsh(&aliases))
+    }
 
+    pub async fn powershell(&self) -> Result<String> {
+        let aliases = self.aliases().await?;
+        Ok(Self::format_powershell(&aliases))
+    }
+
+    fn format_posix(aliases: &[Alias]) -> String {
         let mut config = String::new();
 
         for alias in aliases {
@@ -153,28 +166,39 @@ impl AliasStore {
             config.push_str(&format!("alias {}='{}'\n", alias.name, value));
         }
 
-        Ok(config)
+        config
     }
 
-    pub async fn xonsh(&self) -> Result<String> {
-        let aliases = self.aliases().await?;
-
+    fn format_xonsh(aliases: &[Alias]) -> String {
         let mut config = String::new();
 
         for alias in aliases {
             config.push_str(&format!("aliases['{}'] ='{}'\n", alias.name, alias.value));
         }
 
-        Ok(config)
+        config
+    }
+
+    fn format_powershell(aliases: &[Alias]) -> String {
+        let mut config = String::new();
+
+        for alias in aliases {
+            config.push_str(&crate::shell::powershell::format_alias(alias));
+        }
+
+        config
     }
 
     pub async fn build(&self) -> Result<()> {
         let dir = atuin_common::utils::dotfiles_cache_dir();
         tokio::fs::create_dir_all(dir.clone()).await?;
 
+        let aliases = self.aliases().await?;
+
         // Build for all supported shells
-        let posix = self.posix().await?;
-        let xonsh = self.xonsh().await?;
+        let posix = Self::format_posix(&aliases);
+        let xonsh = Self::format_xonsh(&aliases);
+        let powershell = Self::format_powershell(&aliases);
 
         // All the same contents, maybe optimize in the future or perhaps there will be quirks
         // per-shell
@@ -183,11 +207,13 @@ impl AliasStore {
         let bash = dir.join("aliases.bash");
         let fish = dir.join("aliases.fish");
         let xsh = dir.join("aliases.xsh");
+        let ps1 = dir.join("aliases.ps1");
 
         tokio::fs::write(zsh, &posix).await?;
         tokio::fs::write(bash, &posix).await?;
         tokio::fs::write(fish, &posix).await?;
         tokio::fs::write(xsh, &xonsh).await?;
+        tokio::fs::write(ps1, &powershell).await?;
 
         Ok(())
     }
diff --git a/crates/atuin-dotfiles/src/store/var.rs b/crates/atuin-dotfiles/src/store/var.rs
index 402ba684..9d25b85d 100644
--- a/crates/atuin-dotfiles/src/store/var.rs
+++ b/crates/atuin-dotfiles/src/store/var.rs
@@ -169,7 +169,25 @@ impl VarStore {
 
     pub async fn xonsh(&self) -> Result<String> {
         let env = self.vars().await?;
+        Ok(Self::format_xonsh(&env))
+    }
 
+    pub async fn fish(&self) -> Result<String> {
+        let env = self.vars().await?;
+        Ok(Self::format_fish(&env))
+    }
+
+    pub async fn posix(&self) -> Result<String> {
+        let env = self.vars().await?;
+        Ok(Self::format_posix(&env))
+    }
+
+    pub async fn powershell(&self) -> Result<String> {
+        let env = self.vars().await?;
+        Ok(Self::format_powershell(&env))
+    }
+
+    fn format_xonsh(env: &[Var]) -> String {
         let mut config = String::new();
 
         for env in env {
@@ -177,12 +195,10 @@ impl VarStore {
             config.push_str(&format!("${}={}\n", env.name, escaped_value));
         }
 
-        Ok(config)
+        config
     }
 
-    pub async fn fish(&self) -> Result<String> {
-        let env = self.vars().await?;
-
+    fn format_fish(env: &[Var]) -> String {
         let mut config = String::new();
 
         for env in env {
@@ -190,12 +206,10 @@ impl VarStore {
             config.push_str(&format!("set -gx {} {}\n", env.name, escaped_value));
         }
 
-        Ok(config)
+        config
     }
 
-    pub async fn posix(&self) -> Result<String> {
-        let env = self.vars().await?;
-
+    fn format_posix(env: &[Var]) -> String {
         let mut config = String::new();
 
         for env in env {
@@ -207,17 +221,30 @@ impl VarStore {
             }
         }
 
-        Ok(config)
+        config
+    }
+
+    fn format_powershell(env: &[Var]) -> String {
+        let mut config = String::new();
+
+        for var in env {
+            config.push_str(&crate::shell::powershell::format_var(var));
+        }
+
+        config
     }
 
     pub async fn build(&self) -> Result<()> {
         let dir = atuin_common::utils::dotfiles_cache_dir();
         tokio::fs::create_dir_all(dir.clone()).await?;
 
+        let env = self.vars().await?;
+
         // Build for all supported shells
-        let posix = self.posix().await?;
-        let xonsh = self.xonsh().await?;
-        let fsh = self.fish().await?;
+        let posix = Self::format_posix(&env);
+        let xonsh = Self::format_xonsh(&env);
+        let fsh = Self::format_fish(&env);
+        let powershell = Self::format_powershell(&env);
 
         // All the same contents, maybe optimize in the future or perhaps there will be quirks
         // per-shell
@@ -226,11 +253,13 @@ impl VarStore {
         let bash = dir.join("vars.bash");
         let fish = dir.join("vars.fish");
         let xsh = dir.join("vars.xsh");
+        let ps1 = dir.join("vars.ps1");
 
         tokio::fs::write(zsh, &posix).await?;
         tokio::fs::write(bash, &posix).await?;
         tokio::fs::write(fish, &fsh).await?;
         tokio::fs::write(xsh, &xonsh).await?;
+        tokio::fs::write(ps1, &powershell).await?;
 
         Ok(())
     }
diff --git a/crates/atuin/src/command/client/history.rs b/crates/atuin/src/command/client/history.rs
index afa0f3bc..028db5f1 100644
--- a/crates/atuin/src/command/client/history.rs
+++ b/crates/atuin/src/command/client/history.rs
@@ -34,6 +34,11 @@ use super::search::format_duration_into;
 pub enum Cmd {
     /// Begins a new command in the history
     Start {
+        /// Collects the command from the `ATUIN_COMMAND_LINE` environment variable,
+        /// which does not need escaping and is more compatible between OS and shells
+        #[arg(long = "command-from-env", hide = true)]
+        cmd_env: bool,
+
         command: Vec<String>,
     },
 
@@ -344,13 +349,7 @@ fn parse_fmt(format: &str) -> ParsedFmt<'_> {
 
 impl Cmd {
     #[allow(clippy::too_many_lines, clippy::cast_possible_truncation)]
-    async fn handle_start(
-        db: &impl Database,
-        settings: &Settings,
-        command: &[String],
-    ) -> Result<()> {
-        let command = command.join(" ");
-
+    async fn handle_start(db: &impl Database, settings: &Settings, command: &str) -> Result<()> {
         // It's better for atuin to silently fail here and attempt to
         // store whatever is ran, than to throw an error to the terminal
         let cwd = utils::get_current_dir();
@@ -375,9 +374,7 @@ impl Cmd {
     }
 
     #[cfg(feature = "daemon")]
-    async fn handle_daemon_start(settings: &Settings, command: &[String]) -> Result<()> {
-        let command = command.join(" ");
-
+    async fn handle_daemon_start(settings: &Settings, command: &str) -> Result<()> {
         // It's better for atuin to silently fail here and attempt to
         // store whatever is ran, than to throw an error to the terminal
         let cwd = utils::get_current_dir();
@@ -655,7 +652,8 @@ impl Cmd {
         // Skip initializing any databases for start/end, if the daemon is enabled
         if settings.daemon.enabled {
             match self {
-                Self::Start { command } => {
+                Self::Start { .. } => {
+                    let command = self.get_start_command().unwrap_or_default();
                     return Self::handle_daemon_start(settings, &command).await;
                 }
 
@@ -681,7 +679,10 @@ impl Cmd {
         let history_store = HistoryStore::new(store.clone(), host_id, encryption_key);
 
         match self {
-            Self::Start { command } => Self::handle_start(&db, settings, &command).await,
+            Self::Start { .. } => {
+                let command = self.get_start_command().unwrap_or_default();
+                Self::handle_start(&db, settings, &command).await
+            }
             Self::End { id, exit, duration } => {
                 Self::handle_end(&db, store, history_store, settings, &id, exit, duration).await
             }
@@ -750,6 +751,18 @@ impl Cmd {
             }
         }
     }
+
+    /// Returns the command line to use for the `Start` variant.
+    /// Returns `None` for any other variant.
+    fn get_start_command(&self) -> Option<String> {
+        match self {
+            Self::Start { cmd_env: true, .. } => {
+                Some(std::env::var("ATUIN_COMMAND_LINE").unwrap_or_default())
+            }
+            Self::Start { command, .. } => Some(command.join(" ")),
+            _ => None,
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/crates/atuin/src/command/client/init.rs b/crates/atuin/src/command/client/init.rs
index 516ccd26..410f9b6a 100644
--- a/crates/atuin/src/command/client/init.rs
+++ b/crates/atuin/src/command/client/init.rs
@@ -7,6 +7,7 @@ use eyre::{Result, WrapErr};
 
 mod bash;
 mod fish;
+mod powershell;
 mod xonsh;
 mod zsh;
 
@@ -24,6 +25,8 @@ pub struct Cmd {
 }
 
 #[derive(Clone, Copy, ValueEnum, Debug)]
+#[value(rename_all = "lower")]
+#[allow(clippy::enum_variant_names, clippy::doc_markdown)]
 pub enum Shell {
     /// Zsh setup
     Zsh,
@@ -35,6 +38,8 @@ pub enum Shell {
     Nu,
     /// Xonsh setup
     Xonsh,
+    /// PowerShell setup
+    PowerShell,
 }
 
 impl Cmd {
@@ -100,6 +105,9 @@ $env.config = (
             Shell::Xonsh => {
                 xonsh::init_static(self.disable_up_arrow, self.disable_ctrl_r);
             }
+            Shell::PowerShell => {
+                powershell::init_static(self.disable_up_arrow, self.disable_ctrl_r);
+            }
         }
     }
 
@@ -153,6 +161,15 @@ $env.config = (
                 )
                 .await?;
             }
+            Shell::PowerShell => {
+                powershell::init(
+                    alias_store,
+                    var_store,
+                    self.disable_up_arrow,
+                    self.disable_ctrl_r,
+                )
+                .await?;
+            }
         }
 
         Ok(())
diff --git a/crates/atuin/src/command/client/init/powershell.rs b/crates/atuin/src/command/client/init/powershell.rs
new file mode 100644
index 00000000..7d1b8876
--- /dev/null
+++ b/crates/atuin/src/command/client/init/powershell.rs
@@ -0,0 +1,39 @@
+use atuin_dotfiles::store::{AliasStore, var::VarStore};
+
+pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool) {
+    let base = include_str!("../../../shell/atuin.ps1");
+
+    let (bind_ctrl_r, bind_up_arrow) = if std::env::var("ATUIN_NOBIND").is_ok() {
+        (false, false)
+    } else {
+        (!disable_ctrl_r, !disable_up_arrow)
+    };
+
+    println!("{base}");
+    println!(
+        "Enable-AtuinSearchKeys -CtrlR {} -UpArrow {}",
+        ps_bool(bind_ctrl_r),
+        ps_bool(bind_up_arrow)
+    );
+}
+
+pub async fn init(
+    aliases: AliasStore,
+    vars: VarStore,
+    disable_up_arrow: bool,
+    disable_ctrl_r: bool,
+) -> eyre::Result<()> {
+    init_static(disable_up_arrow, disable_ctrl_r);
+
+    let aliases = atuin_dotfiles::shell::powershell::alias_config(&aliases).await;
+    let vars = atuin_dotfiles::shell::powershell::var_config(&vars).await;
+
+    println!("{aliases}");
+    println!("{vars}");
+
+    Ok(())
+}
+
+fn ps_bool(value: bool) -> &'static str {
+    if value { "$true" } else { "$false" }
+}
diff --git a/crates/atuin/src/command/client/search.rs b/crates/atuin/src/command/client/search.rs
index 8c864e77..4103901a 100644
--- a/crates/atuin/src/command/client/search.rs
+++ b/crates/atuin/src/command/client/search.rs
@@ -1,4 +1,5 @@
-use std::io::{IsTerminal as _, stderr};
+use std::fs::File;
+use std::io::{IsTerminal as _, Write, stderr};
 
 use atuin_common::utils::{self, Escapable as _};
 use clap::Parser;
@@ -131,6 +132,10 @@ pub struct Cmd {
     /// Include duplicate commands in the output (non-interactive only)
     #[arg(long)]
     include_duplicates: bool,
+
+    /// File name to write the result to (hidden from help as this is meant to be used from a script)
+    #[arg(long = "result-file", hide = true)]
+    result_file: Option<String>,
 }
 
 impl Cmd {
@@ -213,7 +218,11 @@ impl Cmd {
 
         if self.interactive {
             let item = interactive::history(&query, settings, db, &history_store, theme).await?;
-            if stderr().is_terminal() {
+
+            if let Some(result_file) = self.result_file {
+                let mut file = File::create(result_file)?;
+                write!(file, "{item}")?;
+            } else if stderr().is_terminal() {
                 eprintln!("{}", item.escape_control());
             } else {
                 eprintln!("{item}");
diff --git a/crates/atuin/src/command/client/search/interactive.rs b/crates/atuin/src/command/client/search/interactive.rs
index bcf456a4..930f634c 100644
--- a/crates/atuin/src/command/client/search/interactive.rs
+++ b/crates/atuin/src/command/client/search/interactive.rs
@@ -1410,18 +1410,22 @@ pub async fn history(
         terminal.clear()?;
     }
 
+    let accept = accept
+        && matches!(
+            Shell::from_env(),
+            Shell::Zsh | Shell::Fish | Shell::Bash | Shell::Xonsh | Shell::Nu | Shell::Powershell
+        );
+
+    let accept_prefix = "__atuin_accept__:";
+
     match result {
         InputAction::AcceptInspecting => {
             match inspecting {
                 Some(result) => {
                     let mut command = result.command;
-                    if accept
-                        && matches!(
-                            Shell::from_env(),
-                            Shell::Zsh | Shell::Fish | Shell::Bash | Shell::Xonsh | Shell::Nu
-                        )
-                    {
-                        command = String::from("__atuin_accept__:") + &command;
+
+                    if accept {
+                        command = String::from(accept_prefix) + &command;
                     }
 
                     // index is in bounds so we return that entry
@@ -1435,13 +1439,8 @@ pub async fn history(
 
             if is_command_chaining {
                 command = format!("{} {}", original_query.trim_end(), command);
-            } else if accept
-                && matches!(
-                    Shell::from_env(),
-                    Shell::Zsh | Shell::Fish | Shell::Bash | Shell::Xonsh | Shell::Nu
-                )
-            {
-                command = String::from("__atuin_accept__:") + &command;
+            } else if accept {
+                command = String::from(accept_prefix) + &command;
             }
 
             // index is in bounds so we return that entry
diff --git a/crates/atuin/src/command/mod.rs b/crates/atuin/src/command/mod.rs
index 95813193..a70ab629 100644
--- a/crates/atuin/src/command/mod.rs
+++ b/crates/atuin/src/command/mod.rs
@@ -18,6 +18,7 @@ mod external;
 
 #[derive(Subcommand)]
 #[command(infer_subcommands = true)]
+#[allow(clippy::large_enum_variant)]
 pub enum AtuinCmd {
     #[cfg(feature = "client")]
     #[command(flatten)]
diff --git a/crates/atuin/src/shell/atuin.ps1 b/crates/atuin/src/shell/atuin.ps1
new file mode 100644
index 00000000..f1caee86
--- /dev/null
+++ b/crates/atuin/src/shell/atuin.ps1
@@ -0,0 +1,178 @@
+# Atuin PowerShell module
+#
+# This should support PowerShell 5.1 (which is shipped with Windows) and later versions, on Windows and Linux.
+#
+# Usage: atuin init powershell | Out-String | Invoke-Expression
+#
+# Settings:
+# - $env:ATUIN_POWERSHELL_PROMPT_OFFSET - Number of lines to offset the prompt position after exiting search.
+#   This is useful when using a multi-line prompt: e.g. set this to -1 when using a 2-line prompt.
+#   It is initialized from the current prompt line count if not set when the first Atuin search is performed.
+
+if (Get-Module Atuin -ErrorAction Ignore) {
+    Write-Warning "The Atuin module is already loaded."
+    return
+}
+
+if (!(Get-Command atuin -ErrorAction Ignore)) {
+    Write-Error "The 'atuin' executable needs to be available in the PATH."
+    return
+}
+
+if (!(Get-Module PSReadLine -ErrorAction Ignore)) {
+    Write-Error "Atuin requires the PSReadLine module to be installed."
+    return
+}
+
+New-Module -Name Atuin -ScriptBlock {
+    $env:ATUIN_SESSION = atuin uuid
+
+    $script:atuinHistoryId = $null
+    $script:previousPSConsoleHostReadLine = $Function:PSConsoleHostReadLine
+
+    # The ReadLine overloads changed with breaking changes over time, make sure the one we expect is available.
+    $script:hasExpectedReadLineOverload = ([Microsoft.PowerShell.PSConsoleReadLine]::ReadLine).OverloadDefinitions.Contains("static string ReadLine(runspace runspace, System.Management.Automation.EngineIntrinsics engineIntrinsics, System.Threading.CancellationToken cancellationToken, System.Nullable[bool] lastRunStatus)")
+
+    # This function name is called by PSReadLine to read the next command line to execute.
+    # We replace it with a custom implementation which adds Atuin support.
+    function PSConsoleHostReadLine {
+        ## 1. Collect the exit code of the previous command.
+
+        # This needs to be done as the first thing because any script run will flush $?.
+        $lastRunStatus = $?
+
+        # Exit statuses are maintained separately for native and PowerShell commands, this needs to be taken into account.
+        $exitCode = if ($lastRunStatus) { 0 } elseif ($global:LASTEXITCODE) { $global:LASTEXITCODE } else { 1 }
+
+        ## 2. Report the status of the previous command to Atuin (atuin history end).
+
+        if ($script:atuinHistoryId) {
+            # The duration is not recorded in old PowerShell versions, let Atuin handle it. $null arguments are ignored.
+            $duration = (Get-History -Count 1).Duration.Ticks * 100
+            $durationArg = if ($duration) { "--duration=$duration" } else { $null }
+
+            atuin history end --exit=$exitCode $durationArg -- $script:atuinHistoryId | Out-Null
+
+            $global:LASTEXITCODE = $exitCode
+            $script:atuinHistoryId = $null
+        }
+
+        ## 3. Read the next command line to execute.
+
+        # PSConsoleHostReadLine implementation from PSReadLine, adjusted to support old versions.
+        Microsoft.PowerShell.Core\Set-StrictMode -Off
+
+        $line = if ($script:hasExpectedReadLineOverload) {
+            # When the overload we expect is available, we can pass $lastRunStatus to it.
+            [Microsoft.PowerShell.PSConsoleReadLine]::ReadLine($Host.Runspace, $ExecutionContext, [System.Threading.CancellationToken]::None, $lastRunStatus)
+        } else {
+            # Either PSReadLine is older than v2.2.0-beta3, or maybe newer than we expect, so use the function from PSReadLine as-is.
+            & $script:previousPSConsoleHostReadLine
+        }
+
+        ## 4. Report the next command line to Atuin (atuin history start).
+
+        # PowerShell doesn't handle double quotes in native command line arguments the same way depending on its version,
+        # and the value of $PSNativeCommandArgumentPassing - see the about_Parsing help page which explains the breaking changes.
+        # This makes it unreliable, so we go through an environment variable, which should always be consistent across versions.
+        try {
+            $env:ATUIN_COMMAND_LINE = $line
+            $script:atuinHistoryId = atuin history start --command-from-env
+        }
+        finally {
+            $env:ATUIN_COMMAND_LINE = $null
+        }
+
+        return $line
+    }
+
+    function Invoke-AtuinSearch {
+        param([string]$ExtraArgs = "")
+
+        $previousOutputEncoding = [System.Console]::OutputEncoding
+        $resultFile = New-TemporaryFile
+
+        try {
+            [System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8
+
+            $query = $null
+            [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$query, [ref]$null)
+
+            # Atuin is started through Start-Process to avoid interfering with the current shell.
+            $env:ATUIN_SHELL = "powershell"
+            $env:ATUIN_QUERY = $query
+            $argString = "search -i --result-file ""$resultFile"" $ExtraArgs"
+            Start-Process -PassThru -NoNewWindow -FilePath atuin -ArgumentList $argString | Wait-Process
+            $suggestion = (Get-Content -Raw $resultFile -Encoding UTF8 | Out-String).Trim()
+
+            # If no shell prompt offset is set, initialize it from the current prompt line count.
+            if ($null -eq $env:ATUIN_POWERSHELL_PROMPT_OFFSET) {
+                try {
+                    $promptLines = (& $Function:prompt | Out-String | Measure-Object -Line).Lines
+                    $env:ATUIN_POWERSHELL_PROMPT_OFFSET = -1 * ($promptLines - 1)
+                }
+                catch {
+                    $env:ATUIN_POWERSHELL_PROMPT_OFFSET = 0
+                }
+            }
+
+            # PSReadLine maintains its own cursor position, which will no longer be valid if Atuin scrolls the display in inline mode.
+            # Fortunately, InvokePrompt can receive a new Y position and reset the internal state.
+            $y = $Host.UI.RawUI.CursorPosition.Y + [int]$env:ATUIN_POWERSHELL_PROMPT_OFFSET
+            $y = [System.Math]::Max([System.Math]::Min($y, [System.Console]::BufferHeight - 1), 0)
+            [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt($null, $y)
+
+            if ($suggestion -eq "") {
+                # The previous input was already rendered by InvokePrompt
+                return
+            }
+
+            $acceptPrefix = "__atuin_accept__:"
+
+            if ( $suggestion.StartsWith($acceptPrefix)) {
+                [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
+                [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion.Substring($acceptPrefix.Length))
+                [Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine()
+            } else {
+                [Microsoft.PowerShell.PSConsoleReadLine]::RevertLine()
+                [Microsoft.PowerShell.PSConsoleReadLine]::Insert($suggestion)
+            }
+        }
+        finally {
+            [System.Console]::OutputEncoding = $previousOutputEncoding
+            $env:ATUIN_SHELL = $null
+            $env:ATUIN_QUERY = $null
+            Remove-Item $resultFile
+        }
+    }
+
+    function Enable-AtuinSearchKeys {
+        param([bool]$CtrlR = $true, [bool]$UpArrow = $true)
+
+        if ($CtrlR) {
+            Set-PSReadLineKeyHandler -Chord "Ctrl+r" -BriefDescription "Runs Atuin search" -ScriptBlock {
+                Invoke-AtuinSearch
+            }
+        }
+
+        if ($UpArrow) {
+            Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock {
+                $line = $null
+                [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$line, [ref]$null)
+
+                if (!$line.Contains("`n")) {
+                    Invoke-AtuinSearch -ExtraArgs "--shell-up-key-binding"
+                } else {
+                    [Microsoft.PowerShell.PSConsoleReadLine]::PreviousLine()
+                }
+            }
+        }
+    }
+
+    $ExecutionContext.SessionState.Module.OnRemove += {
+        $env:ATUIN_SESSION = $null
+        $Function:PSConsoleHostReadLine = $script:previousPSConsoleHostReadLine
+    }
+
+    Export-ModuleMember -Function @("Enable-AtuinSearchKeys", "PSConsoleHostReadLine")
+} | Import-Module -Global