ZIm/crates/project/src/terminals.rs

use crate::{Project, ProjectPath};
use anyhow::Result;
use collections::HashMap;
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, Task, WeakEntity};
use itertools::Itertools;
use language::LanguageName;
use remote::ssh_session::SshArgs;
use settings::{Settings, SettingsLocation};
use smol::channel::bounded;
use std::{
    borrow::Cow,
    path::{Path, PathBuf},
    sync::Arc,
};
use task::{DEFAULT_REMOTE_SHELL, Shell, ShellBuilder, SpawnInTerminal};
use terminal::{
    TaskState, TaskStatus, Terminal, TerminalBuilder,
    terminal_settings::{self, ActivateScript, TerminalSettings, VenvSettings},
};
use util::paths::{PathStyle, RemotePathBuf};

pub struct Terminals {
    pub(crate) local_handles: Vec<WeakEntity<terminal::Terminal>>,
}

/// Terminals are opened either for the users shell, or to run a task.

#[derive(Debug)]
pub enum TerminalKind {
    /// Run a shell at the given path (or $HOME if None)
    Shell(Option<PathBuf>),
    /// Run a task.
    Task(SpawnInTerminal),
}

/// SshCommand describes how to connect to a remote server
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SshCommand {
    pub arguments: Vec<String>,
}

impl SshCommand {
    pub fn add_port_forwarding(&mut self, local_port: u16, host: String, remote_port: u16) {
        self.arguments.push("-L".to_string());
        self.arguments
            .push(format!("{}:{}:{}", local_port, host, remote_port));
    }
}

pub struct SshDetails {
    pub host: String,
    pub ssh_command: SshCommand,
    pub envs: Option<HashMap<String, String>>,
    pub path_style: PathStyle,
}

impl Project {
    pub fn active_project_directory(&self, cx: &App) -> Option<Arc<Path>> {
        let worktree = self
            .active_entry()
            .and_then(|entry_id| self.worktree_for_entry(entry_id, cx))
            .into_iter()
            .chain(self.worktrees(cx))
            .find_map(|tree| tree.read(cx).root_dir());
        worktree
    }

    pub fn first_project_directory(&self, cx: &App) -> Option<PathBuf> {
        let worktree = self.worktrees(cx).next()?;
        let worktree = worktree.read(cx);
        if worktree.root_entry()?.is_dir() {
            Some(worktree.abs_path().to_path_buf())
        } else {
            None
        }
    }

    pub fn ssh_details(&self, cx: &App) -> Option<SshDetails> {
        if let Some(ssh_client) = &self.ssh_client {
            let ssh_client = ssh_client.read(cx);
            if let Some((SshArgs { arguments, envs }, path_style)) = ssh_client.ssh_info() {
                return Some(SshDetails {
                    host: ssh_client.connection_options().host.clone(),
                    ssh_command: SshCommand { arguments },
                    envs,
                    path_style,
                });
            }
        }

        return None;
    }

    pub fn create_terminal(
        &mut self,
        kind: TerminalKind,
        cx: &mut Context<Self>,
    ) -> Task<Result<Entity<Terminal>>> {
        let path: Option<Arc<Path>> = match &kind {
            TerminalKind::Shell(path) => path.as_ref().map(|path| Arc::from(path.as_ref())),
            TerminalKind::Task(spawn_task) => {
                if let Some(cwd) = &spawn_task.cwd {
                    Some(Arc::from(cwd.as_ref()))
                } else {
                    self.active_project_directory(cx)
                }
            }
        };

        let mut settings_location = None;
        if let Some(path) = path.as_ref() {
            if let Some((worktree, _)) = self.find_worktree(path, cx) {
                settings_location = Some(SettingsLocation {
                    worktree_id: worktree.read(cx).id(),
                    path,
                });
            }
        }
        let venv = TerminalSettings::get(settings_location, cx)
            .detect_venv
            .clone();

        cx.spawn(async move |project, cx| {
            let python_venv_directory = if let Some(path) = path {
                match project.upgrade() {
                    Some(project) => {
                        let venv_dir =
                            Self::python_venv_directory(project.clone(), path, venv.clone(), cx)
                                .await?;
                        match venv_dir {
                            Some(venv_dir)
                                if project
                                    .update(cx, |project, cx| {
                                        project.resolve_abs_path(
                                            &venv_dir
                                                .join(BINARY_DIR)
                                                .join("activate")
                                                .to_string_lossy(),
                                            cx,
                                        )
                                    })?
                                    .await
                                    .is_some_and(|m| m.is_file()) =>
                            {
                                Some((venv_dir, venv))
                            }
                            _ => None,
                        }
                    }
                    None => None,
                }
            } else {
                None
            };
            // todo lw: move all this out of here?
            let startup_script = move |shell: &_| {
                let Some((python_venv_directory, venv)) = &python_venv_directory else {
                    return None;
                };
                // todo: If this returns `None` we may don't have a venv, but we may still have a python toolchiain
                // we should modify the PATH here still
                let venv_settings = venv.as_option()?;
                // todo: handle ssh!
                let script_kind = if venv_settings.activate_script
                    == terminal_settings::ActivateScript::Default
                {
                    match shell {
                        Shell::Program(program) => ActivateScript::by_shell(program),
                        Shell::WithArguments {
                            program,
                            args: _,
                            title_override: _,
                        } => ActivateScript::by_shell(program),
                        Shell::System => None,
                    }
                } else {
                    Some(venv_settings.activate_script)
                }
                .or_else(|| ActivateScript::by_env());
                let script_kind = match script_kind {
                    Some(activate) => activate,
                    None => ActivateScript::Default,
                };

                let activate_script_name = match script_kind {
                    terminal_settings::ActivateScript::Default
                    | terminal_settings::ActivateScript::Pyenv => "activate",
                    terminal_settings::ActivateScript::Csh => "activate.csh",
                    terminal_settings::ActivateScript::Fish => "activate.fish",
                    terminal_settings::ActivateScript::Nushell => "activate.nu",
                    terminal_settings::ActivateScript::PowerShell => "activate.ps1",
                };

                let activate_keyword = match script_kind {
                    terminal_settings::ActivateScript::Nushell => "overlay use",
                    terminal_settings::ActivateScript::PowerShell => ".",
                    terminal_settings::ActivateScript::Pyenv => "pyenv",
                    terminal_settings::ActivateScript::Default
                    | terminal_settings::ActivateScript::Csh
                    | terminal_settings::ActivateScript::Fish => "source",
                };

                let line_ending = if cfg!(windows) { '\r' } else { '\n' };

                if venv_settings.venv_name.is_empty() {
                    let path = python_venv_directory
                        .join(BINARY_DIR)
                        .join(activate_script_name)
                        .to_string_lossy()
                        .to_string();
                    let quoted = shlex::try_quote(&path).ok()?;
                    Some(format!("{activate_keyword} {quoted}{line_ending}",))
                } else {
                    Some(format!(
                        "{activate_keyword} {activate_script_name} {name}{line_ending}",
                        name = venv_settings.venv_name
                    ))
                }
            };

            project.update(cx, |project, cx| {
                project.create_terminal_with_startup(kind, startup_script, cx)
            })?
        })
    }

    pub fn terminal_settings<'a>(
        &'a self,
        path: &'a Option<PathBuf>,
        cx: &'a App,
    ) -> &'a TerminalSettings {
        let mut settings_location = None;
        if let Some(path) = path.as_ref() {
            if let Some((worktree, _)) = self.find_worktree(path, cx) {
                settings_location = Some(SettingsLocation {
                    worktree_id: worktree.read(cx).id(),
                    path,
                });
            }
        }
        TerminalSettings::get(settings_location, cx)
    }

    pub fn exec_in_shell(&self, command: String, cx: &App) -> std::process::Command {
        let path = self.first_project_directory(cx);
        let ssh_details = self.ssh_details(cx);
        let settings = self.terminal_settings(&path, cx).clone();

        let builder = ShellBuilder::new(ssh_details.is_none(), &settings.shell).non_interactive();
        let (command, args) = builder.build(Some(command), &Vec::new());

        let mut env = self
            .environment
            .read(cx)
            .get_cli_environment()
            .unwrap_or_default();
        env.extend(settings.env);

        match self.ssh_details(cx) {
            Some(SshDetails {
                ssh_command,
                envs,
                path_style,
                ..
            }) => {
                let (command, args) = wrap_for_ssh(
                    &ssh_command,
                    Some((&command, &args)),
                    path.as_deref(),
                    env,
                    path_style,
                );
                let mut command = std::process::Command::new(command);
                command.args(args);
                if let Some(envs) = envs {
                    command.envs(envs);
                }
                command
            }
            None => {
                let mut command = std::process::Command::new(command);
                command.args(args);
                command.envs(env);
                if let Some(path) = path {
                    command.current_dir(path);
                }
                command
            }
        }
    }

    pub fn clone_terminal(
        &mut self,
        terminal: &Entity<Terminal>,
        cx: &mut Context<Self>,
    ) -> Result<Entity<Terminal>> {
        let terminal = terminal.read(cx);
        let working_directory = terminal
            .working_directory()
            .or_else(|| Some(self.active_project_directory(cx)?.to_path_buf()));
        let startup_script = terminal.startup_script.clone();
        self.create_terminal_with_startup(
            TerminalKind::Shell(working_directory),
            |_| {
                // The shell shouldn't change here
                startup_script
            },
            cx,
        )
    }

    pub fn create_terminal_with_startup(
        &mut self,
        kind: TerminalKind,
        startup_script: impl FnOnce(&Shell) -> Option<String> + 'static,
        cx: &mut Context<Self>,
    ) -> Result<Entity<Terminal>> {
        let this = &mut *self;
        let ssh_details = this.ssh_details(cx);
        let path: Option<Arc<Path>> = match &kind {
            TerminalKind::Shell(path) => path.as_ref().map(|path| Arc::from(path.as_ref())),
            TerminalKind::Task(spawn_task) => {
                if let Some(cwd) = &spawn_task.cwd {
                    if ssh_details.is_some() {
                        Some(Arc::from(cwd.as_ref()))
                    } else {
                        let cwd = cwd.to_string_lossy();
                        let tilde_substituted = shellexpand::tilde(&cwd);
                        Some(Arc::from(Path::new(tilde_substituted.as_ref())))
                    }
                } else {
                    this.active_project_directory(cx)
                }
            }
        };

        let is_ssh_terminal = ssh_details.is_some();

        let mut settings_location = None;
        if let Some(path) = path.as_ref() {
            if let Some((worktree, _)) = this.find_worktree(path, cx) {
                settings_location = Some(SettingsLocation {
                    worktree_id: worktree.read(cx).id(),
                    path,
                });
            }
        }
        let settings = TerminalSettings::get(settings_location, cx).clone();

        let (completion_tx, completion_rx) = bounded(1);

        // Start with the environment that we might have inherited from the Zed CLI.
        let mut env = this
            .environment
            .read(cx)
            .get_cli_environment()
            .unwrap_or_default();
        // Then extend it with the explicit env variables from the settings, so they take
        // precedence.
        env.extend(settings.env);

        let local_path = if is_ssh_terminal { None } else { path.clone() };

        let (spawn_task, shell) = match kind {
            TerminalKind::Shell(_) => {
                match ssh_details {
                    Some(SshDetails {
                        host,
                        ssh_command,
                        envs,
                        path_style,
                    }) => {
                        log::debug!("Connecting to a remote server: {ssh_command:?}");

                        // Alacritty sets its terminfo to `alacritty`, this requiring hosts to have it installed
                        // to properly display colors.
                        // We do not have the luxury of assuming the host has it installed,
                        // so we set it to a default that does not break the highlighting via ssh.
                        env.entry("TERM".to_string())
                            .or_insert_with(|| "xterm-256color".to_string());

                        let (program, args) =
                            wrap_for_ssh(&ssh_command, None, path.as_deref(), env, path_style);
                        env = HashMap::default();
                        if let Some(envs) = envs {
                            env.extend(envs);
                        }
                        (
                            Option::<TaskState>::None,
                            Shell::WithArguments {
                                program,
                                args,
                                title_override: Some(format!("{} — Terminal", host).into()),
                            },
                        )
                    }
                    None => (None, settings.shell),
                }
            }
            TerminalKind::Task(spawn_task) => {
                let task_state = Some(TaskState {
                    id: spawn_task.id,
                    full_label: spawn_task.full_label,
                    label: spawn_task.label,
                    command_label: spawn_task.command_label,
                    hide: spawn_task.hide,
                    status: TaskStatus::Running,
                    show_summary: spawn_task.show_summary,
                    show_command: spawn_task.show_command,
                    show_rerun: spawn_task.show_rerun,
                    completion_rx,
                });

                env.extend(spawn_task.env);

                match ssh_details {
                    Some(SshDetails {
                        host,
                        ssh_command,
                        envs,
                        path_style,
                    }) => {
                        log::debug!("Connecting to a remote server: {ssh_command:?}");
                        env.entry("TERM".to_string())
                            .or_insert_with(|| "xterm-256color".to_string());

                        let (program, args) = wrap_for_ssh(
                            &ssh_command,
                            spawn_task
                                .command
                                .as_ref()
                                .map(|command| (command, &spawn_task.args)),
                            path.as_deref(),
                            env,
                            path_style,
                        );
                        env = HashMap::default();
                        if let Some(envs) = envs {
                            env.extend(envs);
                        }
                        (
                            task_state,
                            Shell::WithArguments {
                                program,
                                args,
                                title_override: Some(format!("{} — Terminal", host).into()),
                            },
                        )
                    }
                    None => {
                        let shell = if let Some(program) = spawn_task.command {
                            Shell::WithArguments {
                                program,
                                args: spawn_task.args,
                                title_override: None,
                            }
                        } else {
                            Shell::System
                        };
                        (task_state, shell)
                    }
                }
            }
        };

        let startup_script = startup_script(&shell);
        TerminalBuilder::new(
            local_path.map(|path| path.to_path_buf()),
            startup_script,
            spawn_task,
            shell,
            env,
            settings.cursor_shape.unwrap_or_default(),
            settings.alternate_scroll,
            settings.max_scroll_history_lines,
            is_ssh_terminal,
            cx.entity_id().as_u64(),
            completion_tx,
            cx,
        )
        .map(|builder| {
            let terminal_handle = cx.new(|cx| builder.subscribe(cx));

            this.terminals
                .local_handles
                .push(terminal_handle.downgrade());

            let id = terminal_handle.entity_id();
            cx.observe_release(&terminal_handle, move |project, _terminal, cx| {
                let handles = &mut project.terminals.local_handles;

                if let Some(index) = handles
                    .iter()
                    .position(|terminal| terminal.entity_id() == id)
                {
                    handles.remove(index);
                    cx.notify();
                }
            })
            .detach();

            terminal_handle
        })
    }

    async fn python_venv_directory(
        this: Entity<Self>,
        abs_path: Arc<Path>,
        venv_settings: VenvSettings,
        cx: &mut AsyncApp,
    ) -> Result<Option<PathBuf>> {
        let Some((worktree, relative_path)) =
            this.update(cx, |this, cx| this.find_worktree(&abs_path, cx))?
        else {
            return Ok(None);
        };
        let toolchain = this
            .update(cx, |this, cx| {
                this.active_toolchain(
                    ProjectPath {
                        worktree_id: worktree.read(cx).id(),
                        path: relative_path.into(),
                    },
                    LanguageName::new("Python"),
                    cx,
                )
            })?
            .await;

        if let Some(toolchain) = toolchain {
            let toolchain_path = Path::new(toolchain.path.as_ref());
            return Ok(toolchain_path
                .parent()
                .and_then(|p| Some(p.parent()?.to_path_buf())));
        }

        let Some(venv_settings) = venv_settings.as_option() else {
            return Ok(None);
        };

        let tool = this.update(cx, |this, cx| {
            venv_settings
                .directories
                .iter()
                .map(|name| abs_path.join(name))
                .find(|venv_path| {
                    let bin_path = venv_path.join(BINARY_DIR);
                    this.find_worktree(&bin_path, cx)
                        .and_then(|(worktree, relative_path)| {
                            worktree.read(cx).entry_for_path(&relative_path)
                        })
                        .is_some_and(|entry| entry.is_dir())
                })
        })?;

        if let Some(toolchain_path) = tool {
            return Ok(Some(toolchain_path));
        }

        let r = this.update(cx, move |_, cx| {
            let map: Vec<_> = venv_settings
                .directories
                .iter()
                .map(|name| abs_path.join(name))
                .collect();
            Some(cx.spawn(async move |project, cx| {
                for venv_path in map {
                    let bin_path = venv_path.join(BINARY_DIR);
                    let exists = project
                        .upgrade()?
                        .update(cx, |project, cx| {
                            project.resolve_abs_path(&bin_path.to_string_lossy(), cx)
                        })
                        .ok()?
                        .await
                        .map_or(false, |meta| meta.is_dir());
                    if exists {
                        return Some(venv_path);
                    }
                }
                None
            }))
        })?;
        Ok(match r {
            Some(task) => task.await,
            None => None,
        })
    }

    pub fn local_terminal_handles(&self) -> &Vec<WeakEntity<terminal::Terminal>> {
        &self.terminals.local_handles
    }
}

const BINARY_DIR: &str = if cfg!(target_os = "windows") {
    "Scripts"
} else {
    "bin"
};

pub fn wrap_for_ssh(
    ssh_command: &SshCommand,
    command: Option<(&String, &Vec<String>)>,
    path: Option<&Path>,
    env: HashMap<String, String>,
    path_style: PathStyle,
) -> (String, Vec<String>) {
    let to_run = if let Some((command, args)) = command {
        // DEFAULT_REMOTE_SHELL is '"${SHELL:-sh}"' so must not be escaped
        let command: Option<Cow<str>> = if command == DEFAULT_REMOTE_SHELL {
            Some(command.into())
        } else {
            shlex::try_quote(command).ok()
        };
        let args = args.iter().filter_map(|arg| shlex::try_quote(arg).ok());
        command.into_iter().chain(args).join(" ")
    } else {
        "exec ${SHELL:-sh} -l".to_string()
    };

    let mut env_changes = String::new();
    for (k, v) in env.iter() {
        if let Some((k, v)) = shlex::try_quote(k).ok().zip(shlex::try_quote(v).ok()) {
            env_changes.push_str(&format!("{k}={v} "));
        }
    }

    let commands = if let Some(path) = path {
        let path = RemotePathBuf::new(path.to_path_buf(), path_style).to_string();
        // shlex will wrap the command in single quotes (''), disabling ~ expansion,
        // replace ith with something that works
        let tilde_prefix = "~/";
        if path.starts_with(tilde_prefix) {
            let trimmed_path = path
                .trim_start_matches("/")
                .trim_start_matches("~")
                .trim_start_matches("/");

            format!("cd \"$HOME/{trimmed_path}\"; {env_changes} {to_run}")
        } else {
            format!("cd \"{path}\"; {env_changes} {to_run}")
        }
    } else {
        format!("cd; {env_changes} {to_run}")
    };
    let shell_invocation = format!("sh -c {}", shlex::try_quote(&commands).unwrap());

    let program = "ssh".to_string();
    let mut args = ssh_command.arguments.clone();

    args.push("-t".to_string());
    args.push(shell_invocation);
    (program, args)
}