//! # rutster-tap-echo — the Rust reference echo brain + test server (spec §2.3, §8.4)
//!
//! Dual-purpose crate:
//! - **Standalone binary** (`cargo run -p rutster-tap-echo`): binds
//!   `ws://127.0.0.1:8081/echo` and echoes `audio_in` → `audio_out` per the
//!   slice-2 protocol (spec §3).
//! - **In-process `start_echo_server`** (lib): used by `rutster`'s integration
//!   tests to drive the tap end-to-end without an external process.
//!
//! ## Why a Rust brain at all (when the canonical brain is Python?)
//!
//! Reuses `rutster-tap`'s protocol types — **the contract test that the wire
//! types are reusable from outside the core**. Any future brain written in
//! Rust (or a step-3 OpenAI adapter in Rust) starts from this shape. The
//! Python brain (`examples/echo_brain/`) proves language-agnosticism; this
//! crate proves reusability + powers the in-process integration tests.
//!
//! ## Stateless contract (spec §5.3)
//!
//! The echo brain holds no per-call state across reconnects. On a fresh
//! `hello` it starts a new logical session; the core treats every reconnect
//! as a resume with the same `session_id`, but the brain just acks and
//! echoes. This is the resilience posture slice-2 proves.

use std::net::SocketAddr;

use futures_util::{SinkExt, StreamExt};
use rutster_tap::protocol::{
    DecodedPayload, HelloPayload, decode_envelope, encode_audio_out, encode_bye, encode_error,
    encode_hello,
};
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tokio_tungstenite::WebSocketStream;
use tracing::{info, warn};

/// Handle returned by `start_echo_server`; drop the `shutdown` sender to stop.
pub struct EchoHandle {
    pub shutdown: oneshot::Sender<()>,
    pub join: JoinHandle<()>,
    pub addr: SocketAddr,
}

/// Start an in-process echo brain bound to `addr`. Returns once the socket
/// is bound; the accept loop runs in a spawned task.
pub async fn start_echo_server(addr: SocketAddr) -> Result<EchoHandle, std::io::Error> {
    let listener = tokio::net::TcpListener::bind(addr).await?;
    let bound_addr = listener.local_addr()?;
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    let join = tokio::spawn(async move {
        accept_loop(listener, shutdown_rx).await;
    });
    Ok(EchoHandle {
        shutdown: shutdown_tx,
        join,
        addr: bound_addr,
    })
}

async fn accept_loop(listener: tokio::net::TcpListener, mut shutdown: oneshot::Receiver<()>) {
    // slice-2 §8.4 reconnect-path test: the integration test kills the
    // brain mid-call by firing `shutdown`. To actually sever the live WS
    // connection (not just stop accepting new ones), we track each spawned
    // per-connection task here and abort them all on shutdown. Without
    // this, the engine's `ws.next()` would never return None (the existing
    // TCP stream stays half-open after the listener drops), and
    // `reconnect_attempts` would never increment — making the mandated
    // reconnect assertion unobservable from a test.
    let mut conns: Vec<JoinHandle<()>> = Vec::new();
    loop {
        tokio::select! {
            _ = &mut shutdown => {
                info!(
                    conn_count = conns.len(),
                    "echo server shutting down; aborting live connections"
                );
                for handle in conns.drain(..) {
                    handle.abort();
                }
                return;
            }
            res = listener.accept() => {
                let (stream, peer) = match res {
                    Ok(s) => s,
                    Err(e) => {
                        warn!(error = %e, "accept failed; continuing");
                        continue;
                    }
                };
                let task = tokio::spawn(async move {
                    let ws = match tokio_tungstenite::accept_async(stream).await {
                        Ok(ws) => ws,
                        Err(e) => {
                            warn!(error = %e, %peer, "ws upgrade failed");
                            return;
                        }
                    };
                    if let Err(e) = echo_one_connection(ws).await {
                        warn!(error = ?e, %peer, "echo connection ended with error");
                    }
                });
                conns.push(task);
                // Periodic cleanup of finished handles so the Vec doesn't
                // grow unbounded across many short-lived connections.
                conns.retain(|h| !h.is_finished());
            }
        }
    }
}

/// Per-connection echo loop. Unit-testable (we can drive a `WebSocketStream`
/// with synthetic frames in a test).
pub async fn echo_one_connection<T>(
    mut ws: WebSocketStream<T>,
) -> Result<(), Box<dyn std::error::Error>>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    let mut seq_egress: u64 = 0;
    // Wait for hello; ack with hello.
    let hello_in = ws
        .next()
        .await
        .ok_or("brain: connection closed before hello")??;
    // tungstenite 0.24: `Message::into_text()` returns `Result<String>`,
    // not `Option<String>`. `map_err` collapses the error to a `&str` for
    // the `Box<dyn Error>` path.
    let hello_text = hello_in.into_text().map_err(|_| "brain: hello not text")?;
    let decoded = decode_envelope(&hello_text)?;
    let session_id = match decoded.payload {
        DecodedPayload::Hello(HelloPayload { session_id, .. }) => session_id,
        _ => return Err("brain: first frame not hello".into()),
    };
    info!(%session_id, "brain: hello received; acking");
    let ack = encode_hello(&session_id, seq_egress, 0)?;
    seq_egress += 1;
    ws.send(tokio_tungstenite::tungstenite::Message::Text(ack))
        .await?;

    // Echo loop: audio_in → audio_out (same PCM).
    while let Some(msg) = ws.next().await {
        let msg = msg?;
        // `into_text()` on 0.24 returns `Result<String>`, so `let Ok(...) else`
        // is the natural pattern for "ignore binary / decode failures and
        // continue the loop" (drop + observe, spec §3.8).
        let Ok(text) = msg.into_text() else {
            // Binary frames ignored (v1 text-JSON only).
            continue;
        };
        let decoded = match decode_envelope(&text) {
            Ok(d) => d,
            Err(e) => {
                let err_frame = encode_error("decode_failed", &e.to_string(), seq_egress, 0)?;
                let _ = ws
                    .send(tokio_tungstenite::tungstenite::Message::Text(err_frame))
                    .await;
                continue;
            }
        };
        match decoded.payload {
            DecodedPayload::AudioIn(audio) => {
                // Echo: same PCM, same samples count.
                let out_frame = rutster_tap::protocol::decode_pcm(&audio.pcm, audio.samples)?;
                let out_str = encode_audio_out(&out_frame, seq_egress, decoded.ts)?;
                seq_egress += 1;
                ws.send(tokio_tungstenite::tungstenite::Message::Text(out_str))
                    .await?;
            }
            DecodedPayload::Bye(p) => {
                info!(reason = %p.reason, "brain: bye received; closing");
                let bye_ack = encode_bye("brain_ack", seq_egress, 0)?;
                let _ = ws
                    .send(tokio_tungstenite::tungstenite::Message::Text(bye_ack))
                    .await;
                let _ = ws.close(None).await;
                return Ok(());
            }
            DecodedPayload::SessionEnd(p) => {
                info!(reason = %p.reason, "brain: session_end received; closing");
                let _ = ws.close(None).await;
                return Ok(());
            }
            _ => {
                // Unknown / unexpected; ignore (drop + observe).
                warn!("brain: ignoring unexpected frame kind");
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    // `PcmFrame` is re-exported by `rutster-tap` from `rutster-media` (spec
    // §3.1 — one canonical home). Test consumes the re-export rather than
    // adding `rutster-media` as a second path dep on this crate.
    use rutster_tap::PcmFrame;
    use rutster_tap::protocol::{encode_audio_in, encode_bye, encode_hello};

    #[tokio::test]
    async fn echo_round_trips_one_audio_frame() {
        // Two ends of an in-memory WS pair. tokio_tungstenite doesn't ship
        // a direct in-mem channel impl, so use a TCP loopback pair.
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (s, _) = listener.accept().await.unwrap();
            let ws = tokio_tungstenite::accept_async(s).await.unwrap();
            echo_one_connection(ws).await.unwrap();
        });

        let client_stream = tokio::net::TcpStream::connect(addr).await.unwrap();
        // tungstenite 0.24: `Request::default()` is not a usable default for
        // `client_async`. Use `IntoClientRequest` impls (`&str` / `String`
        // / `url::Url` all work). `format!("ws://...")` constructs the
        // request in one line — simpler than building a `Request::builder()`
        // chain. The test asserts `echo_one_connection`'s behavior, not the
        // client-side handshake API.
        let (mut client_ws, _resp) =
            tokio_tungstenite::client_async(format!("ws://{addr}/echo"), client_stream)
                .await
                .unwrap();

        // Send hello.
        let hello = encode_hello("test-session-id", 0, 0).unwrap();
        client_ws
            .send(tokio_tungstenite::tungstenite::Message::Text(hello))
            .await
            .unwrap();

        // Receive hello ack.
        let ack = client_ws
            .next()
            .await
            .unwrap()
            .unwrap()
            .into_text()
            .unwrap();
        assert!(ack.contains("\"type\":\"hello\""));
        assert!(ack.contains("test-session-id"));

        // Send audio_in.
        let mut frame = PcmFrame::zeroed();
        frame.samples[0] = 42;
        let audio_in = encode_audio_in(&frame, 1, 100).unwrap();
        client_ws
            .send(tokio_tungstenite::tungstenite::Message::Text(audio_in))
            .await
            .unwrap();

        // Receive audio_out — should be the same PCM.
        let audio_out = client_ws
            .next()
            .await
            .unwrap()
            .unwrap()
            .into_text()
            .unwrap();
        assert!(audio_out.contains("\"type\":\"audio_out\""));
        let decoded = decode_envelope(&audio_out).unwrap();
        match decoded.payload {
            DecodedPayload::AudioOut(p) => {
                let echoed = rutster_tap::protocol::decode_pcm(&p.pcm, p.samples).unwrap();
                assert_eq!(echoed.samples[0], 42);
            }
            _ => panic!("expected audio_out"),
        }

        // Bye.
        let bye = encode_bye("done", 2, 200).unwrap();
        client_ws
            .send(tokio_tungstenite::tungstenite::Message::Text(bye))
            .await
            .unwrap();
        // Wait for the bye-ack so the server's `close(None)` flushes
        // before the spawned task resolves (otherwise the client's next
        // read might race the close). Receive one final message — either
        // the bye-ack text frame or the close frame — then join.
        let _ = client_ws.next().await;
        server.await.unwrap();
    }
}
