use std::time::Duration;

use libp2p::{
    kad::{store::MemoryStore, Behaviour as Kademlia, Config as KademliaConfig},
    PeerId, StreamProtocol,
};

pub type Discovery = Kademlia<MemoryStore>;

pub struct DiscoveryBuilder {
    max_packet_size: usize,
    record_ttl: Option<Duration>,
    provider_ttl: Option<Duration>,
    query_timeout: Duration,
}

impl DiscoveryBuilder {
    pub fn new() -> DiscoveryBuilder {
        DiscoveryBuilder {
            max_packet_size: 8192,
            record_ttl: None,
            provider_ttl: None,
            query_timeout: Duration::from_secs(60),
        }
    }

    pub fn record_ttl(mut self, record_ttl: Option<Duration>) -> Self {
        self.record_ttl = record_ttl;
        self
    }

    pub fn provider_ttl(mut self, provider_ttl: Option<Duration>) -> Self {
        self.provider_ttl = provider_ttl;
        self
    }

    pub fn query_timeout(mut self, query_timeout: Duration) -> Self {
        self.query_timeout = query_timeout;
        self
    }

    pub fn build(self, local_peer_id: PeerId, genesis_hash: &str) -> Discovery {
        let mut config = KademliaConfig::default();
        config.set_max_packet_size(self.max_packet_size);
        config.set_record_ttl(self.record_ttl);
        config.set_provider_record_ttl(self.provider_ttl);
        config.set_query_timeout(self.query_timeout);

        let kademlia_protocols = vec![
            StreamProtocol::try_from_owned(format!("/{genesis_hash}/kad"))
                .expect("Protocol name starts with '/'; qed"),
        ];
        config.set_protocol_names(kademlia_protocols.into_iter().map(Into::into).collect());

        let store = MemoryStore::new(local_peer_id);
        Kademlia::with_config(local_peer_id, store, config)
    }
}
