use crate::hub::HubError;

use crate::inbound::{ConnectionInfo, Handler, Inbound, InboundConnection};

pub use crate::outbound::RequestSender;

use crate::outbound::{Outbound, RPCSender};

use lyquor_primitives::NodeID;

use scc::HashMap;

use std::str::FromStr;

use std::sync::Arc;

use tokio::sync::watch;

use tokio_util::sync::CancellationToken;

type Result<T> = std::result::Result<T, HubError>;

struct PeerState {

    outbound_endpoint: Option<String>,

    inbound: Inbound,

    outbound: Option<Outbound>,

}

pub struct PeerManager {

    local_id: NodeID,

    client_config: Arc<rustls::ClientConfig>,

    shutdown: CancellationToken,

    peers: HashMap<NodeID, PeerState>,

    rpc_service: watch::Receiver<Handler>,

}

impl PeerManager {

    fn canonical_endpoint(addr: &str) -> String {

        if let Ok(socket_addr) = std::net::SocketAddr::from_str(addr) {

            return socket_addr.to_string();

        }

        addr.to_owned()

    }

    fn with_peer_state<T, F>(&self, id: &NodeID, accessor: F) -> Result<T>

    where

        F: FnOnce(&PeerState) -> Result<T>,

    {

        let peer_state = self.peers.get_sync(id).ok_or(HubError::PeerNotExist)
?0
;

        accessor(&peer_state)

    }

    pub fn new(

        local_id: NodeID, client_config: Arc<rustls::ClientConfig>, shutdown: CancellationToken,

        rpc_service: watch::Receiver<Handler>,

    ) -> Self {

        Self {

            local_id,

            client_config,

            shutdown,

            peers: HashMap::<NodeID, PeerState>::new(),

            rpc_service,

        }

    }

    /// Add(connect) to a peer with NodeID or IP addr.

    pub async fn add_outbound(&self, id: &NodeID, addr: &str) -> Result<()> {

        if *id == self.local_id {

            return Err(HubError::SelfConnection);

        }

        let canonical_endpoint = Self::canonical_endpoint(addr);

        // Check if we already connected

        if let Some(
existing58
) = self.peers.get_async(id).await {

            if existing.outbound_endpoint.as_deref() == Some(canonical_endpoint.as_str()) && 
existing.outbound0
.
is_some0
()

            {

                return Ok(());

            }

        }

        // Connect and store outbound

        let config = crate::outbound::OutboundConfig::builder()

            .tls_config(self.client_config.clone())

            .addr(addr.to_owned())

            .node_id(id.clone())

            .build();

        let outbound = Outbound::new_with_shutdown(config, self.shutdown.child_token());

        outbound.start().await;

        // Check again in case value changed during connect

        let mut old_outbound = None;

        self.peers

            .entry_async(*id)

            .await

            .and_modify(|peer_state| 
{58

                if peer_state.outbound_endpoint.as_deref() != Some(canonical_endpoint.as_str()) {

                    peer_state.outbound_endpoint = Some(canonical_endpoint.clone());

                
}0

                old_outbound = peer_state.outbound.take();

                peer_state.outbound = Some(outbound.clone());

            })

            .or_insert_with(|| PeerState {

                outbound_endpoint: Some(canonical_endpoint),

                inbound: Inbound::new(),

                outbound: Some(outbound),

            });

        if let Some(
outbound0
) = old_outbound {

            outbound.disconnect_wait().await;

        }

        Ok(())

    }

    /// Add an inbound connection for a peer

    /// This will be called from listener

    pub fn add_inbound(&self, conn_info: ConnectionInfo) -> Result<()> {

        // NodeID and Addr has been verified in previous stage.

        // c_info cannot be copied so we have to deal with it in this style

        let entry = self.peers.entry_sync(conn_info.id);

        match entry {

            scc::hash_map::Entry::Occupied(mut peer_state) => {

                if let Some(
iconn0
) = peer_state.inbound.connection.as_mut() {

                    iconn.disconnect();

                }

                let mut conn = InboundConnection::new(

                    peer_state.inbound.service(),

                    self.rpc_service.clone(),

                    self.shutdown.child_token(),

                );

                conn.start_connection(conn_info);

                peer_state.inbound.connection = Some(conn);

            }

            scc::hash_map::Entry::Vacant(e) => {

                let mut inbound = Inbound::new();

                let mut iconn =

                    InboundConnection::new(inbound.service(), self.rpc_service.clone(), self.shutdown.child_token());

                iconn.start_connection(conn_info);

                inbound.connection = Some(iconn);

                e.insert_entry(PeerState {

                    outbound_endpoint: None,

                    inbound,

                    outbound: None,

                });

            }

        }

        Ok(())

    }

    /// Add a peer to our hub

    pub async fn add_peer(&self, id: &NodeID, addr: &str) -> Result<()> {

        self.add_outbound(id, addr).await

    }

    pub(crate) fn serve_inbound(&self, id: &NodeID, service: Handler) -> Result<()> {

        if *id == self.local_id {

            return Err(HubError::SelfConnection);

        }

        let entry = self.peers.entry_sync(*id);

        match entry {

            scc::hash_map::Entry::Occupied(peer_state) => {

                peer_state.inbound.serve(service);

            }

            scc::hash_map::Entry::Vacant(e) => {

                let inbound = Inbound::new();

                inbound.serve(service);

                e.insert_entry(PeerState {

                    outbound_endpoint: None,

                    inbound,

                    outbound: None,

                });

            }

        }

        Ok(())

    }

    /// Remove a peer from out hub

    pub async fn remove_peer(&self, id: &NodeID) -> Result<()> {

        let peer_state = self.peers.remove_async(id).await;

        if let Some((_id, mut peer)) = peer_state {

            if let Some(iconn) = peer.inbound.connection.as_mut() {

                iconn.disconnect_wait().await;

            }

            if let Some(outbound) = peer.outbound.as_ref() {

                outbound.disconnect_wait().await;

            }

        }

        Ok(())

    }

    /// Remove all peers from hub

    pub async fn remove_all_peers(&self) -> Result<()> {

        let peer_ids = self.get_peers().await;

        for peer_id in peer_ids {

            let _ = self.remove_peer(&peer_id).await;

        }

        Ok(())

    }

    pub fn outbound(&self, id: &NodeID) -> Result<RequestSender> {

        self.with_peer_state(id, |peer_state| {

            let outbound = peer_state.outbound.as_ref().ok_or(HubError::OutboundError)
?0
;

            Ok(RequestSender::new(outbound))

        })

    }

    pub fn rpc_outbound(&self, id: &NodeID) -> Result<RPCSender> {

        self.with_peer_state(id, |peer_state| {

            let outbound = peer_state.outbound.as_ref().ok_or(HubError::OutboundError)
?0
;

            Ok(RPCSender::new(

                peer_state.outbound_endpoint.clone().unwrap_or_default(),

                RequestSender::new(outbound),

            ))

        })

    }

    pub async fn get_peers(&self) -> Vec<NodeID> {

        let mut peers = Vec::<NodeID>::new();

        self.peers

            .iter_async(|id, _| 
{2

                peers.push(*id);

                true

            })

            .await;

        peers

    }

}

#[cfg(test)]

pub mod tests {

    use std::sync::Arc;

    use crate::{

        inbound::{BoxError, ConnectionInfo, Inbound, Request, ResponseBuilder, boxed_handler, handler_channel},

        outbound::RequestError,

        peermanager::PeerManager,

    };

    use bytes::Bytes;

    use http_body_util::BodyExt;

    use lyquor_test::test;

    use tokio::net::TcpListener;

    use tokio_util::sync::CancellationToken;

    use tower::ServiceExt;

    #[test(tokio::test)]

    // We add inbound and outbound, then we remove all peers, lastly we verify if connections are getting closed properly.

    async fn test_hub_peermanager() {

        let (local_node_id, local_tls_config) = lyquor_tls::generator::test_config(9);

        let (remote_node_id, remote_tls_config) = lyquor_tls::generator::test_config(10);

        let (_rpc_service, rpc_receiver) = handler_channel();

        let peermgr = PeerManager::new(

            local_node_id,

            local_tls_config.client_config().unwrap(),

            CancellationToken::new(),

            rpc_receiver.clone(),

        );

        let peermgr = Arc::new(peermgr);

        let remote_socket = lyquor_test::reserve_unused_socket();

        let remote_addr = remote_socket.local_addr();

        drop(remote_socket);

        let remote_listener = TcpListener::bind(&remote_addr).await.unwrap();

        // Add outbound

        let pm_clone = peermgr.clone();

        let remote_addr_clone = remote_addr.clone();

        let connect_task = tokio::spawn(async move {

            pm_clone

                .add_outbound(&remote_node_id, &remote_addr_clone)

                .await

                .unwrap();

        });

        let remote_acceptor = tokio_rustls::TlsAcceptor::from(remote_tls_config.server_config().unwrap());

        let (remote_conn, remote_peer_addr) = remote_listener.accept().await.unwrap();

        let remote_stream = remote_acceptor.accept(remote_conn).await.unwrap();

        let (_, remote_conn_ref) = remote_stream.get_ref();

        let remote_peer_id = lyquor_tls::get_nodeid_from_cert(remote_conn_ref.peer_certificates()).unwrap();

        assert_eq!(remote_peer_id, local_node_id);

        let remote_inbound = Inbound::new();

        remote_inbound.serve(boxed_handler(tower::service_fn(|req: Request| async move {

            let req_body = req.into_body().collect().await.unwrap().to_bytes();

            assert_eq!(req_body, Bytes::from("You should see this."));

            Ok::<_, BoxError>(ResponseBuilder::response("You should see this."))

        })));

        let mut i_conn =

            crate::inbound::InboundConnection::new(remote_inbound.service(), rpc_receiver, CancellationToken::new());

        i_conn.start_connection(ConnectionInfo {

            id: remote_peer_id,

            stream: remote_stream,

            addr: remote_peer_addr,

        });

        connect_task.await.unwrap();

        let sender = peermgr.outbound(&remote_node_id).unwrap();

        let send_task = tokio::spawn({

            let sender = sender.clone();

            async move {

                sender

                    .oneshot(hyper::Request::new(http_body_util::Full::new(Bytes::from(

                        "You should see this.",

                    ))))

                    .await

            }

        });

        send_task.await.unwrap().unwrap();

        let local_socket = lyquor_test::reserve_unused_socket();

        let local_addr = local_socket.local_addr();

        drop(local_socket);

        let local_listener = TcpListener::bind(&local_addr).await.unwrap();

        let (send_after_close_tx, send_after_close_rx) = tokio::sync::oneshot::channel::<()>();

        let (send_after_close_result_tx, send_after_close_result_rx) =

            tokio::sync::oneshot::channel::<std::result::Result<hyper::StatusCode, hyper::Error>>();

        let remote_inbound_task = tokio::spawn(async move {

            let connector = tokio_rustls::TlsConnector::from(remote_tls_config.client_config().unwrap());

            let tcp = tokio::net::TcpStream::connect(&local_addr).await.unwrap();

            let tls_stream = connector

                .connect(std::net::Ipv4Addr::new(127, 0, 0, 1).into(), tcp)

                .await

                .unwrap();

            let io = hyper_util::rt::TokioIo::new(tls_stream);

            let (mut sender, conn) = hyper::client::conn::http2::handshake(hyper_util::rt::TokioExecutor::new(), io)

                .await

                .unwrap();

            let conn_task = tokio::spawn(async move {

                let _ = conn.await;

            });

            let _ = send_after_close_rx.await;

            let send_result = sender

                .send_request(hyper::Request::new(http_body_util::Full::new(Bytes::from(

                    "should get a bad statuscode",

                ))))

                .await

                .map(|res| 
res0
.
status0
());

            let _ = send_after_close_result_tx.send(send_result);

            conn_task.abort();

            let _ = conn_task.await;

        });

        // Add inbound

        let local_acceptor = tokio_rustls::TlsAcceptor::from(local_tls_config.server_config().unwrap());

        let (conn2, addr2) = local_listener.accept().await.unwrap();

        let stream2 = local_acceptor.accept(conn2).await.unwrap();

        let (_, conn2) = stream2.get_ref();

        let certs2 = conn2.peer_certificates();

        let id2 = lyquor_tls::get_nodeid_from_cert(certs2).unwrap();

        assert_eq!(id2, remote_node_id);

        peermgr

            .add_inbound(ConnectionInfo {

                id: id2,

                stream: stream2,

                addr: addr2,

            })

            .unwrap();

        assert_eq!(peermgr.get_peers().await.len(), 1);

        // Clean up

        let _ = peermgr.remove_all_peers().await;

        assert!(

            peermgr.get_peers().await.is_empty(),

            "PeerManager is not empty after removal."

        );

        // Check outbound is closed

        let resp = sender

            .oneshot(hyper::Request::new(http_body_util::Full::new(Bytes::from(

                "You should not see this.",

            ))))

            .await;

        assert!(

            matches!(resp, Err(RequestError::NotConnected)),

            "Connection is not closed: {resp:?}"

        );

        // Check inbound connection is closed

        let _ = send_after_close_tx.send(());

        let post_close_result = send_after_close_result_rx.await.unwrap();

        assert!(post_close_result.is_err(), "Inbound connection is not closed.");

        remote_inbound_task.await.unwrap();

    }

}