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

use crate::outbound::Outbound;

pub use crate::outbound::RequestSender;

use crate::{hub::HubError, inbound::RequestReceiver};

use lyquor_primitives::NodeID;

use scc::HashMap;

use std::str::FromStr;

use std::sync::Arc;

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>,

    runtime: tokio::runtime::Handle,

    shutdown: CancellationToken,

    peers: HashMap<NodeID, PeerState>,

}

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>, runtime: tokio::runtime::Handle,

        shutdown: CancellationToken,

    ) -> Self {

        Self {

            local_id,

            client_config,

            runtime,

            shutdown,

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

        }

    }

    /// 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(
existing0
) = self.peers.get_async(id).await {

            if existing.outbound_endpoint.as_deref() == Some(canonical_endpoint.as_str()) && existing.outbound.is_some()

            {

                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(self.runtime.clone()).await;

        // Check again in case value changed during connect

        self.peers

            .entry_async(*id)

            .await

            .and_modify(|peer_state| 
{0

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

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

                }

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

            })

            .or_insert_with(|| PeerState {

                outbound_endpoint: Some(canonical_endpoint),

                inbound: Inbound::new(),

                outbound: Some(outbound),

            });

        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.

        let runtime = conn_info.runtime.clone();

        // 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.restart_connection(conn_info, &runtime);

                } else {

                    let mut conn = InboundConnection::new(peer_state.inbound.tx(), 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 tx = inbound.tx();

                let mut iconn = InboundConnection::new(tx, 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

    }

    /// Remove a peer from out hub

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

        let runtime = self.runtime.clone();

        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(&runtime);

            }

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

                outbound.shutdown();

            }

        }

        Ok(())

    }

    /// Remove all peers from hub

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

        let runtime = self.runtime.clone();

        self.peers

            .retain_async(|_, peer_state| 
{0

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

                    inbound_connection.disconnect(&runtime);

                }

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

                    outbound.shutdown();

                }

                false

            })

            .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(self.local_id, outbound))

        })

    }

    pub fn inbound(&self, id: &NodeID) -> Result<RequestReceiver> {

        self.with_peer_state(id, |peer_state| Ok(peer_state.inbound.rx()))

    }

}