use futures::StreamExt;
use libp2p::swarm::{Swarm, SwarmEvent};
use libp2p_swarm_derive::NetworkBehaviour;
use libp2p::identity::Keypair;
use libp2p::ping;
use libp2p::identify;
use libp2p::mdns;
use libp2p::tcp;
use libp2p::noise;
use libp2p::yamux;
use libp2p::PeerId;
use crate::peer::Peer;
use crate::peer_store::{shared_peer_store, SharedPeerStore};
use std::time::Duration;

#[derive(NetworkBehaviour)]
#[behaviour(out_event = "CommsEvent", event_processing = false)]
pub struct CommsBehaviour {
    pub ping: ping::Behaviour,
    pub identify: identify::Behaviour,
    pub mdns: mdns::tokio::Behaviour,
}

#[derive(Debug)]
pub enum CommsEvent {
    Ping(ping::Event),
    Identify(identify::Event),
    Mdns(mdns::Event),
}

impl From<ping::Event> for CommsEvent {
    fn from(event: ping::Event) -> Self {
        CommsEvent::Ping(event)
    }
}

impl From<identify::Event> for CommsEvent {
    fn from(event: identify::Event) -> Self {
        CommsEvent::Identify(event)
    }
}

impl From<mdns::Event> for CommsEvent {
    fn from(event: mdns::Event) -> Self {
        CommsEvent::Mdns(event)
    }
}

pub struct Comms {
    peer_store: SharedPeerStore,
    swarm: Option<Swarm<CommsBehaviour>>,
    local_peer_id: PeerId,
}

impl Comms {
    pub async fn new() -> anyhow::Result<Self> {
        let keypair = Keypair::generate_ed25519();
        let local_peer_id = PeerId::from(keypair.public());
        let peer_store = shared_peer_store();

        let behaviour = CommsBehaviour {
            ping: ping::Behaviour::new(Default::default()),
            identify: identify::Behaviour::new(
                identify::Config::new("p2cm/1.0".to_string(), keypair.public())
                    .with_interval(Duration::from_secs(10)),
            ),
            mdns: mdns::tokio::Behaviour::new(
                Default::default(),
                keypair.public().to_peer_id(),
            )?,
        };

        let swarm = libp2p::SwarmBuilder::with_new_identity()
            .with_tokio()
            .with_tcp(
                tcp::Config::default(),
                noise::Config::new,
                yamux::Config::default,
            )?
            .with_behaviour(|_| behaviour)?
            .with_swarm_config(|cfg| cfg.with_idle_connection_timeout(Duration::from_secs(u64::MAX)))
            .build();

        tracing::info!("Comms created with local peer ID: {}", local_peer_id);

        Ok(Self {
            peer_store,
            swarm: Some(swarm),
            local_peer_id,
        })
    }

    pub fn local_peer_id(&self) -> PeerId {
        self.local_peer_id
    }

    pub fn get_peer_store(&self) -> SharedPeerStore {
        self.peer_store.clone()
    }

    pub fn get_peers(&self) -> Vec<Peer> {
        self.peer_store.get_all_peers()
    }

    pub fn add_peer(&self, peer: Peer) {
        self.peer_store.add_peer(peer);
    }

    pub async fn start_discovery(&mut self) -> anyhow::Result<()> {
        let swarm = self.swarm.take().ok_or(CommsError::AlreadyRunning)?;
        let peer_store = self.peer_store.clone();

        tokio::spawn(async move {
            let mut swarm = swarm;

            loop {
                match swarm.select_next_some().await {
                    SwarmEvent::Behaviour(event) => {
                        match event {
                            CommsEvent::Mdns(mdns::Event::Discovered(discovered)) => {
                                for (peer_id, multiaddr) in discovered {
                                    tracing::info!("Peer discovered via mDNS: {} at {}",
                                        peer_id, multiaddr);
                                    let mut peer = Peer::new(peer_id);
                                    peer.add_address(multiaddr);
                                    peer_store.add_peer(peer);
                                }
                            }
                            CommsEvent::Mdns(mdns::Event::Expired(expired)) => {
                                for (peer_id, _) in expired {
                                    tracing::info!("Peer expired from mDNS: {}", peer_id);
                                    peer_store.remove_peer(&peer_id);
                                }
                            }
                            _ => {
                                tracing::debug!("Swarm event: {:?}", event);
                            }
                        }
                    }
                    SwarmEvent::NewListenAddr { address, .. } => {
                        tracing::info!("Listening on {}", address);
                    }
                    SwarmEvent::ConnectionEstablished { peer_id, .. } => {
                        tracing::info!("Connection established with {}", peer_id);
                    }
                    SwarmEvent::ConnectionClosed { peer_id, .. } => {
                        tracing::info!("Connection closed with {}", peer_id);
                    }
                    _ => {}
                }
            }
        });

        Ok(())
    }
}

#[derive(Debug, Clone, thiserror::Error)]
pub enum CommsError {
    #[error("No peer found")]
    NoPeerFound,
    #[error("Connection failed: {0}")]
    ConnectionFailed(String),
    #[error("Send failed")]
    SendFailed,
    #[error("Already running")]
    AlreadyRunning,
}