use crate::os_input_output::{AsyncSignals, AsyncStdin, ClientOsApi, SignalEvent};
use crate::remote_attach::websockets::MaybeTls;
use crate::remote_attach::WebSocketConnections;
use crate::run_remote_client_terminal_loop;
use crate::web_client::control_message::{
    WebClientToWebServerControlMessage, WebClientToWebServerControlMessagePayload,
    WebServerToWebClientControlMessage,
};
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use serial_test::serial;
use std::io::{self, Write};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use zellij_utils::data::Palette;
use zellij_utils::errors::ErrorContext;
use zellij_utils::ipc::{ClientToServerMsg, ServerToClientMsg};
use zellij_utils::pane_size::Size;

struct MockAsyncStdin {
    rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<Vec<u8>>>>,
}

#[async_trait]
impl AsyncStdin for MockAsyncStdin {
    async fn read(&mut self) -> io::Result<Vec<u8>> {
        match self.rx.lock().await.recv().await {
            Some(data) => Ok(data),
            None => Ok(Vec::new()),
        }
    }
}

struct MockAsyncSignals {
    rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<SignalEvent>>>,
}

#[async_trait]
impl AsyncSignals for MockAsyncSignals {
    async fn recv(&mut self) -> Option<SignalEvent> {
        self.rx.lock().await.recv().await
    }
}

#[derive(Clone)]
struct TestClientOsApi {
    stdout_buffer: Arc<Mutex<Vec<u8>>>,
    stdin_rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<Vec<u8>>>>,
    signal_rx: Arc<tokio::sync::Mutex<mpsc::UnboundedReceiver<SignalEvent>>>,
    terminal_size: Size,
}

impl TestClientOsApi {
    fn new(
        stdin_rx: mpsc::UnboundedReceiver<Vec<u8>>,
        signal_rx: mpsc::UnboundedReceiver<SignalEvent>,
    ) -> Self {
        Self {
            stdout_buffer: Arc::new(Mutex::new(Vec::new())),
            stdin_rx: Arc::new(tokio::sync::Mutex::new(stdin_rx)),
            signal_rx: Arc::new(tokio::sync::Mutex::new(signal_rx)),
            terminal_size: Size { rows: 24, cols: 80 },
        }
    }
}

impl std::fmt::Debug for TestClientOsApi {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TestClientOsApi").finish()
    }
}

impl ClientOsApi for TestClientOsApi {
    fn get_terminal_size(&self) -> Size {
        self.terminal_size
    }

    fn set_raw_mode(&mut self) {}

    fn unset_raw_mode(&self) -> Result<(), std::io::Error> {
        Ok(())
    }

    fn get_stdout_writer(&self) -> Box<dyn Write> {
        Box::new(TestWriter {
            buffer: self.stdout_buffer.clone(),
        })
    }

    fn get_stdin_reader(&self) -> Box<dyn io::BufRead> {
        Box::new(io::Cursor::new(Vec::new()))
    }

    fn update_session_name(&mut self, _new_session_name: String) {}

    fn read_from_stdin(&mut self) -> Result<Vec<u8>, &'static str> {
        Ok(Vec::new())
    }

    fn box_clone(&self) -> Box<dyn ClientOsApi> {
        Box::new(self.clone())
    }

    fn send_to_server(&self, _msg: ClientToServerMsg) {}

    fn recv_from_server(&self) -> Option<(ServerToClientMsg, ErrorContext)> {
        None
    }

    fn handle_signals(
        &self,
        _sigwinch_cb: Box<dyn Fn()>,
        _quit_cb: Box<dyn Fn()>,
        _resize_receiver: Option<std::sync::mpsc::Receiver<()>>,
    ) {
    }

    fn connect_to_server(&self, _path: &std::path::Path) {}

    fn load_palette(&self) -> Palette {
        Palette::default()
    }

    fn enable_mouse(&self) -> anyhow::Result<()> {
        Ok(())
    }

    fn disable_mouse(&self) -> anyhow::Result<()> {
        Ok(())
    }

    fn get_async_stdin_reader(&self) -> Box<dyn AsyncStdin> {
        Box::new(MockAsyncStdin {
            rx: self.stdin_rx.clone(),
        })
    }

    fn get_async_signal_listener(&self) -> io::Result<Box<dyn AsyncSignals>> {
        Ok(Box::new(MockAsyncSignals {
            rx: self.signal_rx.clone(),
        }))
    }
}

struct TestWriter {
    buffer: Arc<Mutex<Vec<u8>>>,
}

impl Write for TestWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buffer.lock().unwrap().extend_from_slice(buf);
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

mod mock_ws_server {
    use super::*;
    use axum::{
        extract::{
            ws::{WebSocket, WebSocketUpgrade},
            State,
        },
        routing::get,
        Router,
    };
    use tokio::net::TcpListener;
    use tokio::task::JoinHandle;

    #[derive(Clone)]
    struct WsState {
        client_tx: Arc<Mutex<mpsc::UnboundedSender<Message>>>,
        server_rx: Arc<Mutex<mpsc::UnboundedReceiver<Message>>>,
    }

    pub struct MockWsServer {
        pub terminal_to_client_tx: mpsc::UnboundedSender<Message>,
        pub control_to_client_tx: mpsc::UnboundedSender<Message>,
        pub client_to_terminal_rx: Arc<Mutex<mpsc::UnboundedReceiver<Message>>>,
        pub client_to_control_rx: Arc<Mutex<mpsc::UnboundedReceiver<Message>>>,
    }

    impl MockWsServer {
        pub async fn start() -> (u16, Self, JoinHandle<()>) {
            let (terminal_to_client_tx, terminal_to_client_rx) = mpsc::unbounded_channel();
            let (client_to_terminal_tx, client_to_terminal_rx) = mpsc::unbounded_channel();
            let (control_to_client_tx, control_to_client_rx) = mpsc::unbounded_channel();
            let (client_to_control_tx, client_to_control_rx) = mpsc::unbounded_channel();

            let terminal_state = WsState {
                client_tx: Arc::new(Mutex::new(client_to_terminal_tx)),
                server_rx: Arc::new(Mutex::new(terminal_to_client_rx)),
            };

            let control_state = WsState {
                client_tx: Arc::new(Mutex::new(client_to_control_tx)),
                server_rx: Arc::new(Mutex::new(control_to_client_rx)),
            };

            let app = Router::new()
                .route(
                    "/ws/terminal",
                    get(terminal_handler).with_state(terminal_state.clone()),
                )
                .route(
                    "/ws/control",
                    get(control_handler).with_state(control_state.clone()),
                );

            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let port = listener.local_addr().unwrap().port();

            let server_handle = tokio::spawn(async move {
                axum::serve(listener, app).await.unwrap();
            });

            tokio::time::sleep(Duration::from_millis(100)).await;

            let server = MockWsServer {
                terminal_to_client_tx,
                control_to_client_tx,
                client_to_terminal_rx: Arc::new(Mutex::new(client_to_terminal_rx)),
                client_to_control_rx: Arc::new(Mutex::new(client_to_control_rx)),
            };

            (port, server, server_handle)
        }
    }

    async fn terminal_handler(
        ws: WebSocketUpgrade,
        State(state): State<WsState>,
    ) -> impl axum::response::IntoResponse {
        ws.on_upgrade(move |socket| handle_websocket(socket, state))
    }

    async fn control_handler(
        ws: WebSocketUpgrade,
        State(state): State<WsState>,
    ) -> impl axum::response::IntoResponse {
        ws.on_upgrade(move |socket| handle_websocket(socket, state))
    }

    async fn handle_websocket(socket: WebSocket, state: WsState) {
        let (mut ws_tx, mut ws_rx) = socket.split();
        let client_tx = state.client_tx.clone();

        let recv_task = tokio::spawn(async move {
            while let Some(Ok(axum_msg)) = ws_rx.next().await {
                let tungstenite_msg = match axum_msg {
                    axum::extract::ws::Message::Text(text) => {
                        Message::Text(text.to_string().into())
                    },
                    axum::extract::ws::Message::Binary(data) => {
                        Message::Binary(data.to_vec().into())
                    },
                    axum::extract::ws::Message::Ping(data) => Message::Ping(data.to_vec().into()),
                    axum::extract::ws::Message::Pong(data) => Message::Pong(data.to_vec().into()),
                    axum::extract::ws::Message::Close(frame) => {
                        if let Some(f) = frame {
                            Message::Close(Some(tokio_tungstenite::tungstenite::protocol::CloseFrame {
                                code: tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::from(f.code),
                                reason: f.reason.to_string().into(),
                            }))
                        } else {
                            Message::Close(None)
                        }
                    },
                };
                let _ = client_tx.lock().unwrap().send(tungstenite_msg);
            }
        });

        loop {
            let msg = state.server_rx.lock().unwrap().try_recv();
            match msg {
                Ok(tungstenite_msg) => {
                    let axum_msg = match tungstenite_msg {
                        Message::Text(text) => {
                            axum::extract::ws::Message::Text(text.to_string().into())
                        },
                        Message::Binary(data) => {
                            axum::extract::ws::Message::Binary(data.to_vec().into())
                        },
                        Message::Ping(data) => {
                            axum::extract::ws::Message::Ping(data.to_vec().into())
                        },
                        Message::Pong(data) => {
                            axum::extract::ws::Message::Pong(data.to_vec().into())
                        },
                        Message::Close(frame) => {
                            if let Some(f) = frame {
                                axum::extract::ws::Message::Close(Some(
                                    axum::extract::ws::CloseFrame {
                                        code: f.code.into(),
                                        reason: f.reason.to_string().into(),
                                    },
                                ))
                            } else {
                                axum::extract::ws::Message::Close(None)
                            }
                        },
                        Message::Frame(_) => continue,
                    };
                    if ws_tx.send(axum_msg).await.is_err() {
                        break;
                    }
                },
                Err(_) => {
                    tokio::time::sleep(Duration::from_millis(10)).await;
                },
            }
        }

        recv_task.abort();
    }
}

#[tokio::test]
#[serial]
async fn test_stdin_forwarded_to_terminal_websocket() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-stdin".to_string(),
    };

    let (stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = Box::new(TestClientOsApi::new(stdin_rx, signal_rx));

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    let test_data = b"hello from stdin\n".to_vec();
    stdin_tx.send(test_data.clone()).unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;
    let received = tokio::time::timeout(
        Duration::from_secs(1),
        server.client_to_terminal_rx.lock().unwrap().recv(),
    )
    .await
    .expect("Timeout")
    .expect("No message");

    match received {
        Message::Binary(data) => assert_eq!(data, test_data),
        _ => panic!("Expected Binary message, got: {:?}", received),
    }

    drop(stdin_tx);
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}

#[tokio::test]
#[serial]
async fn test_terminal_output_written_to_stdout() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-stdout".to_string(),
    };

    let (_stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = TestClientOsApi::new(stdin_rx, signal_rx);
    let stdout_buffer = os_input.stdout_buffer.clone();
    let os_input = Box::new(os_input);

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    let test_output = "Hello from terminal";
    server
        .terminal_to_client_tx
        .send(Message::Text(test_output.to_string().into()))
        .unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;

    let stdout = stdout_buffer.lock().unwrap().clone();
    let stdout_str = String::from_utf8_lossy(&stdout);
    assert!(
        stdout_str.contains(test_output),
        "Expected stdout to contain '{}', got: '{}'",
        test_output,
        stdout_str
    );

    server
        .terminal_to_client_tx
        .send(Message::Close(None))
        .unwrap();
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}

#[tokio::test]
#[serial]
async fn test_resize_signal_sends_control_message() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-resize".to_string(),
    };

    let (_stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = Box::new(TestClientOsApi::new(stdin_rx, signal_rx));

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    tokio::time::sleep(Duration::from_millis(200)).await;

    let _ = tokio::time::timeout(
        Duration::from_millis(500),
        server.client_to_control_rx.lock().unwrap().recv(),
    )
    .await;

    signal_tx.send(SignalEvent::Resize).unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;

    let received = tokio::time::timeout(
        Duration::from_secs(1),
        server.client_to_control_rx.lock().unwrap().recv(),
    )
    .await
    .expect("Timeout")
    .expect("No message");

    match received {
        Message::Text(text) => {
            let parsed: WebClientToWebServerControlMessage =
                serde_json::from_str(&text).expect("Failed to parse");
            assert!(
                matches!(
                    parsed.payload,
                    WebClientToWebServerControlMessagePayload::TerminalResize(_)
                ),
                "Expected TerminalResize, got: {:?}",
                parsed.payload
            );
        },
        _ => panic!("Expected Text message, got: {:?}", received),
    }

    signal_tx.send(SignalEvent::Quit).unwrap();
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}

#[tokio::test]
#[serial]
async fn test_quit_signal_exits_loop() {
    let (port, _server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-quit".to_string(),
    };

    let (_stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = Box::new(TestClientOsApi::new(stdin_rx, signal_rx));

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    signal_tx.send(SignalEvent::Quit).unwrap();

    let result = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit within timeout")
        .expect("Loop panicked");

    assert!(result.is_ok(), "Expected Ok result, got: {:?}", result);
}

#[tokio::test]
#[serial]
async fn test_websocket_close_exits_loop() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-close".to_string(),
    };

    let (_stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = Box::new(TestClientOsApi::new(stdin_rx, signal_rx));

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    server
        .terminal_to_client_tx
        .send(Message::Close(None))
        .unwrap();

    let result = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit within timeout")
        .expect("Loop panicked");

    assert!(result.is_ok(), "Expected Ok result, got: {:?}", result);
}

#[tokio::test]
#[serial]
async fn test_control_message_handling() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;

    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);

    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();

    let connections = WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-control".to_string(),
    };

    let (_stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();

    let os_input = TestClientOsApi::new(stdin_rx, signal_rx);
    let terminal_size = os_input.terminal_size;
    let os_input = Box::new(os_input);

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    tokio::time::sleep(Duration::from_millis(200)).await;

    let _ = tokio::time::timeout(
        Duration::from_millis(500),
        server.client_to_control_rx.lock().unwrap().recv(),
    )
    .await;

    let query_msg = WebServerToWebClientControlMessage::QueryTerminalSize;
    server
        .control_to_client_tx
        .send(Message::Text(
            serde_json::to_string(&query_msg).unwrap().into(),
        ))
        .unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;

    let received = tokio::time::timeout(
        Duration::from_secs(1),
        server.client_to_control_rx.lock().unwrap().recv(),
    )
    .await
    .expect("Timeout")
    .expect("No message");

    match received {
        Message::Text(text) => {
            let parsed: WebClientToWebServerControlMessage =
                serde_json::from_str(&text).expect("Failed to parse");
            let size = match parsed.payload {
                WebClientToWebServerControlMessagePayload::TerminalResize(size) => size,
                WebClientToWebServerControlMessagePayload::TerminalResizeRendering(size) => size,
                other => panic!("Expected TerminalResize, got: {:?}", other),
            };
            assert_eq!(size, terminal_size);
        },
        _ => panic!("Expected Text message, got: {:?}", received),
    }

    let log_msg = WebServerToWebClientControlMessage::Log {
        lines: vec!["Test log".to_string()],
    };
    server
        .control_to_client_tx
        .send(Message::Text(
            serde_json::to_string(&log_msg).unwrap().into(),
        ))
        .unwrap();

    tokio::time::sleep(Duration::from_millis(100)).await;

    server
        .terminal_to_client_tx
        .send(Message::Close(None))
        .unwrap();
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}

async fn connect_mock_ws(port: u16) -> WebSocketConnections {
    let terminal_url = format!("ws://127.0.0.1:{}/ws/terminal", port);
    let control_url = format!("ws://127.0.0.1:{}/ws/control", port);
    let terminal_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (terminal_ws, _) = tokio_tungstenite::client_async_with_config(
        &terminal_url,
        MaybeTls::Plain(terminal_tcp),
        None,
    )
    .await
    .unwrap();
    let control_tcp = TcpStream::connect(format!("127.0.0.1:{}", port))
        .await
        .unwrap();
    let (control_ws, _) = tokio_tungstenite::client_async_with_config(
        &control_url,
        MaybeTls::Plain(control_tcp),
        None,
    )
    .await
    .unwrap();
    WebSocketConnections {
        terminal_ws,
        control_ws,
        web_client_id: "test-nested".to_string(),
    }
}

#[tokio::test]
#[serial]
async fn test_nested_ping_frame_is_answered_locally_and_stripped_from_forwarded_stdin() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;
    let connections = connect_mock_ws(port).await;

    let (stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();
    let os_input = TestClientOsApi::new(stdin_rx, signal_rx);
    let stdout_buffer = os_input.stdout_buffer.clone();
    let os_input = Box::new(os_input);

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    let mut chunk = b"abc".to_vec();
    chunk.extend_from_slice(&zellij_utils::nested_session::encode_frame(
        &zellij_utils::nested_session::NestedSessionMessage::Ping,
    ));
    chunk.extend_from_slice(b"def");
    stdin_tx.send(chunk).unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;
    let received = tokio::time::timeout(
        Duration::from_secs(1),
        server.client_to_terminal_rx.lock().unwrap().recv(),
    )
    .await
    .expect("Timeout")
    .expect("No message");
    match received {
        Message::Binary(data) => assert_eq!(data, b"abcdef".to_vec()),
        _ => panic!("Expected Binary message, got: {:?}", received),
    }

    let pong_frame = zellij_utils::nested_session::encode_frame(
        &zellij_utils::nested_session::NestedSessionMessage::Pong,
    );
    let stdout_bytes = stdout_buffer.lock().unwrap().clone();
    assert!(
        stdout_bytes
            .windows(pong_frame.len())
            .any(|window| window == &pong_frame[..]),
        "expected a pong frame on stdout, got: {:?}",
        stdout_bytes
    );

    drop(stdin_tx);
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}

#[tokio::test]
#[serial]
async fn test_nested_announce_ack_is_relayed_over_control_websocket() {
    let (port, server, _server_handle) = mock_ws_server::MockWsServer::start().await;
    let connections = connect_mock_ws(port).await;

    let (stdin_tx, stdin_rx) = mpsc::unbounded_channel();
    let (_signal_tx, signal_rx) = mpsc::unbounded_channel();
    let os_input = Box::new(TestClientOsApi::new(stdin_rx, signal_rx));

    let loop_handle =
        tokio::spawn(
            async move { run_remote_client_terminal_loop(os_input, connections, None).await },
        );

    let announce_ack = zellij_utils::nested_session::NestedSessionMessage::AnnounceAck {
        ancestry: vec!["outer-host".to_owned()],
        capabilities: vec![zellij_utils::nested_session::NestedSessionCapability::NestedControl],
        descend_keys: vec![],
    };
    stdin_tx
        .send(zellij_utils::nested_session::encode_frame(&announce_ack))
        .unwrap();

    tokio::time::sleep(Duration::from_millis(200)).await;
    let mut relayed_payload = None;
    loop {
        let received = tokio::time::timeout(
            Duration::from_secs(1),
            server.client_to_control_rx.lock().unwrap().recv(),
        )
        .await
        .expect("Timeout")
        .expect("No message");
        if let Message::Text(text) = received {
            let control_msg: WebClientToWebServerControlMessage =
                serde_json::from_str(&text).unwrap();
            if let WebClientToWebServerControlMessagePayload::NestedSessionFrameFromHost {
                payload_bytes,
            } = control_msg.payload
            {
                assert_eq!(control_msg.web_client_id, "test-nested");
                relayed_payload = Some(payload_bytes);
                break;
            }
        }
    }
    assert_eq!(
        zellij_utils::nested_session::decode_payload(&relayed_payload.unwrap()),
        Some(announce_ack)
    );

    drop(stdin_tx);
    let _ = tokio::time::timeout(Duration::from_secs(2), loop_handle)
        .await
        .expect("Loop didn't exit")
        .unwrap();
}