rmux-server 0.6.1

use std::future::pending;
use std::io;
#[cfg(unix)]
use std::time::Duration;

use rmux_core::events::{OutputCursorItem, OutputGap};
use rmux_proto::{
    encode_attach_data, encode_attach_data_into_slice, encode_attach_message, AttachFrameDecoder,
    AttachMessage, ATTACH_DATA_HEADER_LEN,
};
#[cfg(unix)]
use rmux_pty::PtyIo;
#[cfg(unix)]
use rmux_pty::PtyMaster;
#[cfg(unix)]
use tokio::io::unix::AsyncFd;
use tracing::warn;

use crate::outer_terminal::OuterTerminal;

use super::attach_transport::{AttachTransport, TryAttachRead};
use super::types::{AttachTarget, OpenAttachTarget, PaneOutputReceiver};

#[cfg(unix)]
// Eight immediate 64 KiB reads amortize readiness wakeups under sustained PTY
// output while still yielding regularly enough to keep other panes responsive.
const MAX_IMMEDIATE_PANE_READS: usize = 8;
#[cfg(unix)]
const MAX_STARTUP_EIO_READS: usize = 256;
#[cfg(unix)]
const STARTUP_EIO_YIELD_READS: usize = 8;
#[cfg(unix)]
const STARTUP_EIO_BACKOFF: Duration = Duration::from_millis(1);
const STACK_ATTACH_DATA_PAYLOAD: usize = 1024;

#[cfg(unix)]
#[derive(Debug, Default)]
pub(super) struct PaneReadinessState {
    immediate_reads: usize,
    startup_eio_reads: usize,
    established: bool,
    startup_eio_exhausted: bool,
}

pub(super) fn open_attach_target(
    target: AttachTarget,
    render_stream: bool,
) -> io::Result<OpenAttachTarget> {
    let AttachTarget {
        session_name,
        input_target,
        pane_master,
        pane_output,
        pane_output_start_sequence,
        render_frame,
        outer_terminal,
        cursor_style,
        active_pane_geometry,
        raw_passthrough,
        kitty_graphics_passthrough,
        sixel_passthrough,
        persistent_overlay_state_id,
        live_pane,
    } = target;
    Ok(OpenAttachTarget {
        session_name,
        input_target,
        pane_master,
        predicted_echo: Default::default(),
        predicted_echo_started_at: None,
        pane_output: Some(pane_output.subscribe_live_from_sequence(pane_output_start_sequence)),
        render_frame,
        outer_terminal,
        cursor_style,
        active_pane_geometry,
        raw_passthrough,
        kitty_graphics_passthrough,
        sixel_passthrough,
        persistent_overlay_state_id,
        live_pane,
        render_stream,
    })
}

#[cfg(unix)]
pub(super) fn open_pane_writer(pane_master: PtyMaster) -> io::Result<AsyncFd<PtyIo>> {
    let pane_writer = pane_master.into_io();
    pane_writer.set_nonblocking()?;

    AsyncFd::new(pane_writer)
}

pub(super) async fn emit_render_frame(
    stream: &AttachTransport,
    outer_terminal: &OuterTerminal,
    render_frame: &[u8],
) -> io::Result<()> {
    let frame = outer_terminal.wrap_render_frame(render_frame);
    emit_attach_bytes(stream, &frame).await
}

pub(super) async fn emit_coalescible_render_frame(
    stream: &AttachTransport,
    outer_terminal: &OuterTerminal,
    render_frame: &[u8],
    render_stream: bool,
) -> io::Result<()> {
    let frame = outer_terminal.wrap_render_frame(render_frame);
    if frame.is_empty() {
        return Ok(());
    }
    if render_stream {
        emit_attach_frame(stream, &AttachMessage::Render(frame)).await
    } else {
        emit_attach_bytes(stream, &frame).await
    }
}

pub(super) async fn read_socket_bytes(
    stream: &AttachTransport,
    decoder: &mut AttachFrameDecoder,
) -> io::Result<bool> {
    stream.read_into(decoder).await
}

pub(super) fn try_read_socket_bytes(
    stream: &AttachTransport,
    decoder: &mut AttachFrameDecoder,
) -> io::Result<TryAttachRead> {
    stream.try_read_into(decoder)
}

pub(super) async fn emit_attach_message(
    stream: &AttachTransport,
    message: &AttachMessage,
) -> io::Result<()> {
    let frame = encode_attach_message(message).map_err(io::Error::other)?;
    write_all_to_stream(stream, &frame).await
}

pub(super) async fn emit_attach_frame(
    stream: &AttachTransport,
    message: &AttachMessage,
) -> io::Result<()> {
    let frame = encode_attach_message(message).map_err(io::Error::other)?;
    write_all_to_stream(stream, &frame).await
}

pub(super) async fn recv_pane_output(
    pane_output: &mut PaneOutputReceiver,
) -> io::Result<OutputCursorItem> {
    match pane_output.recv().await {
        OutputCursorItem::Event(event) => Ok(OutputCursorItem::Event(event)),
        OutputCursorItem::Gap(gap) => {
            warn_pane_output_gap(&gap);
            Ok(OutputCursorItem::Gap(gap))
        }
    }
}

pub(super) fn warn_pane_output_gap(gap: &OutputGap) {
    warn!(
        expected_sequence = gap.expected_sequence(),
        resume_sequence = gap.resume_sequence(),
        missed_events = gap.missed_events(),
        recent_bytes = gap.recent_snapshot().len(),
        "attach pane output receiver lagged"
    );
}

pub(super) async fn recv_pane_output_optional(
    pane_output: Option<&mut PaneOutputReceiver>,
) -> io::Result<Option<OutputCursorItem>> {
    match pane_output {
        Some(pane_output) => recv_pane_output(pane_output).await.map(Some),
        None => pending().await,
    }
}

pub(super) async fn emit_attach_data_frame(
    stream: &AttachTransport,
    bytes: &[u8],
) -> io::Result<()> {
    if bytes.len() <= STACK_ATTACH_DATA_PAYLOAD {
        let mut frame = [0_u8; STACK_ATTACH_DATA_PAYLOAD + ATTACH_DATA_HEADER_LEN];
        let len = encode_attach_data_into_slice(bytes, &mut frame).map_err(io::Error::other)?;
        return write_all_to_stream(stream, &frame[..len]).await;
    }

    let frame = encode_attach_data(bytes).map_err(io::Error::other)?;
    write_all_to_stream(stream, &frame).await
}

pub(super) async fn emit_attach_bytes(stream: &AttachTransport, bytes: &[u8]) -> io::Result<()> {
    if bytes.is_empty() {
        return Ok(());
    }

    emit_attach_data_frame(stream, bytes).await
}

pub(super) async fn emit_attach_stop(
    stream: &AttachTransport,
    current_target: &OpenAttachTarget,
) -> io::Result<()> {
    emit_attach_bytes(
        stream,
        &current_target.outer_terminal.attach_stop_sequence(),
    )
    .await
}

pub(super) async fn emit_detached_attach_stop(
    stream: &AttachTransport,
    current_target: &OpenAttachTarget,
) -> io::Result<()> {
    let mut bytes = current_target.outer_terminal.attach_stop_sequence();
    bytes.extend_from_slice(
        format!(
            "[detached (from session {})]\r\n",
            current_target.session_name
        )
        .as_bytes(),
    );
    emit_attach_bytes(stream, &bytes).await
}

pub(super) async fn emit_exited_attach_stop(
    stream: &AttachTransport,
    current_target: &OpenAttachTarget,
) -> io::Result<()> {
    let mut bytes = current_target.outer_terminal.attach_stop_sequence();
    bytes.extend_from_slice(b"[exited]\r\n");
    emit_attach_bytes(stream, &bytes).await
}

#[cfg(unix)]
pub(super) async fn read_from_pane(
    pane_reader: &AsyncFd<PtyIo>,
    readiness: &mut PaneReadinessState,
    buffer: &mut [u8],
) -> io::Result<usize> {
    loop {
        if readiness.immediate_reads >= MAX_IMMEDIATE_PANE_READS {
            readiness.immediate_reads = 0;
            tokio::task::yield_now().await;
        }

        // Read once before awaiting readiness. A pane can emit its initial
        // prompt/output before the async task reaches readable().await; this
        // preserves AsyncFd while avoiding dependence on a later readiness edge.
        match try_read_pane_now(pane_reader.get_ref(), buffer)? {
            PaneRead::Bytes(bytes_read) => {
                readiness.record_immediate_bytes(bytes_read);
                return Ok(bytes_read);
            }
            PaneRead::NotReady => {}
            PaneRead::SlaveUnavailable => match readiness.retry_startup_eio() {
                StartupEioReadiness::Retry(delay) => {
                    delay_startup_eio_retry(delay).await;
                    continue;
                }
                StartupEioReadiness::StartupRetriesExhausted
                | StartupEioReadiness::EstablishedEof => return Ok(0),
            },
        }
        readiness.immediate_reads = 0;

        let mut ready = pane_reader.readable().await?;
        match ready.try_io(|inner| inner.get_ref().read(&mut *buffer)) {
            Ok(Ok(bytes_read)) => {
                readiness.record_ready_bytes(bytes_read);
                return Ok(bytes_read);
            }
            Ok(Err(error)) if error.kind() == io::ErrorKind::Interrupted => continue,
            Ok(Err(error))
                if error.raw_os_error() == Some(rustix::io::Errno::IO.raw_os_error()) =>
            {
                match readiness.retry_startup_eio() {
                    StartupEioReadiness::Retry(delay) => {
                        delay_startup_eio_retry(delay).await;
                        continue;
                    }
                    StartupEioReadiness::StartupRetriesExhausted
                    | StartupEioReadiness::EstablishedEof => return Ok(0),
                }
            }
            Ok(Err(error)) => return Err(error),
            Err(_would_block) => continue,
        }
    }
}

#[cfg(unix)]
pub(super) fn try_read_available_from_pane(
    pane_reader: &AsyncFd<PtyIo>,
    buffer: &mut [u8],
) -> io::Result<Option<usize>> {
    match try_read_pane_now(pane_reader.get_ref(), buffer)? {
        PaneRead::Bytes(bytes_read) => Ok(Some(bytes_read)),
        PaneRead::NotReady | PaneRead::SlaveUnavailable => Ok(None),
    }
}

#[cfg(unix)]
async fn delay_startup_eio_retry(delay: StartupEioRetryDelay) {
    match delay {
        StartupEioRetryDelay::Yield => tokio::task::yield_now().await,
        StartupEioRetryDelay::Sleep(duration) => tokio::time::sleep(duration).await,
    }
}

#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StartupEioReadiness {
    Retry(StartupEioRetryDelay),
    StartupRetriesExhausted,
    EstablishedEof,
}

#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StartupEioRetryDelay {
    Yield,
    Sleep(Duration),
}

#[cfg(unix)]
impl StartupEioRetryDelay {
    fn for_attempt(attempt: usize) -> Self {
        if attempt <= STARTUP_EIO_YIELD_READS {
            Self::Yield
        } else {
            Self::Sleep(STARTUP_EIO_BACKOFF)
        }
    }
}

#[cfg(test)]
#[cfg(unix)]
mod startup_eio_retry_delay_tests {
    use super::{
        PaneReadinessState, StartupEioReadiness, StartupEioRetryDelay, STARTUP_EIO_BACKOFF,
        STARTUP_EIO_YIELD_READS,
    };

    #[test]
    fn startup_eio_retries_yield_before_backing_off() {
        for attempt in 1..=STARTUP_EIO_YIELD_READS {
            assert_eq!(
                StartupEioRetryDelay::for_attempt(attempt),
                StartupEioRetryDelay::Yield
            );
        }
        assert_eq!(
            StartupEioRetryDelay::for_attempt(STARTUP_EIO_YIELD_READS + 1),
            StartupEioRetryDelay::Sleep(STARTUP_EIO_BACKOFF)
        );
    }

    #[test]
    fn readiness_state_tracks_startup_eio_until_output_is_established() {
        let mut readiness = PaneReadinessState::default();

        assert_eq!(
            readiness.retry_startup_eio(),
            StartupEioReadiness::Retry(StartupEioRetryDelay::Yield)
        );
        assert_eq!(readiness.startup_eio_reads, 1);

        readiness.record_ready_bytes(1);

        assert!(readiness.established);
        assert_eq!(readiness.startup_eio_reads, 0);
        assert_eq!(
            readiness.retry_startup_eio(),
            StartupEioReadiness::EstablishedEof
        );
        assert!(!readiness.startup_eio_exhausted());
    }

    #[test]
    fn readiness_state_caps_startup_eio_retries() {
        let mut readiness = PaneReadinessState::default();

        for _ in 0..super::MAX_STARTUP_EIO_READS {
            assert!(matches!(
                readiness.retry_startup_eio(),
                StartupEioReadiness::Retry(_)
            ));
        }

        assert_eq!(
            readiness.retry_startup_eio(),
            StartupEioReadiness::StartupRetriesExhausted
        );
        assert!(readiness.startup_eio_exhausted());
    }
}

#[cfg(unix)]
impl PaneReadinessState {
    fn record_immediate_bytes(&mut self, bytes_read: usize) {
        self.startup_eio_reads = 0;
        self.startup_eio_exhausted = false;
        if bytes_read > 0 {
            self.immediate_reads += 1;
            self.established = true;
        } else {
            self.immediate_reads = 0;
        }
    }

    fn record_ready_bytes(&mut self, bytes_read: usize) {
        self.immediate_reads = 0;
        self.startup_eio_reads = 0;
        self.startup_eio_exhausted = false;
        if bytes_read > 0 {
            self.established = true;
        }
    }

    pub(super) fn startup_eio_exhausted(&self) -> bool {
        self.startup_eio_exhausted
    }

    pub(super) fn startup_eio_reads(&self) -> usize {
        self.startup_eio_reads
    }

    fn retry_startup_eio(&mut self) -> StartupEioReadiness {
        self.immediate_reads = 0;
        // Unix PTY masters can report EIO as EOF. Linux can also report it
        // briefly before the slave side has reached a stable post-spawn state.
        // Before the first successful read, treat a bounded run of EIO as
        // startup latency; after output is established, EIO is normal EOF.
        if self.established {
            return StartupEioReadiness::EstablishedEof;
        }
        if self.startup_eio_reads >= MAX_STARTUP_EIO_READS {
            self.startup_eio_exhausted = true;
            return StartupEioReadiness::StartupRetriesExhausted;
        }
        self.startup_eio_reads += 1;
        StartupEioReadiness::Retry(StartupEioRetryDelay::for_attempt(self.startup_eio_reads))
    }
}

#[cfg(unix)]
enum PaneRead {
    Bytes(usize),
    NotReady,
    SlaveUnavailable,
}

#[cfg(unix)]
fn try_read_pane_now(reader: &PtyIo, buffer: &mut [u8]) -> io::Result<PaneRead> {
    match reader.read(buffer) {
        Ok(bytes_read) => Ok(PaneRead::Bytes(bytes_read)),
        Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(PaneRead::NotReady),
        Err(error) if error.kind() == io::ErrorKind::Interrupted => Ok(PaneRead::NotReady),
        Err(error) if error.raw_os_error() == Some(rustix::io::Errno::IO.raw_os_error()) => {
            Ok(PaneRead::SlaveUnavailable)
        }
        Err(error) => Err(error),
    }
}

async fn write_all_to_stream(stream: &AttachTransport, bytes: &[u8]) -> io::Result<()> {
    stream.write_all(bytes).await
}

pub(super) fn invalid_attach_message(error: rmux_proto::RmuxError) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, error)
}

#[cfg(all(test, unix))]
mod unix_tests {
    use std::error::Error;
    use std::io;
    use std::time::{Duration, Instant};

    use rmux_pty::{ChildCommand, TerminalSize as PtyTerminalSize};

    use super::{open_pane_writer, read_from_pane, PaneReadinessState};

    #[tokio::test]
    async fn read_from_pane_consumes_output_written_before_readiness_wait(
    ) -> Result<(), Box<dyn Error>> {
        let mut spawned = ChildCommand::new("sh")
            .args(["-c", "printf PRE_READY; sleep 1"])
            .size(PtyTerminalSize::new(80, 24))
            .spawn()?;
        let output_reader = spawned.master().try_clone()?;

        tokio::time::sleep(Duration::from_millis(100)).await;
        let pane_reader = open_pane_writer(output_reader)?;
        let captured =
            read_until_contains(&pane_reader, "PRE_READY", Duration::from_millis(500)).await?;

        spawned.child().terminate_forcefully()?;
        let _ = spawned.child_mut().wait()?;

        assert!(
            captured.contains("PRE_READY"),
            "expected pre-existing pane output, got {captured:?}"
        );
        Ok(())
    }

    #[tokio::test]
    async fn read_from_pane_consumes_output_written_after_registration_before_wait(
    ) -> Result<(), Box<dyn Error>> {
        let mut spawned = ChildCommand::new("sh")
            .args(["-c", "sleep 0.05; printf POST_REGISTER; sleep 1"])
            .size(PtyTerminalSize::new(80, 24))
            .spawn()?;
        let output_reader = spawned.master().try_clone()?;
        let pane_reader = open_pane_writer(output_reader)?;

        tokio::time::sleep(Duration::from_millis(100)).await;
        let captured =
            read_until_contains(&pane_reader, "POST_REGISTER", Duration::from_millis(500)).await?;

        spawned.child().terminate_forcefully()?;
        let _ = spawned.child_mut().wait()?;

        assert!(
            captured.contains("POST_REGISTER"),
            "expected post-registration pane output, got {captured:?}"
        );
        Ok(())
    }

    async fn read_until_contains(
        pane_reader: &tokio::io::unix::AsyncFd<rmux_pty::PtyIo>,
        needle: &str,
        timeout: Duration,
    ) -> io::Result<String> {
        let deadline = Instant::now() + timeout;
        let mut readiness = PaneReadinessState::default();
        let mut buffer = [0_u8; 256];
        let mut captured = Vec::new();

        while Instant::now() < deadline {
            let remaining = deadline.saturating_duration_since(Instant::now());
            let bytes_read = tokio::time::timeout(
                remaining,
                read_from_pane(pane_reader, &mut readiness, &mut buffer),
            )
            .await
            .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "timed out"))??;
            if bytes_read == 0 {
                break;
            }
            captured.extend_from_slice(&buffer[..bytes_read]);
            let rendered = String::from_utf8_lossy(&captured);
            if rendered.contains(needle) {
                return Ok(rendered.into_owned());
            }
        }

        Ok(String::from_utf8_lossy(&captured).into_owned())
    }
}

#[cfg(test)]
mod emit_render_tests {
    use super::{emit_attach_message, emit_coalescible_render_frame};
    use crate::outer_terminal::{OuterTerminal, OuterTerminalContext};
    use crate::pane_io::attach_transport::AttachTransport;
    use rmux_core::OptionStore;
    use rmux_proto::{AttachFrameDecoder, AttachMessage, AttachShellCommand};
    use tokio::io::AsyncReadExt;

    async fn emitted_message(render_stream: bool) -> AttachMessage {
        let (stream, mut peer) = tokio::io::duplex(1024);
        let stream = AttachTransport::from_io(stream);
        let outer_terminal =
            OuterTerminal::resolve(&OptionStore::default(), OuterTerminalContext::default());

        emit_coalescible_render_frame(&stream, &outer_terminal, b"frame", render_stream)
            .await
            .expect("render frame emits");

        let mut bytes = [0_u8; 128];
        let count = peer.read(&mut bytes).await.expect("read emitted frame");
        let mut decoder = AttachFrameDecoder::new();
        decoder.push_bytes(&bytes[..count]);
        decoder
            .next_message()
            .expect("decode succeeds")
            .expect("message emitted")
    }

    #[tokio::test]
    async fn coalescible_render_falls_back_to_data_for_legacy_clients() {
        assert!(matches!(
            emitted_message(false).await,
            AttachMessage::Data(bytes) if bytes.ends_with(b"frame")
        ));
    }

    #[tokio::test]
    async fn coalescible_render_uses_render_for_capable_clients() {
        assert!(matches!(
            emitted_message(true).await,
            AttachMessage::Render(bytes) if bytes.ends_with(b"frame")
        ));
    }

    #[tokio::test]
    async fn attach_control_messages_are_not_wrapped_as_terminal_data() {
        let (stream, mut peer) = tokio::io::duplex(1024);
        let stream = AttachTransport::from_io(stream);
        let message = AttachMessage::DetachExecShellCommand(AttachShellCommand::new(
            "echo detached".to_owned(),
            "/bin/sh".to_owned(),
            "/tmp".to_owned(),
        ));

        emit_attach_message(&stream, &message)
            .await
            .expect("attach message emits");

        let mut bytes = [0_u8; 256];
        let count = peer.read(&mut bytes).await.expect("read emitted frame");
        let mut decoder = AttachFrameDecoder::new();
        decoder.push_bytes(&bytes[..count]);
        assert_eq!(
            decoder.next_message().expect("decode succeeds"),
            Some(message)
        );
    }
}