use std::net::SocketAddr;
use std::pin::Pin;
use std::time::Duration;

use tokio_stream::Stream;
use tonic::{Request, Response, Status};

pub mod pipe_grpc_test {
    tonic::include_proto!("pipe_grpc_test");
}

use pipe_grpc_test::test_streaming_client::TestStreamingClient;
use pipe_grpc_test::test_streaming_server::{TestStreaming, TestStreamingServer};
use pipe_grpc_test::{StreamItem, StreamRequest};

#[derive(Default)]
struct TestService;

#[tonic::async_trait]
impl TestStreaming for TestService {
    type ServerStreamStream = Pin<Box<dyn Stream<Item = Result<StreamItem, Status>> + Send>>;

    async fn server_stream(
        &self,
        request: Request<StreamRequest>,
    ) -> Result<Response<Self::ServerStreamStream>, Status> {
        let count = request.into_inner().count;
        let (tx, rx) = tokio::sync::mpsc::channel(16);
        tokio::spawn(async move {
            for i in 0..count {
                let item = StreamItem {
                    value: i,
                    label: format!("item-{i}"),
                };
                if tx.send(Ok(item)).await.is_err() {
                    break;
                }
            }
        });
        Ok(Response::new(Box::pin(
            tokio_stream::wrappers::ReceiverStream::new(rx),
        )))
    }
}

async fn start_server() -> SocketAddr {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(TestStreamingServer::new(TestService))
            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
            .await
            .unwrap();
    });
    tokio::time::sleep(Duration::from_millis(50)).await;
    addr
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn collect_all_items() {
    let addr = start_server().await;
    let mut client = TestStreamingClient::connect(format!("http://{addr}"))
        .await
        .unwrap();

    let response = client
        .server_stream(Request::new(StreamRequest { count: 5 }))
        .await
        .unwrap();

    let items: Vec<StreamItem> = lazyflow_grpc::streaming::from_tonic(response.into_inner())
        .collect()
        .await
        .unwrap();

    assert_eq!(items.len(), 5);
    for (i, item) in items.iter().enumerate() {
        assert_eq!(item.value, i as i32);
        assert_eq!(item.label, format!("item-{i}"));
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn pipe_operators_compose() {
    let addr = start_server().await;
    let mut client = TestStreamingClient::connect(format!("http://{addr}"))
        .await
        .unwrap();

    let response = client
        .server_stream(Request::new(StreamRequest { count: 10 }))
        .await
        .unwrap();

    let evens: Vec<i32> = lazyflow_grpc::streaming::from_tonic(response.into_inner())
        .map(|item| item.value)
        .filter(|v| *v % 2 == 0)
        .collect()
        .await
        .unwrap();

    assert_eq!(evens, vec![0, 2, 4, 6, 8]);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn take_cancels_early() {
    let addr = start_server().await;
    let mut client = TestStreamingClient::connect(format!("http://{addr}"))
        .await
        .unwrap();

    let response = client
        .server_stream(Request::new(StreamRequest { count: 10000 }))
        .await
        .unwrap();

    let items: Vec<StreamItem> = lazyflow_grpc::streaming::from_tonic(response.into_inner())
        .take(3)
        .collect()
        .await
        .unwrap();

    assert_eq!(items.len(), 3);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn empty_stream() {
    let addr = start_server().await;
    let mut client = TestStreamingClient::connect(format!("http://{addr}"))
        .await
        .unwrap();

    let response = client
        .server_stream(Request::new(StreamRequest { count: 0 }))
        .await
        .unwrap();

    let items: Vec<StreamItem> = lazyflow_grpc::streaming::from_tonic(response.into_inner())
        .collect()
        .await
        .unwrap();

    assert!(items.is_empty());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn error_maps_to_pipe_error() {
    let status = tonic::Status::not_found("gone");
    let err = lazyflow_grpc::streaming::status_to_pipe_error(status);
    let msg = err.to_string();
    assert!(msg.contains("gone"), "expected 'gone' in: {msg}");
}
