//! Relayers
//!
//! This is the main entry point to the client library.
use crate::{client::ActiveSubscription, Client, Error};
use futures::future::join_all;
use nostr_rs_types::{
    client::{self, subscribe},
    types::SubscriptionId,
    Response,
};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::{mpsc, RwLock};
use url::Url;

type Subscriptions =
    Arc<RwLock<HashMap<SubscriptionId, (subscribe::Subscribe, Vec<ActiveSubscription>)>>>;
/// Clients
///
/// This is a set of outgoing connections to relayers. This struct can connect
/// async to N relayers offering a simple API to talk to all of them at the same
/// time, and to receive messages
#[derive(Debug)]
pub struct Pool {
    clients: RwLock<HashMap<Url, Client>>,
    sender: mpsc::Sender<(Response, Url)>,
    receiver: Option<mpsc::Receiver<(Response, Url)>>,
    subscriptions: Subscriptions,
}

/// Default channel buffer size for the pool
pub const DEFAULT_CHANNEL_BUFFER_SIZE: usize = 1_000;

impl Default for Pool {
    fn default() -> Self {
        let (sender, receiver) = mpsc::channel(DEFAULT_CHANNEL_BUFFER_SIZE);
        Self {
            clients: Default::default(),
            receiver: Some(receiver),
            subscriptions: Default::default(),
            sender,
        }
    }
}

/// Return a subscription that will be removed when dropped
#[derive(Debug)]
pub struct PoolSubscription {
    subscription_id: SubscriptionId,
    subscriptions: Subscriptions,
}

impl Drop for PoolSubscription {
    fn drop(&mut self) {
        let subscriptions = self.subscriptions.clone();
        let subscription_id = self.subscription_id.clone();
        tokio::spawn(async move {
            subscriptions.write().await.remove(&subscription_id);
        });
    }
}

impl Pool {
    /// Creates a new instance with a list of urls
    pub fn new_with_clients(clients: Vec<Url>) -> Self {
        let (sender, receiver) = mpsc::channel(DEFAULT_CHANNEL_BUFFER_SIZE);
        let clients = clients
            .into_iter()
            .map(|url| (url.clone(), Client::new(sender.clone(), url)))
            .collect::<HashMap<_, _>>();

        Self {
            clients: RwLock::new(clients),
            subscriptions: Default::default(),
            receiver: Some(receiver),
            sender,
        }
    }

    /// Splits the pool removing the receiver to be used in a different context
    pub fn split(mut self) -> Result<(mpsc::Receiver<(Response, Url)>, Self), Error> {
        Ok((self.receiver.take().ok_or(Error::AlreadySplitted)?, self))
    }

    /// Tries to receive a message from any of the connected relayers
    pub fn try_recv(&mut self) -> Option<(Response, Url)> {
        self.receiver.as_mut()?.try_recv().ok()
    }

    /// Receives a message from any of the connected relayers
    pub async fn recv(&mut self) -> Option<(Response, Url)> {
        self.receiver.as_mut()?.recv().await
    }

    /// Returns the number of active subscriptions
    pub async fn active_subscriptions(&self) -> usize {
        self.subscriptions.read().await.keys().len()
    }

    /// Subscribes to all the connected relayers
    pub async fn subscribe(
        &self,
        subscription: subscribe::Subscribe,
    ) -> Result<PoolSubscription, Error> {
        let clients = self.clients.read().await;

        let wait_all = clients
            .values()
            .map(|sender| sender.subscribe(subscription.clone()))
            .collect::<Vec<_>>();

        let subscription_id = subscription.subscription_id.clone();

        self.subscriptions.write().await.insert(
            subscription_id.clone(),
            (
                subscription,
                join_all(wait_all)
                    .await
                    .into_iter()
                    .collect::<Result<Vec<_>, _>>()?,
            ),
        );

        Ok(PoolSubscription {
            subscription_id,
            subscriptions: self.subscriptions.clone(),
        })
    }

    /// Sends a request to all the connected relayers
    pub async fn post(&self, request: client::Event) {
        let clients = self.clients.read().await;
        join_all(
            clients
                .values()
                .map(|sender| sender.post(request.clone()))
                .collect::<Vec<_>>(),
        )
        .await;
    }

    /// Returns the number of active connections.
    pub async fn check_active_connections(&self) -> usize {
        self.clients
            .read()
            .await
            .iter()
            .filter(|(_, client)| client.is_connected())
            .collect::<Vec<_>>()
            .len()
    }

    /// Creates a connection to a new relayer.
    ///
    /// This function will open a connection at most once, if a connection
    /// already exists false will be returned
    pub async fn connect_to(&self, url: Url) {
        let mut clients = self.clients.write().await;
        let mut subscriptions = self.subscriptions.write().await;

        if !clients.contains_key(&url) {
            log::warn!("Connecting to {}", url);
            let client = Client::new(self.sender.clone(), url.clone());

            for (filter, sub) in subscriptions.values_mut() {
                let _ = client.subscribe(filter.clone()).await.map(|subscription| {
                    sub.push(subscription);
                });
            }

            clients.insert(url.clone(), client);
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use nostr_rs_memory::Memory;
    use nostr_rs_relayer::Relayer;
    use nostr_rs_types::{account::Account, types::Content};
    use std::time::Duration;
    use tokio::{net::TcpListener, task::JoinHandle, time::sleep};

    async fn dummy_server(port: u16) -> (Url, JoinHandle<()>) {
        let listener = TcpListener::bind(format!("127.0.0.1:{}", port))
            .await
            .unwrap();
        let local_addr = listener.local_addr().expect("addr");

        let relayer = Relayer::new(Some(Memory::default()), None).expect("valid dummy server");
        let stopper = relayer.main(listener).expect("valid main loop");
        (
            Url::parse(&format!("ws://{}", local_addr)).expect("valid url"),
            stopper,
        )
    }

    #[tokio::test]
    async fn droppable_subscription() {
        let client_pool = Pool::default();
        let subscription = client_pool
            .subscribe(Default::default())
            .await
            .expect("valid subscription");

        assert_eq!(client_pool.active_subscriptions().await, 1);
        drop(subscription);
        sleep(Duration::from_millis(10)).await;
        assert_eq!(client_pool.active_subscriptions().await, 0);
    }

    #[tokio::test]
    async fn connect_to_dummy_server() {
        let (addr, stopper) = dummy_server(0).await;
        let client_pool = Pool::new_with_clients(vec![addr]);

        assert_eq!(0, client_pool.check_active_connections().await);

        sleep(Duration::from_millis(10)).await;
        assert_eq!(1, client_pool.check_active_connections().await);

        // stop dummy server
        stopper.abort();

        sleep(Duration::from_millis(100)).await;
        assert_eq!(0, client_pool.check_active_connections().await);
    }

    #[tokio::test]
    async fn two_clients_communication() {
        let (addr, _) = dummy_server(0).await;
        let mut client_pool1 = Pool::new_with_clients(vec![addr.clone()]);
        let client_pool2 = Pool::new_with_clients(vec![addr]);

        let _sub1 = client_pool1
            .subscribe(Default::default())
            .await
            .expect("valid subscription");

        sleep(Duration::from_millis(10)).await;

        assert!(client_pool1
            .try_recv()
            .map(|(r, _)| r)
            .expect("valid message")
            .as_end_of_stored_events()
            .is_some());
        assert!(client_pool1.try_recv().is_none());

        let account1 = Account::default();
        let signed_content = account1
            .sign_content(vec![], Content::ShortTextNote("test 0".to_owned()), None)
            .expect("valid signed content");

        client_pool2.post(signed_content.clone().into()).await;

        sleep(Duration::from_millis(10)).await;

        assert_eq!(
            Some((signed_content.id, signed_content.signature)),
            client_pool1
                .try_recv()
                .and_then(|(r, _)| r.as_event().cloned().map(|x| x.event))
                .map(|x| (x.id, x.signature))
        );
        assert!(client_pool1.try_recv().is_none());
    }

    #[tokio::test]
    async fn reconnect_and_resubscribe() {
        let (addr, stopper) = dummy_server(0).await;
        let mut client_pool1 = Pool::new_with_clients(vec![addr.clone()]);
        let client_pool2 = Pool::new_with_clients(vec![addr.clone()]);

        let _sub1 = client_pool1
            .subscribe(Default::default())
            .await
            .expect("valid subscription");

        sleep(Duration::from_millis(10)).await;

        assert!(client_pool1
            .try_recv()
            .map(|(r, _)| r)
            .expect("valid message")
            .as_end_of_stored_events()
            .is_some());
        assert!(client_pool1.try_recv().is_none());

        let account1 = Account::default();
        let signed_content = account1
            .sign_content(vec![], Content::ShortTextNote("test 1".to_owned()), None)
            .expect("valid signed content");

        client_pool2.post(signed_content.clone().into()).await;

        sleep(Duration::from_millis(10)).await;

        assert_eq!(
            Some((signed_content.id, signed_content.signature)),
            client_pool1
                .try_recv()
                .and_then(|(r, _)| r.as_event().cloned().map(|x| x.event))
                .map(|x| (x.id, x.signature))
        );
        assert!(client_pool1.try_recv().is_none());

        assert_eq!(1, client_pool1.check_active_connections().await);
        assert_eq!(1, client_pool2.check_active_connections().await);

        stopper.abort();

        sleep(Duration::from_millis(10)).await;

        assert_eq!(0, client_pool1.check_active_connections().await);
        assert_eq!(0, client_pool2.check_active_connections().await);

        let (_, stopper) = dummy_server(addr.port().expect("port")).await;

        sleep(Duration::from_millis(2_000)).await;

        assert_eq!(1, client_pool1.check_active_connections().await);
        assert_eq!(1, client_pool2.check_active_connections().await);

        let signed_content = account1
            .sign_content(vec![], Content::ShortTextNote("test 1".to_owned()), None)
            .expect("valid signed content");

        client_pool2.post(signed_content.clone().into()).await;

        sleep(Duration::from_millis(10)).await;

        assert!(client_pool1
            .try_recv()
            .map(|(r, _)| r)
            .expect("valid message")
            .as_end_of_stored_events()
            .is_some());
        assert_eq!(
            Some((signed_content.id, signed_content.signature)),
            client_pool1
                .try_recv()
                .and_then(|(r, _)| r.as_event().cloned().map(|x| x.event))
                .map(|x| (x.id, x.signature))
        );
        assert!(client_pool1.try_recv().is_none());

        stopper.abort();
    }

    #[tokio::test]
    async fn connect_multiple_servers() {
        let (addr1, _) = dummy_server(0).await;
        let (addr2, _) = dummy_server(0).await;
        let mut client_pool1 = Pool::new_with_clients(vec![addr1.clone(), addr2]);
        let client_pool2 = Pool::new_with_clients(vec![addr1]);

        let _sub1 = client_pool1
            .subscribe(Default::default())
            .await
            .expect("valid subscription");

        sleep(Duration::from_millis(10)).await;

        assert!(client_pool1
            .try_recv()
            .map(|(r, _)| r)
            .expect("valid message")
            .as_end_of_stored_events()
            .is_some());

        assert!(client_pool1
            .try_recv()
            .map(|(r, _)| r)
            .expect("valid message")
            .as_end_of_stored_events()
            .is_some());

        let account1 = Account::default();
        let signed_content = account1
            .sign_content(vec![], Content::ShortTextNote("test 0".to_owned()), None)
            .expect("valid signed content");

        client_pool2.post(signed_content.clone().into()).await;

        sleep(Duration::from_millis(10)).await;

        assert_eq!(
            Some((signed_content.id, signed_content.signature)),
            client_pool1
                .try_recv()
                .and_then(|(r, _)| r.as_event().cloned().map(|x| x.event))
                .map(|x| (x.id, x.signature))
        );
        assert!(client_pool1.try_recv().is_none());
    }
}
