#![allow(dead_code)]

use crate::p2p::notifications::{
    handler::{
        NotificationsHandler, NotificationsHandlerFromBehavior, 
        NotificationsHandlerToBehavior,
    },
    messages::BlockHash,
    messages::ProtocolRole,
};

use bytes::BytesMut;
use futures::channel::mpsc;
use libp2p::{
    core::{ConnectedPoint, Endpoint},
    swarm::{
        derive_prelude::ConnectionEstablished, ConnectionClosed, ConnectionDenied, 
        ConnectionId, NetworkBehaviour, NotifyHandler, ToSwarm,
    },
    Multiaddr, PeerId,
};
use std::{
    collections::{HashMap, HashSet, VecDeque},
    task::{Poll, Waker},
};

const LOG_TARGET: &str = "ghost-wlaker-behavior";

#[derive(Debug)]
pub enum NotificationsToSwarm {
    CustomProtocolOpen {
        peer_id: PeerId,
        index: usize,
        received_handshake: Vec<u8>,
        inbound: bool,
        sender: mpsc::Sender<Vec<u8>>,
    },

    CustomProtocolClosed {
        peer_id: PeerId,
        index: usize,
    },

    Notification {
        peer_id: PeerId,
        index: usize,
        message: BytesMut,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProtocolsData {
    pub genesis_hash: BlockHash,
    pub node_role: ProtocolRole,
}

pub struct Notifications {
    events: VecDeque<ToSwarm<NotificationsToSwarm, NotificationsHandlerFromBehavior>>,
    peers_details: HashMap<PeerId, HashSet<ConnectionId>>,
    data: ProtocolsData,
    waker: Option<Waker>,
}

impl Notifications {
    pub fn new(data: ProtocolsData) -> Self {
        Notifications {
            events: VecDeque::with_capacity(16),
            peers_details: HashMap::default(),
            data,
            waker: None,
        }
    }

    fn propagate_event(
        &mut self,
        event: ToSwarm<NotificationsToSwarm, NotificationsHandlerFromBehavior>,
    ) {
        if let Some(waker) = self.waker.take() {
            waker.wake();
        }

        self.events.push_back(event);
    }
}

impl NetworkBehaviour for Notifications {
    type ConnectionHandler = NotificationsHandler;
    type ToSwarm = NotificationsToSwarm;

    fn handle_pending_inbound_connection(
        &mut self,
        _connection_id: ConnectionId,
        _local_addr: &Multiaddr,
        _remote_addr: &Multiaddr,
    ) -> Result<(), ConnectionDenied> {
        Ok(())
    }

    fn handle_pending_outbound_connection(
        &mut self,
        _connection_id: ConnectionId,
        _maybe_peer: Option<PeerId>,
        _addresses: &[Multiaddr],
        _effective_role: Endpoint,
    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
        Ok(Vec::new())
    }

    fn handle_established_inbound_connection(
        &mut self,
        _connection_id: libp2p::swarm::ConnectionId,
        peer: libp2p::PeerId,
        local_addr: &libp2p::Multiaddr,
        remote_addr: &libp2p::Multiaddr,
    ) -> Result<libp2p::swarm::THandler<Self>, libp2p::swarm::ConnectionDenied> {
        tracing::debug!(target: LOG_TARGET, "Notifications new inbound for peer={:?}", peer);

        let handler = NotificationsHandler::new(
            peer,
            ConnectedPoint::Listener {
                local_addr: local_addr.clone(),
                send_back_addr: remote_addr.clone(),
            },
            self.data.clone(),
        );

        Ok(handler)
    }

    fn handle_established_outbound_connection(
        &mut self,
        _connection_id: libp2p::swarm::ConnectionId,
        peer: libp2p::PeerId,
        addr: &libp2p::Multiaddr,
        _role_override: libp2p::core::Endpoint,
    ) -> Result<libp2p::swarm::THandler<Self>, libp2p::swarm::ConnectionDenied> {
        tracing::debug!(target: LOG_TARGET, "Notifications new outbound for peer={:?}", peer);

        let handler = NotificationsHandler::new(
            peer,
            ConnectedPoint::Dialer {
                role_override: Endpoint::Dialer,
                address: addr.clone(),
            },
            self.data.clone(),
        );

        Ok(handler)
    }

    fn on_swarm_event(&mut self, event: libp2p::swarm::FromSwarm<Self::ConnectionHandler>) {
        match event {
            libp2p::swarm::FromSwarm::ConnectionEstablished(ConnectionEstablished {
                peer_id,
                connection_id,
                ..
            }) => {
                tracing::debug!(target: LOG_TARGET,
                    "Notifications swarm connection established peer={:?} connection={:?}",
                    peer_id,
                    connection_id
                );

                self.peers_details
                    .entry(peer_id)
                    .and_modify(|entry| {
                        let _ = entry.insert(connection_id);
                    })
                .or_insert_with(|| {
                    let mut hash = HashSet::new();
                    hash.insert(connection_id);
                    hash
                });

                for index in 0..2 {
                    self.propagate_event(ToSwarm::NotifyHandler {
                        peer_id,
                        handler: NotifyHandler::One(connection_id),
                        event: NotificationsHandlerFromBehavior::Open { index },
                    });
                }
            }
            libp2p::swarm::FromSwarm::ConnectionClosed(ConnectionClosed {
                peer_id,
                connection_id,
                ..
            }) => {
                tracing::debug!(target: LOG_TARGET,
                    "Notifications swarm connection closed peer={:?} connection={:?}",
                    peer_id,
                    connection_id
                );

                if let Some(details) = self.peers_details.get_mut(&peer_id) {
                    let removed = details.remove(&connection_id);
                    if !removed {
                        tracing::trace!(target: LOG_TARGET,
                            "Notifications swarm connection closed for untracked connection peer={:?} connection={:?}",
                            peer_id,
                            connection_id
                        );
                    }
                } else {
                    tracing::trace!(target: LOG_TARGET,
                        "Notifications swarm connection closed for untracked peer, peer={:?} connection={:?}",
                        peer_id,
                        connection_id
                    );
                }

                for index in 0..2 {
                    self.propagate_event(ToSwarm::NotifyHandler {
                        peer_id,
                        handler: NotifyHandler::One(connection_id),
                        event: NotificationsHandlerFromBehavior::Close { index },
                    });
                }
            }
            _ => (),
        }
    }

    fn on_connection_handler_event(
        &mut self,
        peer_id: libp2p::PeerId,
        connection_id: libp2p::swarm::ConnectionId,
        event: libp2p::swarm::THandlerOutEvent<Self>,
    ) {
        tracing::debug!(target: LOG_TARGET,
            "Notifications new substream for peer {:?} {:?}",
            peer_id,
            event
        );

        match event {
            NotificationsHandlerToBehavior::HandshakeCompleted {
                index,
                handshake,
                is_inbound,
                sender,
                ..
            } => {
                tracing::trace!(target: LOG_TARGET,
                    "Notifications handler complited handshake peer={:?} connection={:?} index={:?} handshake={:?}",
                    peer_id,
                    connection_id,
                    index,
                    handshake,
                );

                self.propagate_event(ToSwarm::GenerateEvent(
                        NotificationsToSwarm::CustomProtocolOpen {
                            index,
                            peer_id,
                            received_handshake: handshake,
                            inbound: is_inbound,
                            sender,
                        },
                ));
            }
            NotificationsHandlerToBehavior::HandshakeError { index } => {
                tracing::trace!(target: LOG_TARGET,
                    "Notifications handler error handshake peer={:?} connection={:?} index={:?}",
                    peer_id,
                    connection_id,
                    index,
                );
            }
            NotificationsHandlerToBehavior::OpenDesiredByRemote { index } => {
                // Note: extend to reject protocols for specific peers in the future.
                self.propagate_event(ToSwarm::NotifyHandler {
                    peer_id,
                    handler: NotifyHandler::One(connection_id),
                    event: NotificationsHandlerFromBehavior::Open { index },
                });
            }
            NotificationsHandlerToBehavior::CloseDesired { index } => {
                self.propagate_event(ToSwarm::NotifyHandler {
                    peer_id,
                    handler: NotifyHandler::One(connection_id),
                    event: NotificationsHandlerFromBehavior::Close { index },
                });
            }
            NotificationsHandlerToBehavior::Close { .. } => {}
            NotificationsHandlerToBehavior::Notification { bytes, index } => {
                self.propagate_event(ToSwarm::GenerateEvent(NotificationsToSwarm::Notification {
                    peer_id,
                    index,
                    message: bytes,
                }));
            }
        }
    }

    fn poll(
        &mut self,
        cx: &mut std::task::Context<'_>,
        _params: &mut impl libp2p::swarm::PollParameters,
    ) -> std::task::Poll<ToSwarm<Self::ToSwarm, libp2p::swarm::THandlerInEvent<Self>>> {
        self.waker = Some(cx.waker().clone());

        if let Some(event) = self.events.pop_front() {
            return Poll::Ready(event);
        }

        Poll::Pending
    }
}
