launcher/src/tmux/session.rs

use anyhow::Result;
use std::time::Duration;
use tokio::time::sleep;
use tracing::{info, warn};

use crate::config::{Cluster, Config, Node, Recipe};
use crate::ssh::*;
use crate::tmux::{SessionHealth, TmuxManager};

impl TmuxManager {
    /// Executes a list of commands via SSH on a node to create a tmux session.
    pub async fn create_session(
        &self,
        config: &Config,
        node: &Node,
        commands: &[String],
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!("  ssh {}@{}:", config.ssh_user, node.ssh_ip);
            for cmd in commands {
                info!("    {}", cmd);
            }
            return Ok(());
        }

        let script = commands.join(" && ");
        let output = run_ssh(config, node, script, dry_run).await?;

        if output.trim().is_empty() {
            info!("  ✅ tmux session created on {}", node.name);
        } else {
            info!("  tmux session output on {}: {}", node.name, output.trim());
        }

        Ok(())
    }

    /// Checks whether a tmux session exists on the given node.
    pub async fn check_session(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<bool> {
        if dry_run {
            info!(
                "  ssh {}@{}: tmux list-sessions 2>/dev/null | grep '{}'",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(true);
        }

        let cmd = format!(
            "tmux list-sessions 2>/dev/null | grep -c '{}'",
            session_name
        );
        let output = run_ssh(config, node, cmd, dry_run).await?;

        let found = output.trim().parse::<usize>().unwrap_or(0) > 0;
        if found {
            info!("  Session '{}' found on {}", session_name, node.name);
        } else {
            warn!("  Session '{}' NOT found on {}", session_name, node.name);
        }

        Ok(found)
    }

    /// Performs a deep health check on a tmux session.
    pub async fn deep_check_session(
        &self,
        config: &Config,
        node: &Node,
        _cluster: &Cluster,
        session_name: &str,
        recipe: &Recipe,
        dry_run: bool,
    ) -> Result<SessionHealth> {
        if dry_run {
            return Ok(SessionHealth {
                session_exists: true,
                process_alive: true,
                gpu_ok: true,
                api_ok: true,
            });
        }

        // 1. Check session exists
        let exists = self
            .check_session(config, node, session_name, dry_run)
            .await?;

        // 2. Check process is alive inside tmux
        let process_alive = if exists {
            let pid_cmd = format!(
                "tmux list-panes -t {}:0 -F '#{{pane_pid}}' | head -1",
                session_name
            );
            match run_ssh(config, node, pid_cmd, dry_run).await {
                Ok(pid_str) => {
                    let pid = pid_str.trim();
                    if !pid.is_empty() && pid != "0" {
                        let ps_cmd = format!("ps -p {} -o pid= 2>/dev/null", pid);
                        run_ssh(config, node, ps_cmd, dry_run)
                            .await
                            .map(|o| !o.trim().is_empty())
                            .unwrap_or(false)
                    } else {
                        false
                    }
                }
                Err(_) => false,
            }
        } else {
            false
        };

        // 3. Check GPU is available
        let gpu_ok = match run_ssh(
            config,
            node,
            "nvidia-smi --query-gpu=index,utilization.gpu,memory.used --format=csv,noheader 2>&1 | head -5"
                .to_string(),
            dry_run,
        )
        .await
        {
            Ok(out) => !out.contains("failed") && !out.contains("error") && !out.trim().is_empty(),
            Err(_) => false,
        };

        // 4. Check API endpoint (head node only)
        let api_ok = if node.rank == 0 && process_alive {
            let api_check = format!(
                "curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{}/v1/models 2>/dev/null || echo 'fail'",
                recipe.port
            );
            match run_ssh(config, node, api_check, dry_run).await {
                Ok(out) => {
                    let code = out.trim();
                    code == "200" || !code.contains("fail")
                }
                Err(_) => false,
            }
        } else {
            node.rank != 0 || process_alive
        };

        Ok(SessionHealth {
            session_exists: exists,
            process_alive,
            gpu_ok,
            api_ok,
        })
    }

    /// Gracefully stops a tmux session with timeout escalation.
    pub async fn graceful_stop(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!(
                "  ssh {}@{}: graceful stop of session '{}'",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(());
        }

        info!(
            "  Sending Ctrl+C to session '{}' on {}",
            session_name, node.name
        );
        let _ = run_ssh(
            config,
            node,
            format!("tmux send-keys -t {}:0 C-c", session_name),
            dry_run,
        )
        .await;

        let mut exited = false;
        for i in 0..15 {
            sleep(Duration::from_secs(2)).await;
            let check_cmd = format!(
                "tmux list-panes -t {}:0 -F '#{{pane_pid}}' 2>/dev/null | head -1",
                session_name
            );
            match run_ssh(config, node, check_cmd, dry_run).await {
                Ok(pid_str) => {
                    let pid = pid_str.trim();
                    if pid.is_empty() || pid == "0" || pid == "No" {
                        exited = true;
                        info!(
                            "  Session '{}' exited gracefully (attempt {})",
                            session_name,
                            i + 1
                        );
                        break;
                    }
                    let ps_cmd = format!("ps -p {} -o pid= 2>/dev/null", pid);
                    match run_ssh(config, node, ps_cmd, dry_run).await {
                        Ok(ps_out) if ps_out.trim().is_empty() => {
                            exited = true;
                            info!(
                                "  Session '{}' process exited (attempt {})",
                                session_name,
                                i + 1
                            );
                            break;
                        }
                        _ => continue,
                    }
                }
                Err(_) => {
                    exited = true;
                    break;
                }
            }
        }

        if !exited {
            warn!(
                "  Session '{}' did not exit gracefully, force killing...",
                session_name
            );
            let _ = run_ssh(
                config,
                node,
                format!("tmux kill-session -t {}", session_name),
                dry_run,
            )
            .await;
        }

        Ok(())
    }

    /// Streams the content of a tmux pane, optionally following new output.
    pub async fn stream_pane(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        window: u32,
        follow: bool,
        dry_run: bool,
    ) -> Result<String> {
        if dry_run {
            if follow {
                info!(
                    "  ssh {}@{}: tail -f (streaming from session '{}' window {})",
                    config.ssh_user, node.ssh_ip, session_name, window
                );
            } else {
                info!(
                    "  ssh {}@{}: tmux capture-pane -t {}:{} -p -S -100",
                    config.ssh_user, node.ssh_ip, session_name, window
                );
            }
            return Ok(String::new());
        }

        let cmd = if follow {
            format!(
                "tmux pipe-pane -o -t {}:{} 'cat > /tmp/sm121_pipe_{}_{}.log' && tail -f /tmp/sm121_pipe_{}_{}.log",
                session_name,
                window,
                session_name.replace('-', "_"),
                window,
                session_name.replace('-', "_"),
                window
            )
        } else {
            format!(
                "tmux capture-pane -t {}:{} -p -S -200 2>/dev/null || echo 'no output'",
                session_name, window
            )
        };

        run_ssh_with_retry(config, node, cmd, dry_run, 1).await
    }

    /// Attaches to a tmux session on a node via SSH (interactive).
    pub async fn attach_session(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!(
                "  Would attach: ssh -t {}@{} tmux attach-session -t {}",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(());
        }

        info!(
            "  Attaching to '{}' on {}... Press Ctrl+B D to detach.",
            session_name, node.name
        );
        info!("  (russh does not support interactive PTY. Use the following command directly:)");
        println!(
            "  ssh -t {}@{} tmux attach-session -t {}",
            config.ssh_user, node.ssh_ip, session_name
        );
        println!();
        info!("  Streaming pane content instead:");
        let output = self
            .stream_pane(config, node, session_name, 0, true, dry_run)
            .await?;
        println!("{}", output);

        Ok(())
    }
}