//! PTY pump shared by the attach endpoints.
//!
//! - Binary frames carry raw PTY bytes both directions
//! - Text frames carry [`maestro_protocol::ClientControl`] (Resize/Ping/Pong)

use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::ClientControl;
use tokio::sync::broadcast;
use tokio::sync::broadcast::error::RecvError;
use uuid::Uuid;

use crate::session::SessionManager;

/// Pump a WebSocket against a session's PTY, behind `WS /api/attach/:id`.
///
/// The subscription is acquired by the handler *before* the upgrade and handed
/// in here, so there is no window in which the session can disappear between
/// the existence check and the subscribe.
pub(crate) async fn serve(
    socket: WebSocket,
    manager: SessionManager,
    id: Uuid,
    replay: Vec<u8>,
    mut rx: broadcast::Receiver<Vec<u8>>,
) -> Result<(), axum::Error> {
    let (mut sink, mut stream) = socket.split();

    // NAP-229: send buffered scrollback up-front so the client sees the last
    // screenful immediately instead of a blank terminal. The replay lock the
    // bridge held during snapshot+subscribe guarantees no byte appears both
    // here and on the live receiver, and none is dropped at the seam.
    if !replay.is_empty() && sink.send(Message::Binary(replay)).await.is_err() {
        return Ok(());
    }

    // Set to true when the session itself terminates (broadcast closes); we
    // owe the client a Close frame in that case. Local breakouts (e.g. the
    // client sent its own Close, or the socket died) skip it.
    let mut session_ended = false;

    loop {
        tokio::select! {
            // PTY output → client
            received = rx.recv() => match received {
                Ok(bytes) => {
                    if sink.send(Message::Binary(bytes)).await.is_err() {
                        break;
                    }
                }
                Err(RecvError::Lagged(_)) => {
                    // dropped some frames under back-pressure; keep going
                    continue;
                }
                Err(RecvError::Closed) => {
                    session_ended = true;
                    break;
                }
            },

            // Client input → tmux
            from_client = stream.next() => match from_client {
                None => break,
                Some(Err(_)) => break,
                Some(Ok(Message::Close(_))) => break,
                Some(Ok(Message::Binary(bytes))) => {
                    let _ = manager.send_keys(id, &bytes).await;
                }
                Some(Ok(Message::Text(text))) => {
                    if let Ok(ctrl) = serde_json::from_str::<ClientControl>(&text) {
                        match ctrl {
                            ClientControl::Resize { cols, rows } => {
                                let _ = manager.resize(id, cols, rows).await;
                            }
                            ClientControl::Ping => {
                                let pong = serde_json::to_string(&ClientControl::Pong)
                                    .expect("serialize pong");
                                if sink.send(Message::Text(pong)).await.is_err() {
                                    break;
                                }
                            }
                            ClientControl::Pong | ClientControl::Error(_) => {}
                        }
                    }
                }
                Some(Ok(_)) => {}
            },
        }
    }

    if session_ended {
        let _ = sink.send(Message::Close(None)).await;
    }
    Ok(())
}