dittolive-ditto 5.0.3

use std::{str::FromStr, sync::Mutex};

use anyhow::Result;
use dittolive_ditto::{
    fs::TempRoot,
    prelude::*,
    preview::{datastreams::*, peer_pubkey::PeerPubkey},
};
use rand::Rng;
use safer_ffi::bytes::Bytes;
use tokio::sync::mpsc::unbounded_channel;

pub fn get_ditto(database_id: Option<DatabaseId>) -> Result<(TempRoot, Ditto)> {
    let database_id = database_id.unwrap_or_else(DatabaseId::generate);
    let temp_root = TempRoot::new();
    let config = DittoConfig::new(
        database_id.to_string(),
        DittoConfigConnect::SmallPeersOnly { private_key: None },
    )
    .with_persistence_directory(temp_root.root_path());
    let ditto = Ditto::open_sync(config)?;
    ditto.set_license_from_env("DITTO_LICENSE")?;
    Ok((temp_root, ditto))
}

#[tokio::test]
#[ignore]
async fn non_existent_peer_fails() -> Result<()> {
    let (_temp_root, mut dittos) = setup(1).await?;
    let ditto = dittos.remove(0);
    let dsep = ditto.datastreams();

    let target = PeerPubkey::from_str("pkAg").unwrap();

    let res = dsep.connect(target, "any").finish_async().await;
    assert_eq!(res.unwrap_err(), ConnectionError::PeerNotFound);
    Ok(())
}

fn get_local_key(ditto: &Ditto) -> PeerPubkey {
    PeerPubkey::from_str(&ditto.presence().graph().local_peer.peer_key).unwrap()
}

#[tokio::test]
#[ignore]
// todo(frankie.foston): clarify expected behaviour here
async fn send_message_to_self() -> Result<()> {
    let (_temp_root, mut dittos) = setup(1).await?;
    let ditto = dittos.remove(0);
    let pk = get_local_key(&ditto);

    let dsep = ditto.datastreams();
    let mut acceptor = dsep.bind_topic("self").finish(unbounded_channel()).unwrap();

    let stream = dsep
        .connect(pk.clone(), "self")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .unwrap();
    let mut rx = acceptor.recv().await.unwrap().open(unbounded_channel());
    assert_eq!(stream.peer_pubkey(), pk);
    let res = stream.message("Hi").send().changed().await;
    assert_eq!(res, SendStatus::Sent);

    let msg = rx.recv().await.unwrap();
    assert_eq!(&*msg, b"Hello world!");

    Ok(())
}

async fn setup(n: usize) -> Result<(Vec<TempRoot>, Vec<Ditto>)> {
    let mut roots = vec![];
    let mut output = vec![];

    let database_id = DatabaseId::generate();

    let mut tc = TransportConfig::default();
    tc.peer_to_peer.bluetooth_le.enabled = false;
    tc.listen.tcp.enabled = true;
    tc.listen.tcp.interface_ip = "127.0.0.1".to_string();

    let port = rand::rng().random_range(10000..(65535 - n as u16));

    for i in 0..(n as u16) {
        let (_root, ditto) = get_ditto(Some(database_id.clone()))?;

        let mut tc = tc.clone();
        tc.listen.tcp.port = port + i;

        if i != 0 {
            let remote_port = port + i - 1;
            tc.connect
                .tcp_servers
                .insert(format!("127.0.0.1:{remote_port}"));
        }

        ditto.set_transport_config(tc.clone());
        ditto.sync().start()?;

        roots.push(_root);
        output.push(ditto);
    }

    for ditto in &output {
        while ditto.presence().graph().remote_peers.len() != n - 1 {
            // wait for all peers to connect
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    }

    Ok((roots, output))
}

#[tokio::test]
async fn on_connecting_cannot_refuse_link_connection() -> Result<()> {
    const N: usize = 3;
    let (_temp_root, ditto) = setup(N).await?;

    let pk1 = get_local_key(&ditto[0]);
    let pk2 = get_local_key(&ditto[N - 1]);
    assert_ne!(pk1, pk2);

    let dsep1 = ditto[0].datastreams();

    ditto[N - 1]
        .presence()
        .set_connection_request_handler(move |cr: ConnectionRequest| {
            let remote_pk = PeerPubkey::from_str(&cr.peer_key()).unwrap();

            // Will not be hit for a Link
            if remote_pk == pk1 {
                ConnectionRequestAuthorization::Deny
            } else {
                ConnectionRequestAuthorization::Allow
            }
        });
    let _acceptor = ditto[N - 1]
        .datastreams()
        .bind_topic("on_connect")
        .on_receive_factory(WriteOnly)
        .finish_with(|stream| core::mem::drop(stream))
        .unwrap();

    let _stream = dsep1
        .connect(pk2, "on_connect")
        .finish_async()
        .await
        .expect("Connection should succeed despite the `connection_request_handler`");

    Ok(())
}

#[tokio::test]
async fn setup_builds_correctly() -> Result<()> {
    let (_temp_root, ditto) = setup(5).await?;

    assert_eq!(5, ditto.len());
    Ok(())
}

#[tokio::test]
async fn big_message_fails() -> Result<()> {
    let (_temp_root, setup) = setup(2).await.unwrap();
    let a = setup.first().unwrap();
    let a_dsep = a.datastreams();
    let a_key = get_local_key(&a);
    let b = setup.last().unwrap();
    let b_dsep = b.datastreams();
    let mut a_acceptor = a_dsep
        .bind_topic("large")
        .on_receive_factory(tokio::sync::mpsc::unbounded_channel)
        .finish(tokio::sync::mpsc::unbounded_channel())
        .unwrap();
    let b_stream = b_dsep
        .connect(a_key.clone(), "large")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .unwrap();
    let mut a_rx = a_acceptor.recv().await.unwrap();

    let mut rng = rand::rng();
    // 100MB payload
    let mut payload = vec![0u8; 1024 * 1024 * 100];
    dbg!(payload.len());
    rng.fill(&mut payload[..]);

    let res = b_stream.message(Payload::from(payload));
    dbg!("Sending");
    let res = res.send();
    dbg!("Sent");
    let res = res.changed().await;
    assert_eq!(res, SendStatus::Sent);
    dbg!("Closing stream");
    b_stream.close(()).unwrap().await;
    dbg!("Closed stream");

    let Ok(recv) = tokio::time::timeout(std::time::Duration::from_millis(100), a_rx.recv()).await
    else {
        panic!("Stream closed on the other end didn't destroy its closure in time")
    };
    assert!(
        recv.is_none(),
        "The stream should get closed without having received the oversized message"
    );

    Ok(())
}

#[tokio::test]
async fn multiple_streams_between_two_peers() {
    let (_temp_root, setup) = setup(2).await.unwrap();
    let a = setup.first().unwrap();
    let a_dsep = a.datastreams();
    let a_key = get_local_key(&a);
    let b = setup.last().unwrap();
    let b_dsep = b.datastreams();
    let _acceptor = a_dsep
        .bind_topic("stream1")
        .on_receive_factory(WriteOnly)
        .finish(tokio::sync::mpsc::unbounded_channel());
    let _acceptor = a_dsep
        .bind_topic("stream2")
        .on_receive_factory(WriteOnly)
        .finish(tokio::sync::mpsc::unbounded_channel());
    let _b_stream = b_dsep
        .connect(a_key.clone(), "stream1")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .unwrap();
    let _b_stream = b_dsep
        .connect(a_key.clone(), "stream2")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .unwrap();
}

#[tokio::test]
async fn multiple_peers_sharing_one_topic() {
    let (_temp_root, setup) = setup(3).await.unwrap();
    for s in &setup {
        get_local_key(s);
    }
    let a = &setup[0];
    let a_dsep = a.datastreams();
    let a_key = get_local_key(&a);
    let b_dsep = setup[1].datastreams();
    let c_dsep = setup[2].datastreams();
    let _acceptor = a_dsep
        .bind_topic("sharing".as_bytes())
        .finish_with(|stream| {
            std::thread::spawn(move || {
                let stream = stream.open(std::sync::mpsc::channel());
                while let Ok(msg) = stream.recv() {
                    stream.message(msg).send();
                }
            });
        });
    let b_stream = b_dsep
        .connect(a_key.clone(), "sharing".as_bytes())
        .finish_async()
        .await
        .unwrap();
    let mut b_stream = b_stream.open(tokio::sync::mpsc::unbounded_channel());
    let c_stream = c_dsep
        .connect(a_key.clone(), "sharing".as_bytes())
        .finish_async()
        .await
        .expect("C and B should not interfere with each other");
    let mut c_stream = c_stream.open(tokio::sync::mpsc::unbounded_channel());
    b_stream.message("b".as_bytes()).send();
    c_stream.message("c".as_bytes()).send();
    assert_eq!(b_stream.recv().await.unwrap().as_slice(), b"b");
    assert_eq!(c_stream.recv().await.unwrap().as_slice(), b"c");
}

#[tokio::test]
async fn unique_stream_per_topic_for_two_peers() {
    let (_temp_root, setup) = setup(2).await.unwrap();
    let a = setup.first().unwrap();
    let a_dsep = a.datastreams();
    let a_key = get_local_key(&a);
    let b = setup.last().unwrap();
    let b_dsep = b.datastreams();
    let _acceptor = a_dsep
        .bind_topic("only_once")
        .on_receive_factory(WriteOnly)
        .finish(tokio::sync::mpsc::unbounded_channel());
    let b_stream = b_dsep
        .connect(a_key.clone(), "only_once")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .unwrap();
    _ = b_dsep
        .connect(a_key.clone(), "only_once")
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .expect_err(
            "Only a single connection should be accepted between two peers on a given topic",
        );
    b_stream.close(()).unwrap().await;
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
    _ = b_dsep
        .connect(a_key.clone(), "only_once".as_bytes())
        .on_receive_factory(WriteOnly)
        .finish_async()
        .await
        .expect("After closing the stream, reconnection should be possible");
}

async fn sequential<const N: usize>() {
    use std::{future::Future, pin::Pin};

    let (_temp_root, setup) = setup(N).await.unwrap();
    let a = setup.first().unwrap();
    let a_dsep = a.datastreams();
    let a_key = get_local_key(&a);
    let b = setup.last().unwrap();
    let b_dsep = b.datastreams();
    let (handles_tx, mut handles_rx) = tokio::sync::mpsc::unbounded_channel();
    let handles_tx = Mutex::new(Some(handles_tx));
    let _a_bytes_acceptor = a_dsep
        .bind_topic("bytes")
        .on_receive_factory(WriteOnly)
        .finish_with(move |stream| {
            let handles_tx = handles_tx.lock().unwrap().take().unwrap();
            for i in (0..2000u32).step_by(10) {
                let handle = stream
                    .message(Bytes::copied_from_slice(i.to_le_bytes().as_slice()))
                    .send();
                handles_tx.send(handle).unwrap();
            }
        })
        .unwrap();
    let mut bytes_sub = b_dsep
        .connect(a_key.clone(), "bytes")
        .finish_async()
        .await
        .unwrap()
        .open(tokio::sync::mpsc::unbounded_channel());
    let mut bytes_expected = 0;
    let mut handle_future: Option<SendHandleFuture> = None;
    let mut pending = std::future::pending();
    loop {
        if bytes_expected == 2000
            && handle_future.is_none()
            && handles_rx.is_empty()
            && handles_rx.is_closed()
        {
            break;
        }
        let needs_handle = handle_future.is_none();
        let handle_fut: Pin<&mut dyn Future<Output = SendStatus>> = unsafe {
            Pin::new_unchecked(
                handle_future
                    .as_mut()
                    .map_or(&mut pending as &mut dyn Future<Output = SendStatus>, |f| {
                        f as &mut dyn Future<Output = SendStatus>
                    }),
            )
        };
        tokio::select! {
            bytes = bytes_sub.recv(), if bytes_expected < 2000 => {
                let mut payload = u32::MAX.to_le_bytes();
                payload.copy_from_slice(&bytes.unwrap());
                assert_eq!(u32::from_le_bytes(payload), bytes_expected, "Out of sequence packet on stream `bytes`");
                bytes_expected += 10;
            }
            handle = handles_rx.recv(), if needs_handle => {
                handle_future = handle.map(|h|h.changed());
            }
            status = handle_fut => {
                assert_eq!(dbg!(status), SendStatus::Sent);
                handle_future = None;
            }
        }
    }
    core::mem::drop(bytes_sub);
}

#[tokio::test]
async fn singlehop_streams_are_sequential() {
    sequential::<2>().await
}

#[tokio::test]
async fn multihop_streams_are_sequential() {
    sequential::<4>().await
}

/// Ensures that attempting to connect to an acceptor that expects a different reliability than
/// expected will fail.
#[tokio::test]
async fn reliability_modes() {
    let (_temp_root, setup) = setup(2).await.unwrap();
    let _acceptor = setup[0]
        .datastreams()
        .bind_topic("reliable")
        .on_receive_factory(WriteOnly)
        .finish(tokio::sync::mpsc::unbounded_channel());
    assert_eq!(
        setup[1]
            .datastreams()
            .connect(get_local_key(&setup[0]), "reliable")
            .reliability(Reliability::Unreliable)
            .finish_async()
            .await
            .unwrap_err(),
        ConnectionError::ConnectionRejected,
        "Expected a rejection when connecting unreliably to a reliable stream"
    );
    core::mem::drop(
        setup[1]
            .datastreams()
            .connect(get_local_key(&setup[0]), "reliable")
            .finish_async()
            .await
            .unwrap(),
    );
    let _acceptor = setup[0]
        .datastreams()
        .bind_topic("unreliable")
        .reliability(Reliability::Unreliable)
        .on_receive_factory(WriteOnly)
        .finish(tokio::sync::mpsc::unbounded_channel());
    assert_eq!(
        setup[1]
            .datastreams()
            .connect(get_local_key(&setup[0]), "unreliable")
            .finish_async()
            .await
            .unwrap_err(),
        ConnectionError::ConnectionRejected,
        "Expected a rejection when connecting reliably to an unreliable stream"
    );
    core::mem::drop(
        setup[1]
            .datastreams()
            .connect(get_local_key(&setup[0]), "unreliable")
            .reliability(Reliability::Unreliable)
            .finish_async()
            .await
            .unwrap(),
    );
}