use super::test_utils::random_peer_url;
use crate::{
    default_test_builder,
    factories::core_fetch::{
        CoreFetch, CoreFetchConfig, test::test_utils::make_op,
    },
};
use kitsune2_api::*;
use kitsune2_test_utils::{
    enable_tracing,
    id::{create_op_id_list, random_op_id},
    iter_check,
    space::TEST_SPACE_ID,
};
use prost::Message;
use std::{
    collections::HashSet,
    sync::{Arc, Mutex},
    time::Duration,
};

type RequestsSent = Vec<(OpId, Url)>;

struct TestCase {
    fetch: CoreFetch,
    requests_sent: Arc<Mutex<RequestsSent>>,
    op_store: DynOpStore,
    peer_meta_store: DynPeerMetaStore,
    _transport: DynTransport,
}

async fn setup_test(config: &CoreFetchConfig) -> TestCase {
    let builder =
        Arc::new(default_test_builder().with_default_config().unwrap());
    let op_store = builder
        .op_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();
    let peer_meta_store = builder
        .peer_meta_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();
    let requests_sent = Arc::new(Mutex::new(Vec::new()));
    let mock_transport = make_mock_transport(requests_sent.clone());
    let report = builder
        .report
        .create(builder.clone(), mock_transport.clone())
        .await
        .unwrap();

    let fetch = CoreFetch::new(
        config.clone(),
        TEST_SPACE_ID,
        report,
        op_store.clone(),
        peer_meta_store.clone(),
        mock_transport.clone(),
    );

    TestCase {
        fetch,
        requests_sent,
        op_store,
        peer_meta_store,
        _transport: mock_transport,
    }
}

fn make_mock_transport(
    requests_sent: Arc<Mutex<RequestsSent>>,
) -> Arc<MockTransport> {
    let mut mock_transport = MockTransport::new();
    mock_transport.expect_send_module().returning({
        let requests_sent = requests_sent.clone();
        move |peer, space_id, module_id, data| {
            assert_eq!(space_id, TEST_SPACE_ID);
            assert_eq!(module_id, crate::factories::core_fetch::MOD_NAME);
            Box::pin({
                let requests_sent = requests_sent.clone();
                async move {
                    let fetch_message = K2FetchMessage::decode(data).unwrap();
                    let op_ids: Vec<OpId> =
                        match fetch_message.fetch_message_type() {
                            FetchMessageType::Request => {
                                let request =
                                    FetchRequest::decode(fetch_message.data)
                                        .unwrap();
                                request.into()
                            }
                            _ => panic!("unexpected fetch message type"),
                        };
                    let mut lock = requests_sent.lock().unwrap();
                    op_ids.into_iter().for_each(|op_id| {
                        lock.push((op_id, peer.clone()));
                    });

                    Ok(())
                }
            })
        }
    });
    mock_transport
        .expect_register_module_handler()
        .returning(|_, _, _| ());
    Arc::new(mock_transport)
}

// Test that ops that couldn't be added to the request queue are purged from state.
#[tokio::test(flavor = "multi_thread")]
async fn request_op_failures_are_purged_from_state() {
    enable_tracing();
    let TestCase { fetch, .. } = setup_test(&CoreFetchConfig::default()).await;

    // Request an op that will get added to state.
    let successful_op_id = random_op_id();
    let successful_peer_url = random_peer_url();
    fetch
        .request_ops(
            vec![successful_op_id.clone()],
            successful_peer_url.clone(),
        )
        .await
        .unwrap();

    let mut expected_state = HashSet::new();
    expected_state.insert((successful_op_id, successful_peer_url));
    assert_eq!(fetch.state.lock().unwrap().requests, expected_state);

    // Drop the channel receiver, which makes sending to the channel fail.
    fetch.tasks.iter().for_each(|t| t.abort());
    // Wait until tasks are actually aborted.
    iter_check!({
        if fetch.tasks.iter().all(|t| t.is_finished()) {
            break;
        }
    });

    // Request an op that will get added to state.
    fetch
        .request_ops(create_op_id_list(1), random_peer_url())
        .await
        .unwrap();

    // State should still only contain the successfully added op id.
    assert_eq!(fetch.state.lock().unwrap().requests, expected_state);
}

#[tokio::test(flavor = "multi_thread")]
async fn happy_op_fetch_from_multiple_agents() {
    let config = CoreFetchConfig {
        parallel_request_count: 5,
    };
    let TestCase {
        fetch,
        requests_sent,
        _transport,
        ..
    } = setup_test(&config).await;

    let op_list_1 = create_op_id_list(10);
    let op_list_2 = create_op_id_list(20);
    let op_list_3 = create_op_id_list(30);
    let total_ops = op_list_1.len() + op_list_2.len() + op_list_3.len();

    let peer_url_1 = random_peer_url();
    let peer_url_2 = random_peer_url();
    let peer_url_3 = random_peer_url();

    let mut expected_requests = Vec::new();
    op_list_1
        .clone()
        .into_iter()
        .for_each(|op_id| expected_requests.push((op_id, peer_url_1.clone())));
    op_list_2
        .clone()
        .into_iter()
        .for_each(|op_id| expected_requests.push((op_id, peer_url_2.clone())));
    op_list_3
        .clone()
        .into_iter()
        .for_each(|op_id| expected_requests.push((op_id, peer_url_3.clone())));
    let mut expected_ops = Vec::new();
    op_list_1
        .clone()
        .into_iter()
        .for_each(|op_id| expected_ops.push((op_id, peer_url_1.clone())));
    op_list_2
        .clone()
        .into_iter()
        .for_each(|op_id| expected_ops.push((op_id, peer_url_2.clone())));
    op_list_3
        .clone()
        .into_iter()
        .for_each(|op_id| expected_ops.push((op_id, peer_url_3.clone())));

    futures::future::join_all([
        fetch.request_ops(op_list_1.clone(), peer_url_1.clone()),
        fetch.request_ops(op_list_2.clone(), peer_url_2.clone()),
        fetch.request_ops(op_list_3.clone(), peer_url_3.clone()),
    ])
    .await;

    // Check that at least one request was sent for each op.
    iter_check!({
        let requests_sent = requests_sent.lock().unwrap().clone();
        if requests_sent.len() >= total_ops
            && expected_requests
                .iter()
                .all(|expected_op| requests_sent.contains(expected_op))
        {
            break;
        }
    });

    // Check that op ids are still part of ops to fetch.
    let lock = fetch.state.lock().unwrap();
    assert!(expected_ops.iter().all(|v| lock.requests.contains(v)));
}

#[tokio::test(flavor = "multi_thread")]
async fn filter_requests_for_held_ops() {
    let TestCase {
        fetch,
        requests_sent,
        op_store,
        _transport,
        ..
    } = setup_test(&CoreFetchConfig::default()).await;

    let held_op_1 = make_op(vec![1; 64]);
    let held_op_id_1 = held_op_1.compute_op_id();
    let held_op_2 = make_op(vec![2; 64]);
    let held_op_id_2 = held_op_2.compute_op_id();
    let peer_url = random_peer_url();

    op_store
        .process_incoming_ops(vec![held_op_1.into(), held_op_2.into()])
        .await
        .unwrap();

    let new_op = make_op(vec![3; 64]);
    let new_op_id = new_op.compute_op_id();
    let op_list = vec![
        held_op_id_1.clone(),
        held_op_id_2.clone(),
        new_op_id.clone(),
    ];

    fetch.request_ops(op_list, peer_url.clone()).await.unwrap();
    iter_check!(1000, 100, {
        let lock = requests_sent.lock().unwrap();
        if lock.contains(&(new_op_id.clone(), peer_url.clone())) {
            assert!(!lock.contains(&(held_op_id_1.clone(), peer_url.clone())));
            assert!(!lock.contains(&(held_op_id_2.clone(), peer_url.clone())));
            break;
        }
    });
}

#[tokio::test(flavor = "multi_thread")]
async fn unresponsive_urls_are_filtered() {
    let responsive_url = random_peer_url();
    let unresponsive_url = random_peer_url();

    let builder =
        Arc::new(default_test_builder().with_default_config().unwrap());
    let op_store = builder
        .op_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();
    // Create a mock transport to track where fetch requests are sent.
    // Set up a channel to know when the intended fetch request has been sent.
    let (expected_fetch_request_sent_tx, mut expected_fetch_request_sent_rx) =
        tokio::sync::mpsc::channel(1);
    let mut transport = MockTransport::new();
    let responsive_url_clone = responsive_url.clone();
    let unresponsive_url_clone = unresponsive_url.clone();
    transport.expect_send_module().returning({
        move |url, _, _, _| {
            if url == unresponsive_url_clone {
                panic!("fetch request sent to unresponsive URL");
            }
            if url == responsive_url_clone {
                expected_fetch_request_sent_tx.try_send(()).unwrap();
            }
            Box::pin(async move { Ok(()) })
        }
    });
    transport
        .expect_register_module_handler()
        .returning(|_, _, _| ());
    let transport = Arc::new(transport);
    let report = builder
        .report
        .create(builder.clone(), transport.clone())
        .await
        .unwrap();
    let peer_meta_store = builder
        .peer_meta_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();
    let fetch = CoreFetch::new(
        CoreFetchConfig::default(),
        TEST_SPACE_ID,
        report,
        op_store.clone(),
        peer_meta_store.clone(),
        transport.clone(),
    );

    let op_list_1 = create_op_id_list(1);
    let op_list_2 = create_op_id_list(1);
    // Set unresponsive URL as unresponsive in peer meta store.
    peer_meta_store
        .set_unresponsive(
            unresponsive_url.clone(),
            Timestamp::now(),
            Timestamp::now(),
        )
        .await
        .unwrap();

    // Request ops from unresponsive URL, which should be omitted.
    fetch
        .request_ops(op_list_1, responsive_url.clone())
        .await
        .unwrap();
    // Request ops from another URL which is not unresponsive, so that this request
    // can be awaited to have happened, which implies that the previous
    // publish to the unresponsive URL has been omitted.
    fetch
        .request_ops(op_list_2, unresponsive_url.clone())
        .await
        .unwrap();

    // Timeout after waiting 100 ms for fetch to request from the responsive URL.
    tokio::select! {
        _ = expected_fetch_request_sent_rx.recv() => {},
        _ = tokio::time::sleep(Duration::from_millis(100)) => {
            panic!("publish not sent to responsive URL")
        },
    };
}

#[tokio::test(flavor = "multi_thread")]
async fn requests_are_dropped_when_peer_url_unresponsive() {
    let config = CoreFetchConfig::default();
    let TestCase {
        fetch,
        peer_meta_store,
        _transport,
        ..
    } = setup_test(&config).await;

    let op_list_1 = create_op_id_list(2);
    let unresponsive_url = random_peer_url();

    // Create a second agent to check that their ops have not been removed.
    let op_list_2 = create_op_id_list(2);
    let responsive_url = random_peer_url();

    peer_meta_store
        .set_unresponsive(
            unresponsive_url.clone(),
            Timestamp::now(),
            Timestamp::now(),
        )
        .await
        .unwrap();

    fetch
        .request_ops(op_list_1.clone(), unresponsive_url.clone())
        .await
        .unwrap();

    // Add control agent's ops to set.
    fetch
        .request_ops(op_list_2, responsive_url.clone())
        .await
        .unwrap();

    // Wait for a request attempt to the unresponsive URL, which should remove all of its requests
    // from the set.
    iter_check!({
        if fetch
            .state
            .lock()
            .unwrap()
            .requests
            .iter()
            .all(|(_, peer_url)| *peer_url == responsive_url)
        {
            break;
        }
    });
}

#[tokio::test(flavor = "multi_thread")]
async fn empty_fetch_queue_notifies_drained() {
    let config = CoreFetchConfig::default();
    let TestCase {
        fetch, _transport, ..
    } = setup_test(&config).await;

    let (tx, mut rx) = futures::channel::oneshot::channel();
    fetch.notify_on_drained(tx);

    // The fetch queue should be immediately empty, so the notification should be sent
    // immediately.
    assert!(rx.try_recv().unwrap().is_some());
}

#[tokio::test(flavor = "multi_thread")]
async fn fetch_queue_notify_on_last_op_fetched() {
    enable_tracing();
    let config = CoreFetchConfig::default();
    let TestCase {
        fetch, _transport, ..
    } = setup_test(&config).await;

    let op_list = create_op_id_list(10);
    let peer_url = random_peer_url();

    let mut expected_requests = Vec::new();
    op_list
        .clone()
        .into_iter()
        .for_each(|op_id| expected_requests.push((op_id, peer_url.clone())));
    let mut expected_ops = Vec::new();
    op_list
        .clone()
        .into_iter()
        .for_each(|op_id| expected_ops.push((op_id, peer_url.clone())));

    fetch
        .request_ops(op_list.clone(), peer_url.clone())
        .await
        .unwrap();

    let (tx, rx) = futures::channel::oneshot::channel();
    fetch.notify_on_drained(tx);

    // Successfully fetched ops will clear the requests set.
    // That should trigger the notification.
    fetch.state.lock().unwrap().requests.clear();

    tokio::time::timeout(Duration::from_secs(1), rx)
        .await
        .expect("Timed out")
        .unwrap();
}

#[tokio::test(flavor = "multi_thread")]
async fn fetch_queue_notify_when_all_peers_unresponsive() {
    enable_tracing();
    let config = CoreFetchConfig::default();
    let builder =
        Arc::new(default_test_builder().with_default_config().unwrap());
    let op_store = builder
        .op_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();
    let peer_meta_store = builder
        .peer_meta_store
        .create(builder.clone(), TEST_SPACE_ID)
        .await
        .unwrap();

    let mut mock_transport = MockTransport::new();
    mock_transport.expect_send_module().returning({
        move |_, _, _, _| {
            Box::pin({
                async move { Err(K2Error::other("peer unresponsive")) }
            })
        }
    });
    mock_transport
        .expect_register_module_handler()
        .returning(|_, _, _| ());
    let mock_transport = Arc::new(mock_transport);
    let report = builder
        .report
        .create(builder.clone(), mock_transport.clone())
        .await
        .unwrap();

    let fetch = CoreFetch::new(
        config.clone(),
        TEST_SPACE_ID,
        report,
        op_store.clone(),
        peer_meta_store.clone(),
        mock_transport.clone(),
    );

    let op_list = create_op_id_list(10);
    let peer_url = random_peer_url();

    let mut expected_requests = Vec::new();
    op_list
        .clone()
        .into_iter()
        .for_each(|op_id| expected_requests.push((op_id, peer_url.clone())));
    let mut expected_ops = Vec::new();
    op_list
        .clone()
        .into_iter()
        .for_each(|op_id| expected_ops.push((op_id, peer_url.clone())));

    fetch
        .request_ops(op_list.clone(), peer_url.clone())
        .await
        .unwrap();

    let (tx, rx) = futures::channel::oneshot::channel();
    fetch.notify_on_drained(tx);

    tokio::time::timeout(Duration::from_secs(1), rx)
        .await
        .expect("Timed out")
        .unwrap();
}
