use std::path::{Path, PathBuf};
use tracing::Instrument;

use super::AutoStart;
use super::util::{
    discover_daemon_probes, format_age, format_timestamp, resolve_session_id, server_request,
};

fn client_config(
    name: &str,
    session_id: u32,
    ctl_path: &Path,
    settings: &gritty::config::SessionSettings,
    server_id: u64,
) -> gritty::client::ClientConfig {
    gritty::client::ClientConfig {
        session: name.to_string(),
        session_id,
        ctl_path: ctl_path.to_path_buf(),
        env_vars: vec![],
        no_escape: settings.no_escape,
        forward_agent: settings.forward_agent,
        forward_open: settings.forward_open,
        oauth_redirect: settings.oauth_redirect,
        oauth_timeout: settings.oauth_timeout,
        heartbeat_interval: settings.heartbeat_interval,
        heartbeat_timeout: settings.heartbeat_timeout,
        client_name: settings.client_name.clone(),
        expected_server_id: server_id,
        device_id: gritty::get_or_create_device_id(),
    }
}

/// Boolean options for [`connect_session`], grouped into one struct because
/// seven adjacent bare `bool` parameters are silently transposable at the call
/// site (mirrors the existing `SessionSettings` grouping pattern).
pub(crate) struct ConnectFlags {
    pub(crate) detach: bool,
    pub(crate) no_create: bool,
    pub(crate) force: bool,
    pub(crate) pick: bool,
    pub(crate) no_pick: bool,
    pub(crate) new_session: bool,
    pub(crate) wait: bool,
}

pub(crate) async fn connect_session(
    session: Option<String>,
    command: Option<String>,
    flags: ConnectFlags,
    settings: gritty::config::SessionSettings,
    ctl_path: PathBuf,
    auto_start_mode: AutoStart,
) -> anyhow::Result<()> {
    use futures_util::{SinkExt, StreamExt};
    use gritty::protocol::{Frame, FrameCodec};
    use tokio_util::codec::Framed;

    let ConnectFlags { detach, no_create, force, pick, no_pick, new_session, wait } = flags;

    // Resolve the session name. An explicit name or the interactive picker
    // can be resolved up front; -n/--new must wait until connect_or_start has
    // proven the socket reachable.
    let (preliminary_name, picker_force) = match &session {
        Some(name) => (Some(name.clone()), false),
        None if new_session => (None, false),
        None => {
            let (n, _picked, pf) = pick_session(pick, no_pick, &ctl_path).await;
            (Some(n), pf)
        }
    };
    let force = force || picker_force;
    let session_command = command.unwrap_or_default();

    let (stream, _auto_started) =
        super::util::connect_or_start(&ctl_path, &auto_start_mode, wait).await?;

    // -n/--new: resolve the auto-name only now. Resolving it before
    // connect_or_start collapses to "default" whenever the local tunnel
    // socket is down (e.g. after a laptop reboot), and the attach-first path
    // would then silently attach to a pre-existing remote `default` instead
    // of creating a fresh session-N -- violating the --new contract.
    let name = match preliminary_name {
        Some(n) => n,
        None => auto_new_session_name(&ctl_path).await,
    };

    let mut framed = Framed::new(stream, FrameCodec);
    let info = gritty::handshake(&mut framed, gritty::get_or_create_device_id()).await?;
    gritty::require_matched_version(&info)?;
    let server_id = info.server_id;

    // Carry current terminal size so the server can resize the PTY before
    // replaying scrollback/ring buffer on reconnect. Zero for probe-only.
    let (attach_cols, attach_rows) =
        if detach { (0, 0) } else { gritty::client::get_terminal_size() };

    // Try attach first
    framed
        .send(Frame::Attach {
            session: name.clone(),
            client_name: settings.client_name.clone(),
            force,
            no_replay: detach,
            cols: attach_cols,
            rows: attach_rows,
            attach_token: 0,
            // Explicit connect: no prior stream position. attach_token == 0
            // already signals "fresh viewer" to the server, so it replays
            // scrollback context rather than an incremental resume.
            rendered_offset: 0,
            line_dirty: false,
        })
        .await?;

    match Frame::expect_from(framed.next().await)? {
        Frame::Ok if detach => {
            // Probe succeeded (connect -d): existence confirmed, not attaching.
            eprintln!("\x1b[32m\u{25b8} session {name} exists (not attaching, -d)\x1b[0m");
            return Ok(());
        }
        Frame::AttachAck { token: _, session_id } => {
            eprintln!("\x1b[32m\u{25b8} attached {name}\x1b[0m");
            let client_span = tracing::info_span!("client", session = %name, session_id);
            let code = gritty::client::run(
                framed,
                client_config(&name, session_id, &ctl_path, &settings, server_id),
            )
            .instrument(client_span)
            .await?;
            std::process::exit(code);
        }
        Frame::Error { code: gritty::protocol::ErrorCode::NoSuchSession, .. } => {
            if name == "-" {
                // `-` means "last-attached session"; creating one named `-`
                // is reserved and would fail with a misleading error.
                let host = host_from_ctl_path(&ctl_path);
                anyhow::bail!("no previously-attached session on {host}");
            }
            if no_create {
                anyhow::bail!("no such session: {name}");
            }
            // Fall through to create
        }
        Frame::Error { code: gritty::protocol::ErrorCode::AlreadyAttached, message, .. } => {
            let host = host_from_ctl_path(&ctl_path);
            eprintln!("error: {message}");
            eprintln!("  gritty connect {host}:{name} --force   to take over");
            std::process::exit(1);
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }

    // Create new session -- need a fresh connection since the previous one
    // was consumed by the failed attach
    let (stream, _) = super::util::connect_or_start(&ctl_path, &auto_start_mode, wait).await?;
    let mut framed = Framed::new(stream, FrameCodec);
    let info = gritty::handshake(&mut framed, gritty::get_or_create_device_id()).await?;
    gritty::require_matched_version(&info)?;
    let server_id = info.server_id;
    // Get terminal size for initial PTY dimensions
    let (cols, rows) = crossterm::terminal::size().unwrap_or((0, 0));
    framed
        .send(Frame::NewSession {
            name: name.clone(),
            command: session_command,
            cwd: String::new(),
            cols,
            rows,
            client_name: settings.client_name.clone(),
        })
        .await?;

    match Frame::expect_from(framed.next().await)? {
        Frame::SessionCreated { id } => {
            eprintln!("\x1b[32m\u{25b8} session {name}\x1b[0m");

            // Daemon emits SessionCreated immediately followed by AttachAck.
            // The token echoes the device_id (client ignores it -- ownership
            // is tracked by the persistent device_id, not an ephemeral token).
            match Frame::expect_from(framed.next().await)? {
                Frame::AttachAck { .. } => {}
                other => anyhow::bail!("unexpected response from server: {other:?}"),
            }

            // Send Env immediately so the server's 2s deferred-spawn deadline
            // is satisfied before -d returns or alert_detached_sessions runs
            // its own multi-RTT round-trip.
            framed.send(Frame::Env { vars: gritty::collect_env_vars() }).await?;

            if detach {
                return Ok(());
            }

            // Alert about other detached sessions
            alert_detached_sessions(&name, &ctl_path).await;

            let client_span = tracing::info_span!("client", session = %name, session_id = id);
            let code = gritty::client::run(
                framed,
                client_config(&name, id, &ctl_path, &settings, server_id),
            )
            .instrument(client_span)
            .await?;
            std::process::exit(code);
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }
}

/// Resolve session name when none was explicitly given.
/// Returns `(name, picked)` where `picked` is true if the user chose interactively.
/// Resolve `-n`/`--new` to the next auto-named session slot. Uses the same
/// rule the picker applies: `default` if unused, otherwise the first free
/// `session-N`. Falls back to `default` if the server can't be reached --
/// the connect flow auto-starts the daemon and will re-attempt anyway.
async fn auto_new_session_name(ctl_path: &Path) -> String {
    use gritty::protocol::Frame;

    let sessions = match server_request(ctl_path, Frame::ListSessions).await {
        Ok(Frame::SessionInfo { sessions }) => sessions,
        _ => return "default".to_string(),
    };
    suggest_name(&build_rows(&sessions))
}

async fn pick_session(pick: bool, no_pick: bool, ctl_path: &Path) -> (String, bool, bool) {
    use gritty::protocol::Frame;

    if no_pick {
        return ("default".to_string(), false, false);
    }

    let sessions = match server_request(ctl_path, Frame::ListSessions).await {
        Ok(Frame::SessionInfo { sessions }) => sessions,
        _ => return ("default".to_string(), false, false),
    };

    if sessions.is_empty() {
        return ("default".to_string(), false, false);
    }

    let host = host_from_ctl_path(ctl_path);

    if pick {
        return pick_or_list(&host, &sessions, ctl_path).await;
    }

    let detached: Vec<_> = sessions.iter().filter(|s| !s.attached).collect();

    // One session, detached: attach directly
    if sessions.len() == 1 && detached.len() == 1 {
        return (session_display_name(&sessions[0]), false, false);
    }

    // Multiple sessions, exactly one detached: attach to the detached one
    if detached.len() == 1 {
        return (session_display_name(detached[0]), false, false);
    }

    // Ambiguous (multiple detached) or all attached: show picker
    pick_or_list(&host, &sessions, ctl_path).await
}

fn session_display_name(s: &gritty::protocol::SessionEntry) -> String {
    if s.name.is_empty() { s.id.to_string() } else { s.name.clone() }
}

/// Suggest a name for a new session. Returns "default" if unused, otherwise "session-N".
fn suggest_name(rows: &[Row]) -> String {
    if !rows.iter().any(|r| r.name == "default") {
        return "default".to_string();
    }
    for n in 2.. {
        let candidate = format!("session-{n}");
        if !rows.iter().any(|r| r.name == candidate) {
            return candidate;
        }
    }
    unreachable!()
}

fn print_session_list(host: &str, sessions: &[gritty::protocol::SessionEntry]) {
    if sessions.len() == 1 {
        eprintln!("error: session on {host} is already attached:");
    } else {
        eprintln!("error: multiple sessions on {host} -- specify one:");
    }
    for s in sessions {
        let name = session_display_name(s);
        let suffix = if s.attached { "     (attached)" } else { "" };
        eprintln!("  gritty connect {host}:{name}{suffix}");
    }
}

/// Show picker (TUI if stderr is a TTY, static list otherwise).
/// Returns `(name, true)` for interactive pick, or exits on abort/non-TTY.
async fn pick_or_list(
    host: &str,
    sessions: &[gritty::protocol::SessionEntry],
    ctl_path: &Path,
) -> (String, bool, bool) {
    if std::io::IsTerminal::is_terminal(&std::io::stderr()) {
        match tui_pick_session(host, sessions, ctl_path).await {
            Some((name, force)) => (name, true, force),
            None => std::process::exit(1),
        }
    } else {
        print_session_list(host, sessions);
        std::process::exit(1);
    }
}

struct Row {
    name: String,
    attached: bool,
    age: String,
    cmd: String,
    cwd: String,
    client: String,
    hotkey: Option<char>, // '1'-'9' for first 9 rows
}

fn build_rows(sessions: &[gritty::protocol::SessionEntry]) -> Vec<Row> {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    let home = std::env::var("HOME").unwrap_or_default();
    sessions
        .iter()
        .enumerate()
        .map(|(i, s)| Row {
            name: session_display_name(s),
            attached: s.attached,
            age: format_age(now, s.created_at),
            cmd: s.foreground_cmd.clone(),
            cwd: shorten_home(&s.cwd, &home),
            client: s.client_name.clone(),
            hotkey: if i < 9 { Some((b'1' + i as u8) as char) } else { None },
        })
        .collect()
}

fn shorten_home(path: &str, home: &str) -> String {
    if !home.is_empty() && path.starts_with(home) {
        let rest = &path[home.len()..];
        if rest.is_empty() || rest.starts_with('/') {
            return format!("~{rest}");
        }
    }
    path.to_string()
}

/// Interactive session picker. Returns selected session name, or None on abort.
async fn tui_pick_session(
    host: &str,
    sessions: &[gritty::protocol::SessionEntry],
    ctl_path: &Path,
) -> Option<(String, bool)> {
    use crossterm::{
        event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
        terminal,
    };
    use std::io::Write;

    let mut stderr = std::io::stderr();

    // Find first detached session for initial cursor position
    let initial = sessions.iter().position(|s| !s.attached).unwrap_or(0);
    let mut cursor = initial;

    let mut rows = build_rows(sessions);

    enum Mode {
        Pick,
        Input { buf: String, cursor_pos: usize, rename_of: Option<String> },
        ConfirmKill { name: String },
    }

    let mut mode = Mode::Pick;
    // Last server error from an in-picker rename/kill, shown until the next
    // attempt. Without it a rejected rename/kill is a silent no-op.
    let mut status: Option<String> = None;
    let mut prev_total_lines: usize = 0;

    struct PickerTermGuard;
    impl Drop for PickerTermGuard {
        fn drop(&mut self) {
            use std::io::Write;
            let _ = write!(std::io::stderr(), "\x1b[?25h");
            let _ = std::io::stderr().flush();
            let _ = crossterm::terminal::disable_raw_mode();
        }
    }
    // Enter raw mode
    let _ = terminal::enable_raw_mode();
    let _ = write!(stderr, "\x1b[?25l"); // hide cursor
    let _term_guard = PickerTermGuard;

    fn prev_char_boundary(s: &str, mut pos: usize) -> usize {
        while pos > 0 {
            pos -= 1;
            if s.is_char_boundary(pos) {
                break;
            }
        }
        pos
    }
    fn next_char_boundary(s: &str, pos: usize) -> usize {
        s[pos..].chars().next().map_or(pos, |c| pos + c.len_utf8())
    }

    let render = |stderr: &mut std::io::Stderr,
                  rows: &[Row],
                  cursor: usize,
                  mode: &Mode,
                  status: Option<&str>,
                  prev_total_lines: usize| {
        let name_w = rows.iter().map(|r| r.name.len()).max().unwrap_or(0).max(3);
        let tag_w = 10; // "(attached)" is 10 chars
        let age_w = rows.iter().map(|r| r.age.len()).max().unwrap_or(0);
        let cmd_w = rows.iter().map(|r| r.cmd.len()).max().unwrap_or(0);
        let client_w = rows.iter().map(|r| r.client.len()).max().unwrap_or(0);
        // header + rows + new-session + hint, plus an optional status line.
        let total_lines = rows.len() + 3 + usize::from(status.is_some());

        // If we drew before, erase old output first
        if prev_total_lines > 0 {
            let _ = write!(stderr, "\x1b[{}A", prev_total_lines);
            for _ in 0..prev_total_lines {
                let _ = write!(stderr, "\x1b[2K\r\n");
            }
            let _ = write!(stderr, "\x1b[{}A", prev_total_lines);
        }

        // Show/hide cursor based on mode
        match mode {
            Mode::Pick | Mode::ConfirmKill { .. } => {
                let _ = write!(stderr, "\x1b[?25l");
            }
            Mode::Input { .. } => {
                let _ = write!(stderr, "\x1b[?25h");
            }
        }

        // Header
        let _ = write!(stderr, "\x1b[36;1mPick a session on {host}:\x1b[0m\r\n");
        for (i, row) in rows.iter().enumerate() {
            let marker = if i == cursor && matches!(mode, Mode::Pick | Mode::ConfirmKill { .. }) {
                "\x1b[32;1m\u{25b8}\x1b[0m"
            } else {
                " "
            };
            let hk = row.hotkey.map_or(String::from("  "), |c| format!("{c})"));

            if i == cursor && matches!(mode, Mode::Pick | Mode::ConfirmKill { .. }) {
                let tag = if row.attached { "(attached)" } else { "" };
                let _ = write!(
                    stderr,
                    "{marker} \x1b[2m{hk}\x1b[0m \x1b[32;1m{:<name_w$}\x1b[0m  {:<tag_w$}  \x1b[32m{:<age_w$}\x1b[0m  \x1b[32m{:<cmd_w$}\x1b[0m  \x1b[32m{:<client_w$}\x1b[0m  \x1b[32m{}\x1b[0m\r\n",
                    row.name, tag, row.age, row.cmd, row.client, row.cwd,
                );
            } else if row.attached {
                let _ = write!(
                    stderr,
                    "{marker} \x1b[2m{hk} {:<name_w$}  {:<tag_w$}  {:<age_w$}  {:<cmd_w$}  {:<client_w$}  {}\x1b[0m\r\n",
                    row.name, "(attached)", row.age, row.cmd, row.client, row.cwd,
                );
            } else {
                let _ = write!(
                    stderr,
                    "{marker} \x1b[2m{hk}\x1b[0m {:<name_w$}  {:<tag_w$}  {:<age_w$}  {:<cmd_w$}  {:<client_w$}  {}\r\n",
                    row.name, "", row.age, row.cmd, row.client, row.cwd,
                );
            }
        }
        // "New session" row / input line
        let suggested = suggest_name(rows);
        match mode {
            Mode::Pick | Mode::ConfirmKill { .. } => {
                let marker = if cursor == rows.len() { "\x1b[32;1m\u{25b8}\x1b[0m" } else { " " };
                if cursor == rows.len() {
                    let _ = write!(
                        stderr,
                        "{marker} \x1b[2m+)\x1b[0m \x1b[32;1mnew session \x1b[2m({suggested})\x1b[0m\r\n"
                    );
                } else {
                    let _ =
                        write!(stderr, "{marker} \x1b[2m+) new session ({suggested})\x1b[0m\r\n");
                }
            }
            Mode::Input { buf, cursor_pos, rename_of } => {
                let prefix = if rename_of.is_some() { "r)" } else { "+)" };
                let (before, after) = buf.split_at(*cursor_pos);
                let cursor_ch = after.chars().next().unwrap_or(' ');
                let rest = if after.is_empty() { "" } else { &after[cursor_ch.len_utf8()..] };
                let _ = write!(
                    stderr,
                    "\x1b[32;1m\u{25b8}\x1b[0m \x1b[2m{prefix}\x1b[0m \x1b[32;1m{before}\x1b[7m{cursor_ch}\x1b[27m{rest}\x1b[0m\r\n"
                );
            }
        }
        // Hint line
        let hints = match mode {
            Mode::Pick => {
                "1-9 jump  enter select  f force  n new  c/+ new (named)  d default  r rename  x kill  esc quit"
                    .to_string()
            }
            Mode::Input { rename_of: Some(_), .. } => "enter rename  esc back".to_string(),
            Mode::Input { .. } => "enter create  esc back".to_string(),
            Mode::ConfirmKill { name } => format!("kill {name}? y/n"),
        };
        let _ = write!(stderr, "\x1b[2m  {hints}\x1b[0m\r\n");
        if let Some(msg) = status {
            let _ = write!(stderr, "\x1b[31m  {msg}\x1b[0m\r\n");
        }
        let _ = stderr.flush();
        total_lines
    };

    // Initial render
    prev_total_lines =
        render(&mut stderr, &rows, cursor, &mode, status.as_deref(), prev_total_lines);

    let result = loop {
        // Poll-based loop so we can yield to the async runtime
        if !event::poll(std::time::Duration::from_millis(100)).unwrap_or(false) {
            tokio::task::yield_now().await;
            continue;
        }
        let Ok(ev) = event::read() else {
            break None;
        };
        match &mut mode {
            Mode::Pick => match ev {
                Event::Key(KeyEvent {
                    code: KeyCode::Up | KeyCode::Char('k'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    cursor = cursor.saturating_sub(1);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down | KeyCode::Char('j'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    if cursor < rows.len() {
                        cursor += 1;
                    }
                }
                Event::Key(KeyEvent { code: KeyCode::Enter, .. }) => {
                    if cursor < rows.len() {
                        break Some((rows[cursor].name.clone(), false));
                    }
                    // On the "new session" row: create immediately with suggested name
                    break Some((suggest_name(&rows), false));
                }
                // Hotkeys 1-9
                Event::Key(KeyEvent {
                    code: KeyCode::Char(ch @ '1'..='9'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    let idx = (ch as u8 - b'1') as usize;
                    if idx < rows.len() {
                        break Some((rows[idx].name.clone(), false));
                    }
                }
                // 'd' -> select or create "default"
                Event::Key(KeyEvent {
                    code: KeyCode::Char('d'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    break Some(("default".to_string(), false));
                }
                // 'n' -> create new session immediately with the suggested name
                Event::Key(KeyEvent {
                    code: KeyCode::Char('n'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    break Some((suggest_name(&rows), false));
                }
                // 'c' or '+' -> new session input, prompting to edit the name
                Event::Key(KeyEvent {
                    code: KeyCode::Char('c') | KeyCode::Char('+'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    let name = suggest_name(&rows);
                    let len = name.len();
                    cursor = rows.len();
                    status = None;
                    mode = Mode::Input { buf: name, cursor_pos: len, rename_of: None };
                }
                // 'x' or Delete -> kill selected session
                Event::Key(KeyEvent {
                    code: KeyCode::Char('x') | KeyCode::Delete,
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    if cursor < rows.len() {
                        status = None;
                        mode = Mode::ConfirmKill { name: rows[cursor].name.clone() };
                    }
                }
                // 'r' -> rename selected session
                Event::Key(KeyEvent {
                    code: KeyCode::Char('r'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    if cursor < rows.len() {
                        let name = rows[cursor].name.clone();
                        let len = name.len();
                        status = None;
                        mode = Mode::Input {
                            buf: name.clone(),
                            cursor_pos: len,
                            rename_of: Some(name),
                        };
                    }
                }
                // 'f' -> force-attach (take over) selected session
                Event::Key(KeyEvent {
                    code: KeyCode::Char('f'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    if cursor < rows.len() {
                        break Some((rows[cursor].name.clone(), true));
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Esc | KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                    ..
                }) => {
                    break None;
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char('c') | KeyCode::Char('g'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    break None;
                }
                _ => continue,
            },
            Mode::ConfirmKill { name } => match ev {
                Event::Key(KeyEvent { code: KeyCode::Char('y'), .. }) => {
                    let kill_name = name.clone();
                    let ctl = ctl_path.to_path_buf();
                    match server_request(
                        &ctl,
                        gritty::protocol::Frame::KillSession { session: kill_name },
                    )
                    .await
                    {
                        Ok(gritty::protocol::Frame::Ok) => status = None,
                        Ok(gritty::protocol::Frame::Error { message, .. }) => {
                            status = Some(format!("kill failed: {message}"));
                        }
                        other => {
                            status = Some(format!("kill failed: unexpected response {other:?}"));
                        }
                    }
                    // Refresh session list
                    if let Ok(gritty::protocol::Frame::SessionInfo { sessions: fresh }) =
                        server_request(&ctl, gritty::protocol::Frame::ListSessions).await
                    {
                        if fresh.is_empty() {
                            break Some(("default".to_string(), false));
                        }
                        rows = build_rows(&fresh);
                        cursor = cursor.min(rows.len().saturating_sub(1));
                    }
                    mode = Mode::Pick;
                }
                _ => {
                    mode = Mode::Pick;
                }
            },
            Mode::Input { buf, cursor_pos, rename_of } => match ev {
                Event::Key(KeyEvent { code: KeyCode::Enter, .. }) => {
                    let new_name = buf.trim().to_string();
                    // clone, don't take: on a rejected rename we stay in Input
                    // mode, and `rename_of` must still say this is a rename.
                    if let Some(old_name) = rename_of.clone() {
                        // Rename mode
                        if new_name.is_empty() || new_name == old_name {
                            mode = Mode::Pick;
                        } else {
                            let ctl = ctl_path.to_path_buf();
                            match server_request(
                                &ctl,
                                gritty::protocol::Frame::RenameSession {
                                    session: old_name,
                                    new_name,
                                },
                            )
                            .await
                            {
                                Ok(gritty::protocol::Frame::Ok) => {
                                    status = None;
                                    // Refresh
                                    if let Ok(gritty::protocol::Frame::SessionInfo {
                                        sessions: fresh,
                                    }) =
                                        server_request(&ctl, gritty::protocol::Frame::ListSessions)
                                            .await
                                    {
                                        rows = build_rows(&fresh);
                                        cursor = cursor.min(rows.len().saturating_sub(1));
                                    }
                                    mode = Mode::Pick;
                                }
                                // Rejected (name collision, invalid name): keep
                                // the user in rename mode with the message.
                                Ok(gritty::protocol::Frame::Error { message, .. }) => {
                                    status = Some(format!("rename failed: {message}"));
                                }
                                other => {
                                    status = Some(format!(
                                        "rename failed: unexpected response {other:?}"
                                    ));
                                }
                            }
                        }
                    } else if new_name.is_empty() {
                        mode = Mode::Pick;
                        cursor = rows.len();
                    } else {
                        break Some((new_name, false));
                    }
                }
                Event::Key(KeyEvent { code: KeyCode::Esc, .. }) => {
                    let back_to_new = rename_of.is_none();
                    status = None;
                    mode = Mode::Pick;
                    if back_to_new {
                        cursor = rows.len();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char('c') | KeyCode::Char('g'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    let back_to_new = rename_of.is_none();
                    status = None;
                    mode = Mode::Pick;
                    if back_to_new {
                        cursor = rows.len();
                    }
                }
                // Readline: Ctrl+A -> beginning
                Event::Key(KeyEvent {
                    code: KeyCode::Char('a'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                })
                | Event::Key(KeyEvent { code: KeyCode::Home, .. }) => {
                    *cursor_pos = 0;
                }
                // Readline: Ctrl+E -> end
                Event::Key(KeyEvent {
                    code: KeyCode::Char('e'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                })
                | Event::Key(KeyEvent { code: KeyCode::End, .. }) => {
                    *cursor_pos = buf.len();
                }
                // Readline: Ctrl+U -> kill line
                Event::Key(KeyEvent {
                    code: KeyCode::Char('u'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    buf.drain(..*cursor_pos);
                    *cursor_pos = 0;
                }
                // Readline: Ctrl+W -> kill word backward
                Event::Key(KeyEvent {
                    code: KeyCode::Char('w'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    let before = &buf[..*cursor_pos];
                    let trimmed = before.trim_end();
                    let word_start = trimmed.rfind(' ').map_or(0, |i| i + 1);
                    buf.drain(word_start..*cursor_pos);
                    *cursor_pos = word_start;
                }
                // Backspace
                Event::Key(KeyEvent { code: KeyCode::Backspace, .. }) => {
                    if *cursor_pos > 0 {
                        let prev = prev_char_boundary(buf, *cursor_pos);
                        buf.remove(prev);
                        *cursor_pos = prev;
                    }
                }
                // Delete
                Event::Key(KeyEvent { code: KeyCode::Delete, .. }) => {
                    if *cursor_pos < buf.len() {
                        buf.remove(*cursor_pos);
                    }
                }
                // Left arrow / Ctrl+B
                Event::Key(KeyEvent { code: KeyCode::Left, .. })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('b'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    *cursor_pos = prev_char_boundary(buf, *cursor_pos);
                }
                // Right arrow / Ctrl+F
                Event::Key(KeyEvent { code: KeyCode::Right, .. })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('f'),
                    modifiers: KeyModifiers::CONTROL,
                    ..
                }) => {
                    *cursor_pos = next_char_boundary(buf, *cursor_pos);
                }
                // Typing
                Event::Key(KeyEvent {
                    code: KeyCode::Char(ch),
                    modifiers: KeyModifiers::NONE | KeyModifiers::SHIFT,
                    ..
                }) => {
                    buf.insert(*cursor_pos, ch);
                    *cursor_pos += ch.len_utf8();
                }
                _ => continue,
            },
        }
        prev_total_lines =
            render(&mut stderr, &rows, cursor, &mode, status.as_deref(), prev_total_lines);
    };

    // Cleanup: erase picker lines (terminal restore handled by PickerTermGuard)
    if prev_total_lines > 0 {
        let _ = write!(stderr, "\x1b[{}A", prev_total_lines);
        for _ in 0..prev_total_lines {
            let _ = write!(stderr, "\x1b[2K\r\n");
        }
        let _ = write!(stderr, "\x1b[{}A", prev_total_lines);
    }
    let _ = stderr.flush();

    result
}

/// Extract a display-friendly host name from a ctl socket path.
fn host_from_ctl_path(ctl_path: &Path) -> String {
    // Tunnel sockets: .../connect-<host>.sock -> host is <host>
    // Local daemon: .../ctl.sock -> host is "local"
    let file = ctl_path.file_stem().and_then(|s| s.to_str()).unwrap_or("local");
    if let Some(host) = file.strip_prefix("connect-") {
        host.to_string()
    } else {
        "local".to_string()
    }
}

/// After creating a new session, show a hint if there are other detached
/// sessions the user might have forgotten about.
async fn alert_detached_sessions(current_name: &str, ctl_path: &Path) {
    use gritty::protocol::Frame;

    let ctl_path_buf = ctl_path.to_path_buf();
    let Ok(resp) = server_request(&ctl_path_buf, Frame::ListSessions).await else {
        return;
    };
    let Frame::SessionInfo { sessions } = resp else {
        return;
    };
    let detached: Vec<_> = sessions
        .iter()
        .filter(|s| {
            !s.attached && {
                let display = if s.name.is_empty() { s.id.to_string() } else { s.name.clone() };
                display != current_name
            }
        })
        .collect();
    if detached.is_empty() {
        return;
    }
    let names: Vec<_> = detached
        .iter()
        .map(|s| if s.name.is_empty() { s.id.to_string() } else { s.name.clone() })
        .collect();
    eprintln!("\x1b[2;33m\u{25b8} detached sessions: {}\x1b[0m", names.join(", "));
}

pub(crate) async fn tail_session(target: String, ctl_path: PathBuf) -> anyhow::Result<i32> {
    use futures_util::{SinkExt, StreamExt};
    use gritty::protocol::Frame;

    // Resolve the target to a numeric id before opening the tail stream so
    // reconnect can reuse that id (the original target string may be `-`
    // or a name that can shift while we're tailing).
    let session_id = resolve_session_id(&ctl_path, &target).await?;

    let (mut framed, info) = super::util::connect_handshaked(&ctl_path, true).await?;
    let server_id = info.server_id;
    framed.send(Frame::Tail { session: session_id.to_string() }).await?;

    match Frame::expect_from(framed.next().await)? {
        Frame::Ok => {
            eprintln!("\x1b[2;33m\u{25b8} tailing {target}\x1b[0m");
            gritty::client::tail(
                session_id,
                framed,
                &ctl_path,
                server_id,
                gritty::get_or_create_device_id(),
            )
            .await
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }
}

pub(crate) async fn rename_session(
    target: String,
    new_name: String,
    ctl_path: PathBuf,
) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    match server_request(
        &ctl_path,
        Frame::RenameSession { session: target.clone(), new_name: new_name.clone() },
    )
    .await?
    {
        Frame::Ok => {
            eprintln!("\x1b[32m\u{25b8} renamed {target} -> {new_name}\x1b[0m");
            Ok(())
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }
}

pub(crate) async fn kill_session(target: String, ctl_path: PathBuf) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    match server_request(&ctl_path, Frame::KillSession { session: target.clone() }).await? {
        Frame::Ok => {
            eprintln!("\x1b[32m\u{25b8} session {target} killed\x1b[0m");
            Ok(())
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }
}

pub(crate) async fn kill_server(ctl_path: PathBuf) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    // Deliberately tolerant of protocol version mismatch: killing the server
    // is the first step of the upgrade recovery ritual, so it MUST work
    // across a mismatched handshake.
    match super::util::server_request_any_version(&ctl_path, Frame::KillServer).await? {
        Frame::Ok => {
            eprintln!("\x1b[32m\u{25b8} server killed\x1b[0m");
            Ok(())
        }
        Frame::Error { message, .. } => anyhow::bail!("{message}"),
        other => anyhow::bail!("unexpected response from server: {other:?}"),
    }
}

/// Orchestrate a full daemon (+ tunnel, for remote hosts) restart. The
/// canonical recovery ritual after upgrading the `gritty` binary on one or
/// both sides: old daemon gets killed across the version mismatch, tunnel
/// is torn down and respawned (which bootstraps the remote daemon with the
/// new binary), and the local server is started if we targeted `local`.
pub(crate) async fn restart(
    host: Option<String>,
    ctl_socket: Option<PathBuf>,
) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    let host = host.unwrap_or_else(|| "local".to_string());
    // Snapshot the override before resolve_ctl_path consumes it -- the local
    // respawn below must land on the same socket, not the default path.
    let ctl_socket_arg = ctl_socket.as_ref().map(|p| p.to_string_lossy().into_owned());
    let ctl_path = super::util::resolve_ctl_path(ctl_socket, Some(&host))?;

    // Step 1: kill-server, tolerant of both "nothing running" and
    // protocol-version mismatch. We don't want to fail the whole restart
    // just because the daemon was already gone.
    match super::util::server_request_any_version(&ctl_path, Frame::KillServer).await {
        Ok(Frame::Ok) => {
            eprintln!("\x1b[32m\u{25b8} server killed\x1b[0m");
        }
        Ok(Frame::Error { message, .. }) => {
            eprintln!("\x1b[2;33m\u{25b8} kill-server: {message} (continuing)\x1b[0m");
        }
        Ok(other) => {
            eprintln!(
                "\x1b[2;33m\u{25b8} kill-server: unexpected response {other:?} (continuing)\x1b[0m"
            );
        }
        Err(_) => {
            // Connect failed -- no daemon to kill, that's fine.
            eprintln!("\x1b[2;33m\u{25b8} no server running\x1b[0m");
        }
    }

    if host == "local" {
        // For local, kick off a fresh `gritty server`.
        eprintln!("\x1b[2;33m\u{25b8} starting server...\x1b[0m");
        super::util::auto_start(&super::util::server_auto_start_args(ctl_socket_arg.as_deref()))?;
        eprintln!("\x1b[32m\u{25b8} server restarted\x1b[0m");
    } else {
        // Remote: capture the original destination from the .dest sidecar
        // before disconnect wipes it. Using just `host` here would collapse
        // `user@server.example.com:2222` down to the friendly connection
        // name and break SSH.
        let destination = gritty::connect::resolve_destination(&host);
        // Capture the recreate args (destination + persisted CLI -o options)
        // *before* disconnect wipes the sidecar files.
        let recreate = gritty::connect::tunnel_recreate_args(&host, &destination);
        // Tear down the tunnel (the supervisor may already be exiting
        // because the ctl socket vanished when the daemon died, but
        // `disconnect` is idempotent for the "already stopped" case).
        gritty::connect::disconnect(&host).await?;
        eprintln!("\x1b[2;33m\u{25b8} starting tunnel {host}...\x1b[0m");
        let recreate: Vec<&str> = recreate.iter().map(String::as_str).collect();
        super::util::auto_start(&recreate)?;
        eprintln!("\x1b[32m\u{25b8} {host} restarted\x1b[0m");
    }
    Ok(())
}

/// The nine shared session-table columns (ID..Status) for one row. Identical
/// between `list_sessions` and `list_all_sessions`; the latter just prepends a
/// Host column. Single-sourced so the "starting"/attached/heartbeat/detached
/// status logic and column order cannot drift between the two listings.
fn session_status_cols(s: &gritty::protocol::SessionEntry, now: u64) -> Vec<String> {
    let name = if s.name.is_empty() { "-".to_string() } else { s.name.clone() };
    let (pty, pid, created, status) = if s.shell_pid == 0 {
        ("-".to_string(), "-".to_string(), "-".to_string(), "starting".to_string())
    } else {
        let status = if s.attached {
            if s.last_heartbeat > 0 {
                let ago = now.saturating_sub(s.last_heartbeat);
                format!("attached (heartbeat {ago}s ago)")
            } else {
                "attached".to_string()
            }
        } else {
            "detached".to_string()
        };
        (s.pty_path.clone(), s.shell_pid.to_string(), format_timestamp(s.created_at), status)
    };
    vec![
        s.id.to_string(),
        name,
        s.foreground_cmd.clone(),
        s.cwd.clone(),
        s.client_name.clone(),
        pty,
        pid,
        created,
        status,
    ]
}

pub(crate) async fn list_sessions(ctl_path: PathBuf) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    let resp = server_request(&ctl_path, Frame::ListSessions).await?;
    match resp {
        Frame::SessionInfo { sessions } => {
            if sessions.is_empty() {
                println!("no active sessions");
            } else {
                let now = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs();

                // Build row data
                let rows: Vec<Vec<String>> =
                    sessions.iter().map(|s| session_status_cols(s, now)).collect();

                gritty::table::print_table(
                    &["ID", "Name", "Cmd", "CWD", "Client", "PTY", "PID", "Created", "Status"],
                    &rows,
                );
            }
            Ok(())
        }
        other => {
            anyhow::bail!("unexpected response from server: {other:?}");
        }
    }
}

pub(crate) async fn list_all_sessions() -> anyhow::Result<()> {
    use gritty::protocol::{Frame, FrameCodec, SessionEntry};

    let probes = discover_daemon_probes();

    if probes.is_empty() {
        anyhow::bail!("no server running");
    }

    let futures: Vec<_> = probes
        .into_iter()
        .map(|(host, path)| async move {
            let result = tokio::time::timeout(std::time::Duration::from_secs(2), async {
                let stream = gritty::security::connect_verified(&path)
                    .await
                    .map_err(|e| format!("connect: {e}"))?;
                let mut framed = tokio_util::codec::Framed::new(stream, FrameCodec);
                let info = gritty::handshake(&mut framed, gritty::get_or_create_device_id())
                    .await
                    .map_err(|e| e.to_string())?;
                gritty::require_matched_version(&info).map_err(|e| e.to_string())?;
                futures_util::SinkExt::send(&mut framed, Frame::ListSessions)
                    .await
                    .map_err(|e| format!("send ListSessions: {e}"))?;
                match Frame::expect_from(futures_util::StreamExt::next(&mut framed).await) {
                    Ok(Frame::SessionInfo { sessions }) => Ok(sessions),
                    Ok(Frame::Error { message, .. }) => Err(message),
                    Ok(other) => Err(format!("unexpected response: {other:?}")),
                    Err(e) => Err(e.to_string()),
                }
            })
            .await;
            let outcome: Result<Vec<SessionEntry>, String> = match result {
                Ok(inner) => inner,
                Err(_) => Err("probe timed out after 2s".to_string()),
            };
            (host, outcome)
        })
        .collect();

    let results: Vec<(String, Result<Vec<SessionEntry>, String>)> =
        futures_util::future::join_all(futures).await;

    let errors: Vec<(&String, &String)> =
        results.iter().filter_map(|(h, r)| r.as_ref().err().map(|e| (h, e))).collect();
    let has_sessions =
        results.iter().any(|(_, r)| r.as_ref().map(|s| !s.is_empty()).unwrap_or(false));

    if !has_sessions {
        if errors.is_empty() {
            println!("no active sessions");
        } else {
            println!("no active sessions");
            for (host, err) in &errors {
                eprintln!("\x1b[2;33m\u{26a0} {host}: {err}\x1b[0m");
            }
        }
        return Ok(());
    }

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();

    let ok_hosts: Vec<(&String, &Vec<SessionEntry>)> = results
        .iter()
        .filter_map(|(h, r)| match r {
            Ok(s) if !s.is_empty() => Some((h, s)),
            _ => None,
        })
        .collect();
    let multi_host = ok_hosts.len() > 1;

    // Build row data: [host, id, name, cmd, cwd, client, pty, pid, created, status]
    let rows: Vec<Vec<String>> = ok_hosts
        .iter()
        .flat_map(|(host, sessions)| {
            sessions.iter().map(move |s| {
                let mut row = Vec::with_capacity(10);
                row.push((*host).clone());
                row.extend(session_status_cols(s, now));
                row
            })
        })
        .collect();

    if multi_host {
        gritty::table::print_table(
            &["Host", "ID", "Name", "Cmd", "CWD", "Client", "PTY", "PID", "Created", "Status"],
            &rows,
        );
    } else {
        let host = &rows[0][0];
        println!("Host: {host}");
        let trimmed: Vec<Vec<String>> = rows.iter().map(|r| r[1..].to_vec()).collect();
        gritty::table::print_table(
            &["ID", "Name", "Cmd", "CWD", "Client", "PTY", "PID", "Created", "Status"],
            &trimmed,
        );
    }
    for (host, err) in &errors {
        eprintln!("\x1b[2;33m\u{26a0} {host}: {err}\x1b[0m");
    }
    Ok(())
}

/// Print available sessions and exit with an error when a session-requiring
/// command is invoked without the session part (e.g. `gritty tail local`
/// instead of `gritty tail local:session`).
pub(crate) async fn suggest_session(cmd: &str, host: &str, ctl_path: &Path) -> anyhow::Result<()> {
    use gritty::protocol::Frame;

    let ctl_path_buf = ctl_path.to_path_buf();
    let resp = match server_request(&ctl_path_buf, Frame::ListSessions).await {
        Ok(resp) => resp,
        Err(_) => {
            anyhow::bail!("specify a session: gritty {cmd} {host}:<session>");
        }
    };

    match resp {
        Frame::SessionInfo { sessions } if sessions.is_empty() => {
            anyhow::bail!("no active sessions on {host}");
        }
        Frame::SessionInfo { sessions } => {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs();

            let mut msg = format!("specify a session: gritty {cmd} {host}:<session>\n\n");
            msg.push_str("  ID  Name     Age\n");
            for s in &sessions {
                let name = if s.name.is_empty() { "-".to_string() } else { s.name.clone() };
                let age = format_age(now, s.created_at);
                msg.push_str(&format!("  {}   {:<8} {}\n", s.id, name, age));
            }
            anyhow::bail!("{msg}");
        }
        _ => anyhow::bail!("specify a session: gritty {cmd} {host}:<session>"),
    }
}

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

    fn row(name: &str) -> Row {
        Row {
            name: name.to_string(),
            attached: false,
            age: String::new(),
            cmd: String::new(),
            cwd: String::new(),
            client: String::new(),
            hotkey: None,
        }
    }

    #[test]
    fn suggest_name_uses_default_when_free() {
        assert_eq!(suggest_name(&[]), "default");
        assert_eq!(suggest_name(&[row("work")]), "default");
    }

    #[test]
    fn suggest_name_increments_past_default() {
        assert_eq!(suggest_name(&[row("default")]), "session-2");
        assert_eq!(suggest_name(&[row("default"), row("session-2")]), "session-3");
    }

    #[test]
    fn suggest_name_fills_first_free_slot() {
        // session-2 missing -- the first free slot, not the max+1.
        assert_eq!(suggest_name(&[row("default"), row("session-3")]), "session-2");
    }

    fn entry() -> gritty::protocol::SessionEntry {
        gritty::protocol::SessionEntry {
            id: 3,
            name: String::new(),
            pty_path: "/dev/pts/7".to_string(),
            shell_pid: 1234,
            created_at: 0,
            attached: false,
            last_heartbeat: 0,
            foreground_cmd: "vim".to_string(),
            cwd: "/home/x".to_string(),
            client_name: "laptop".to_string(),
            agent_forwarding_active: false,
            is_last_attached: false,
        }
    }

    #[test]
    fn session_status_cols_starting_when_no_shell() {
        let mut s = entry();
        s.shell_pid = 0;
        let cols = session_status_cols(&s, 100);
        // id, name(-), cmd, cwd, client, pty(-), pid(-), created(-), status
        assert_eq!(cols[0], "3");
        assert_eq!(cols[1], "-"); // empty name renders as "-"
        assert_eq!(cols[5], "-"); // pty
        assert_eq!(cols[6], "-"); // pid
        assert_eq!(cols[8], "starting");
    }

    #[test]
    fn session_status_cols_attached_reports_heartbeat_age() {
        let mut s = entry();
        s.attached = true;
        s.last_heartbeat = 90;
        let cols = session_status_cols(&s, 100);
        assert_eq!(cols[8], "attached (heartbeat 10s ago)");
        assert_eq!(cols[6], "1234"); // pid
    }

    #[test]
    fn session_status_cols_detached_and_attached_no_heartbeat() {
        let s = entry();
        assert_eq!(session_status_cols(&s, 100)[8], "detached");
        let mut s2 = entry();
        s2.attached = true;
        assert_eq!(session_status_cols(&s2, 100)[8], "attached");
    }
}
