use anyhow::Result;
use futures_util::{SinkExt, StreamExt};
use tokio::net::TcpStream;
use tokio_tungstenite::{
    connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream,
};
use tracing::info;

/// A WebSocket connection (either client or server side).
pub type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;

/// Connect to a WebSocket server as a client.
pub async fn connect_to_server(url: &str) -> Result<WsStream> {
    info!("Connecting to WebSocket server: {}", url);
    let (ws_stream, response) = connect_async(url).await?;
    info!("WebSocket connected, status: {}", response.status());
    Ok(ws_stream)
}

/// Send a binary message over the WebSocket.
pub async fn send_binary(ws: &mut WsStream, data: Vec<u8>) -> Result<()> {
    ws.send(Message::Binary(data.into())).await?;
    Ok(())
}

/// Receive the next binary message from the WebSocket.
/// Returns None if the connection is closed.
pub async fn recv_binary(ws: &mut WsStream) -> Result<Option<Vec<u8>>> {
    loop {
        match ws.next().await {
            Some(Ok(Message::Binary(data))) => return Ok(Some(data.to_vec())),
            Some(Ok(Message::Close(_))) => return Ok(None),
            Some(Ok(Message::Ping(data))) => {
                ws.send(Message::Pong(data)).await?;
            }
            Some(Ok(_)) => continue,
            Some(Err(e)) => return Err(anyhow::anyhow!("WebSocket error: {}", e)),
            None => return Ok(None),
        }
    }
}

/// Send a close frame.
pub async fn close(ws: &mut WsStream) -> Result<()> {
    ws.close(None).await?;
    Ok(())
}

/// WebSocket server acceptor — accepts a TcpStream and performs the WebSocket upgrade.
pub async fn accept_connection(
    stream: TcpStream,
) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
    let ws = tokio_tungstenite::accept_async(MaybeTlsStream::Plain(stream)).await?;
    Ok(ws)
}