//! libp2p as a transport (`--transports libp2p`, in builds with the `libp2p` feature):
//! direct connections over TCP (Noise, Yamux) and QUIC, to addresses this node knows.
//!
//! A node's transport identity is its libp2p peer ID; where it is, its listen addresses
//! (multiaddrs ending in the peer ID, `locations`). Seeds are such multiaddrs, or ones
//! without the peer ID, which a connection finds out. Messages go one per stream of the
//! `/continuum/1` protocol: their length (4 bytes, big endian), then the message as
//! `wire` encodes it. A send is done once the stream is written and closed; a connection
//! that can't be made, or a stream that can't be written, refuses it.
//!
//! What this transport doesn't do yet (`docs/future-work.md`): learn addresses from the
//! peers it talks to (`identify`), find peers on the LAN (mDNS), or reach nodes behind
//! NAT (relays, hole punching). It never uses public libp2p infrastructure.

use std::collections::HashMap;
use std::path::Path;
use std::str::FromStr;
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use ::libp2p::multiaddr::Protocol;
use ::libp2p::swarm::dial_opts::{DialOpts, PeerCondition};
use ::libp2p::swarm::{ConnectionId, SwarmEvent};
use ::libp2p::{Multiaddr, PeerId, StreamProtocol, Swarm, identity::Keypair};
use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use futures::{AsyncReadExt, AsyncWriteExt, StreamExt};
use tokio::sync::{Semaphore, mpsc, oneshot, watch};
use tracing::{debug, info, warn};

use crate::locks::MutexExt;
use crate::transport::{EncodedMessage, MessageHandler, ParsedMessage, Transport, wire};

const PROTOCOL: StreamProtocol = StreamProtocol::new("/continuum/1");
/// The largest message read: anyone who can connect can send anything, so this bounds
/// what one stream costs. Sync messages above 16 chunks travel as blobs.
const MAX_MESSAGE: usize = 128 << 20;
/// How long a connection may take to come up.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// A connection nobody uses is closed after this; neighbours keep theirs busy.
const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
/// How long a stream may take to open and to carry a message, before what its size adds
/// (`stream_timeout`).
const STREAM_TIMEOUT: Duration = Duration::from_secs(10);
/// The slowest a message is expected to go, in bytes per second.
const MIN_RATE: u64 = 1 << 20;
/// Inbound streams read at once.
const READS: usize = 32;
/// Addresses kept per peer, and peers with addresses at most.
const ADDRESSES_PER_PEER: usize = 8;
const MAX_PEERS: usize = 4096;

/// How long a message of `bytes` may take on a stream.
fn stream_timeout(bytes: usize) -> Duration {
    STREAM_TIMEOUT + Duration::from_millis(bytes as u64 * 1000 / MIN_RATE)
}

/// What a libp2p transport is made with.
pub(crate) struct Libp2pConfig {
    pub(crate) key: Keypair,
    /// Multiaddrs to listen on.
    pub(crate) listen: Vec<Multiaddr>,
    /// libp2p's in-memory transport instead of TCP and QUIC: for tests.
    pub(crate) memory: bool,
}

impl Libp2pConfig {
    /// The key in `path` (32 raw bytes, an Ed25519 secret key), or a new one each start
    /// without a file: the peer ID only says where a node is, so a new one costs nothing
    /// but the WHO that finds it again (`transport::identity`).
    pub(crate) fn load_key(path: Option<&Path>) -> Result<Keypair> {
        let Some(path) = path else {
            return Ok(Keypair::generate_ed25519());
        };
        let mut bytes = std::fs::read(path)
            .with_context(|| format!("Reading the libp2p key from {}", path.display()))?;
        crate::swarm::auth::warn_if_readable_by_others(path, "libp2p key");
        if bytes.len() != 32 {
            bail!(
                "The libp2p key in {} has {} bytes; it has 32",
                path.display(),
                bytes.len()
            );
        }
        Keypair::ed25519_from_bytes(&mut bytes)
            .map_err(|error| anyhow!("The libp2p key in {}: {error}", path.display()))
    }
}

enum Command {
    /// Connect to `peer` (at `addresses`) unless connected.
    Connect {
        peer: PeerId,
        addresses: Vec<Multiaddr>,
        done: oneshot::Sender<Result<(), String>>,
    },
    /// Connect to `address` to find out which peer is there.
    Identify {
        address: Multiaddr,
        done: oneshot::Sender<Result<PeerId, String>>,
    },
}

/// Where peers are, and where we listen.
#[derive(Default)]
struct Book {
    addresses: HashMap<PeerId, Vec<Multiaddr>>,
    listening: Vec<Multiaddr>,
}

impl Book {
    /// `peer` is at `address` (without its `/p2p` part): kept first.
    fn remember(&mut self, peer: PeerId, mut address: Multiaddr) {
        if let Some(Protocol::P2p(_)) = address.iter().last() {
            address.pop();
        }
        if !self.addresses.contains_key(&peer) && self.addresses.len() >= MAX_PEERS {
            return;
        }
        let addresses = self.addresses.entry(peer).or_default();
        addresses.retain(|known| *known != address);
        addresses.insert(0, address);
        addresses.truncate(ADDRESSES_PER_PEER);
    }
}

/// `address` with the peer ID it ends in, if it ends in one.
fn split_peer(address: &Multiaddr) -> Option<PeerId> {
    match address.iter().last() {
        Some(Protocol::P2p(peer)) => Some(peer),
        _ => None,
    }
}

pub(crate) struct Libp2pTransport {
    local_id: String,
    peer_id: PeerId,
    control: libp2p_stream::Control,
    commands: mpsc::UnboundedSender<Command>,
    book: Arc<Mutex<Book>>,
    incoming: tokio::sync::Mutex<Option<libp2p_stream::IncomingStreams>>,
    stopped: watch::Receiver<bool>,
}

fn lock(book: &Mutex<Book>) -> MutexGuard<'_, Book> {
    book.lock_or_recover()
}

impl Libp2pTransport {
    /// Builds the swarm, listens, and runs it in a task of its own.
    pub(crate) fn start(config: Libp2pConfig) -> Result<Self> {
        let peer_id = config.key.public().to_peer_id();
        let mut swarm = if config.memory {
            memory_swarm(config.key)?
        } else {
            network_swarm(config.key)?
        };
        for address in config.listen {
            swarm
                .listen_on(address.clone())
                .with_context(|| format!("Listening on {address}"))?;
        }
        let mut control = swarm.behaviour().new_control();
        let incoming = control
            .accept(PROTOCOL)
            .map_err(|error| anyhow!("Accepting {PROTOCOL}: {error}"))?;
        let book = Arc::new(Mutex::new(Book::default()));
        let (commands, receiver) = mpsc::unbounded_channel();
        let (stopping, stopped) = watch::channel(false);
        tokio::spawn(run(swarm, receiver, Arc::clone(&book), stopping));
        info!(peer_id = %peer_id, "libp2p transport started");
        Ok(Self {
            local_id: peer_id.to_string(),
            peer_id,
            control,
            commands,
            book,
            incoming: tokio::sync::Mutex::new(Some(incoming)),
            stopped,
        })
    }

    /// Connected to `peer`, if it can be had in time.
    async fn connect(&self, peer: PeerId) -> Result<()> {
        let addresses = lock(&self.book)
            .addresses
            .get(&peer)
            .cloned()
            .unwrap_or_default();
        let (done, connected) = oneshot::channel();
        self.commands
            .send(Command::Connect {
                peer,
                addresses,
                done,
            })
            .map_err(|_| anyhow!("the libp2p swarm stopped"))?;
        match tokio::time::timeout(CONNECT_TIMEOUT, connected).await {
            Ok(Ok(Ok(()))) => Ok(()),
            Ok(Ok(Err(error))) => bail!("couldn't connect to {peer}: {error}"),
            Ok(Err(_)) => bail!("the libp2p swarm stopped"),
            Err(_) => bail!("connecting to {peer} took longer than {CONNECT_TIMEOUT:?}"),
        }
    }
}

type Behaviour = libp2p_stream::Behaviour;

fn network_swarm(key: Keypair) -> Result<Swarm<Behaviour>> {
    Ok(::libp2p::SwarmBuilder::with_existing_identity(key)
        .with_tokio()
        .with_tcp(
            ::libp2p::tcp::Config::default(),
            ::libp2p::noise::Config::new,
            ::libp2p::yamux::Config::default,
        )?
        .with_quic()
        .with_dns()?
        .with_behaviour(|_| Behaviour::new())?
        .with_swarm_config(|config| config.with_idle_connection_timeout(IDLE_TIMEOUT))
        .build())
}

fn memory_swarm(key: Keypair) -> Result<Swarm<Behaviour>> {
    use ::libp2p::Transport as _;
    use ::libp2p::core::transport::MemoryTransport;
    use ::libp2p::core::upgrade::Version;

    let peer_id = key.public().to_peer_id();
    let transport = MemoryTransport::default()
        .upgrade(Version::V1)
        .authenticate(::libp2p::noise::Config::new(&key)?)
        .multiplex(::libp2p::yamux::Config::default())
        .timeout(CONNECT_TIMEOUT)
        .boxed();
    Ok(Swarm::new(
        transport,
        Behaviour::new(),
        peer_id,
        ::libp2p::swarm::Config::with_tokio_executor().with_idle_connection_timeout(IDLE_TIMEOUT),
    ))
}

/// Drives the swarm: dials what `commands` ask for and notes where we listen, until the
/// transport is dropped.
async fn run(
    mut swarm: Swarm<Behaviour>,
    mut commands: mpsc::UnboundedReceiver<Command>,
    book: Arc<Mutex<Book>>,
    stopping: watch::Sender<bool>,
) {
    let mut connecting: HashMap<PeerId, Vec<oneshot::Sender<Result<(), String>>>> = HashMap::new();
    let mut identifying: HashMap<
        ConnectionId,
        (Multiaddr, oneshot::Sender<Result<PeerId, String>>),
    > = HashMap::new();
    loop {
        tokio::select! {
            command = commands.recv() => match command {
                None => break,
                Some(Command::Connect { peer, addresses, done }) => {
                    if swarm.is_connected(&peer) {
                        let _ = done.send(Ok(()));
                        continue;
                    }
                    let waiting = connecting.entry(peer).or_default();
                    waiting.push(done);
                    if waiting.len() > 1 {
                        continue; // Already dialing.
                    }
                    let opts = DialOpts::peer_id(peer)
                        .addresses(addresses)
                        .condition(PeerCondition::DisconnectedAndNotDialing)
                        .build();
                    if let Err(error) = swarm.dial(opts) {
                        for done in connecting.remove(&peer).into_iter().flatten() {
                            let _ = done.send(Err(error.to_string()));
                        }
                    }
                }
                Some(Command::Identify { address, done }) => {
                    let opts = DialOpts::unknown_peer_id().address(address.clone()).build();
                    let id = opts.connection_id();
                    match swarm.dial(opts) {
                        Ok(()) => {
                            identifying.insert(id, (address, done));
                        }
                        Err(error) => {
                            let _ = done.send(Err(error.to_string()));
                        }
                    }
                }
            },
            event = swarm.select_next_some() => match event {
                SwarmEvent::ConnectionEstablished { peer_id, connection_id, .. } => {
                    for done in connecting.remove(&peer_id).into_iter().flatten() {
                        let _ = done.send(Ok(()));
                    }
                    if let Some((address, done)) = identifying.remove(&connection_id) {
                        lock(&book).remember(peer_id, address);
                        let _ = done.send(Ok(peer_id));
                    }
                }
                SwarmEvent::OutgoingConnectionError { connection_id, peer_id, error } => {
                    if let Some((_, done)) = identifying.remove(&connection_id) {
                        let _ = done.send(Err(error.to_string()));
                    }
                    if let Some(peer) = peer_id
                        && !swarm.is_connected(&peer)
                    {
                        for done in connecting.remove(&peer).into_iter().flatten() {
                            let _ = done.send(Err(error.to_string()));
                        }
                    }
                }
                SwarmEvent::NewListenAddr { address, .. } => {
                    info!(address = %address, "libp2p listening");
                    let mut book = lock(&book);
                    if !book.listening.contains(&address) {
                        book.listening.push(address);
                    }
                }
                SwarmEvent::ExpiredListenAddr { address, .. } => {
                    info!(address = %address, "libp2p no longer listening");
                    lock(&book).listening.retain(|known| *known != address);
                }
                SwarmEvent::ListenerError { error, .. } => {
                    warn!(error = %error, "libp2p listener failed");
                }
                _ => {}
            },
        }
    }
    let _ = stopping.send(true);
}

/// One message from an inbound stream.
async fn read_message(mut stream: ::libp2p::Stream) -> Result<Vec<u8>> {
    let mut length = [0u8; 4];
    tokio::time::timeout(STREAM_TIMEOUT, stream.read_exact(&mut length))
        .await
        .map_err(|_| anyhow!("no length within {STREAM_TIMEOUT:?}"))??;
    let length = u32::from_be_bytes(length) as usize;
    if length > MAX_MESSAGE {
        bail!("a message of {length} bytes; at most {MAX_MESSAGE} are read");
    }
    let mut body = vec![0u8; length];
    tokio::time::timeout(stream_timeout(length), stream.read_exact(&mut body))
        .await
        .map_err(|_| {
            anyhow!(
                "{length} bytes took longer than {:?}",
                stream_timeout(length)
            )
        })??;
    Ok(body)
}

#[async_trait]
impl Transport for Libp2pTransport {
    fn local_id(&self) -> &str {
        &self.local_id
    }

    fn is_peer_id(&self, id: &str) -> bool {
        PeerId::from_str(id).is_ok()
    }

    /// The longest a message may take on a stream.
    fn delivery_limit(&self) -> Duration {
        stream_timeout(MAX_MESSAGE)
    }

    /// Where we listen, each with our peer ID.
    fn local_addresses(&self) -> Vec<String> {
        lock(&self.book)
            .listening
            .iter()
            .filter_map(|address| address.clone().with_p2p(self.peer_id).ok())
            .map(|address| address.to_string())
            .collect()
    }

    /// A peer ID, or a multiaddr ending in one (noted as where that peer is).
    fn locate(&self, target: &str) -> Option<String> {
        if let Ok(peer) = PeerId::from_str(target) {
            return Some(peer.to_string());
        }
        let address = Multiaddr::from_str(target).ok()?;
        let peer = split_peer(&address)?;
        if peer != self.peer_id {
            lock(&self.book).remember(peer, address);
        }
        Some(peer.to_string())
    }

    /// Where `target` is, each address with its peer ID.
    fn locations(&self, target: &str) -> Vec<String> {
        let Ok(peer) = PeerId::from_str(target) else {
            return Vec::new();
        };
        lock(&self.book)
            .addresses
            .get(&peer)
            .into_iter()
            .flatten()
            .filter_map(|address| address.clone().with_p2p(peer).ok())
            .map(|address| address.to_string())
            .collect()
    }

    /// A multiaddr ending in a peer ID is that peer; one without is connected to, to
    /// find out. A bare peer ID resolves if this transport knows where it is.
    async fn resolve_seed(&self, seed: &str) -> Result<String> {
        if let Ok(peer) = PeerId::from_str(seed) {
            if lock(&self.book).addresses.contains_key(&peer) {
                return Ok(peer.to_string());
            }
            bail!("no address known for libp2p peer {peer}; give a multiaddr");
        }
        let address =
            Multiaddr::from_str(seed).with_context(|| format!("{seed} is no multiaddr"))?;
        if let Some(peer) = self.locate(seed) {
            return Ok(peer);
        }
        let (done, identified) = oneshot::channel();
        self.commands
            .send(Command::Identify { address, done })
            .map_err(|_| anyhow!("the libp2p swarm stopped"))?;
        match tokio::time::timeout(CONNECT_TIMEOUT, identified).await {
            Ok(Ok(Ok(peer))) => Ok(peer.to_string()),
            Ok(Ok(Err(error))) => bail!("couldn't connect to {seed}: {error}"),
            Ok(Err(_)) => bail!("the libp2p swarm stopped"),
            Err(_) => bail!("connecting to {seed} took longer than {CONNECT_TIMEOUT:?}"),
        }
    }

    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
        let body = wire::encode(message)?;
        Ok(EncodedMessage {
            peer: peer.to_owned(),
            address: peer.to_owned(),
            payload_len: body.len(),
            body: body.into(),
            kind: message.kind(),
        })
    }

    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
        let peer = PeerId::from_str(&message.peer)
            .with_context(|| format!("{} is no libp2p peer ID", message.peer))?;
        self.connect(peer).await?;
        let mut control = self.control.clone();
        let body = message.body.clone();
        let write = async move {
            let mut stream = control.open_stream(peer, PROTOCOL).await?;
            stream.write_all(&(body.len() as u32).to_be_bytes()).await?;
            stream.write_all(&body).await?;
            stream.close().await?;
            Ok::<(), anyhow::Error>(())
        };
        let timeout = stream_timeout(message.body.len());
        tokio::time::timeout(timeout, write)
            .await
            .map_err(|_| anyhow!("sending to {peer} took longer than {timeout:?}"))?
            .with_context(|| format!("Sending to {peer}"))
    }

    /// Reads inbound streams (`READS` at once) and handles their messages one at a
    /// time, in the order they were read.
    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
        let Some(mut incoming) = self.incoming.lock().await.take() else {
            bail!("libp2p consumer already started");
        };
        let (sender, mut received) = mpsc::channel::<(PeerId, Vec<u8>)>(READS);
        let reads = Arc::new(Semaphore::new(READS));
        tokio::spawn(async move {
            while let Some((peer, stream)) = incoming.next().await {
                let Ok(permit) = Arc::clone(&reads).acquire_owned().await else {
                    break;
                };
                let sender = sender.clone();
                tokio::spawn(async move {
                    match read_message(stream).await {
                        Ok(body) => {
                            let _ = sender.send((peer, body)).await;
                        }
                        Err(error) => {
                            debug!(peer = %peer, error = %format!("{error:#}"), "Dropping an unreadable libp2p stream");
                        }
                    }
                    drop(permit);
                });
            }
        });
        while let Some((peer, body)) = received.recv().await {
            match wire::decode(&body) {
                Ok(message) => {
                    if let Err(error) = handler(peer.to_string(), message, body.len()).await {
                        warn!(peer = %peer, error = %format!("{error:#}"), "Failed to process inbound message");
                    }
                }
                Err(error) => {
                    debug!(peer = %peer, error = %error, "Dropping a libp2p message that doesn't decode");
                }
            }
        }
        Ok(())
    }

    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
        let mut stopped = self.stopped.clone();
        let _ = stopped.wait_for(|stopped| *stopped).await;
        Err(std::io::Error::other("the libp2p swarm stopped"))
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicU64, Ordering};

    use super::*;
    use crate::document::replica::Epoch;
    use crate::transport::{HeadsSummary, Keepalive};

    /// Memory ports are process-wide.
    static NEXT_PORT: AtomicU64 = AtomicU64::new(40_000);

    fn memory_node() -> (Libp2pTransport, Multiaddr) {
        let port = NEXT_PORT.fetch_add(1, Ordering::Relaxed);
        let key = Keypair::generate_ed25519();
        let address: Multiaddr = format!("/memory/{port}").parse().unwrap();
        let peer = key.public().to_peer_id();
        let transport = Libp2pTransport::start(Libp2pConfig {
            key,
            listen: vec![address.clone()],
            memory: true,
        })
        .unwrap();
        (transport, address.with_p2p(peer).unwrap())
    }

    fn ping() -> ParsedMessage {
        ParsedMessage::Ping(Keepalive {
            heads: HeadsSummary::of(Vec::new()),
            epoch: Epoch::GENESIS,
            epoch_seq: 3,
            ready: true,
            syncs_with_you: false,
        })
    }

    /// Starts `transport`'s consumer, forwarding what it receives.
    fn received(
        transport: Arc<Libp2pTransport>,
    ) -> mpsc::UnboundedReceiver<(String, ParsedMessage)> {
        let (sender, receiver) = mpsc::unbounded_channel();
        tokio::spawn(async move {
            transport
                .start_consumer(Box::new(move |from, message, _| {
                    let sender = sender.clone();
                    Box::pin(async move {
                        let _ = sender.send((from, message));
                        Ok(())
                    })
                }))
                .await
        });
        receiver
    }

    #[tokio::test]
    async fn messages_both_ways_over_one_connection() {
        let (a, a_address) = memory_node();
        let (b, _) = memory_node();
        let (a, b) = (Arc::new(a), Arc::new(b));
        let mut at_a = received(Arc::clone(&a));
        let mut at_b = received(Arc::clone(&b));

        assert_eq!(
            b.resolve_seed(&a_address.to_string()).await.unwrap(),
            a.local_id()
        );
        b.send_message(a.local_id(), &ping()).await.unwrap();
        let (from, message) = at_a.recv().await.unwrap();
        assert_eq!(from, b.local_id());
        assert!(matches!(message, ParsedMessage::Ping(keepalive) if keepalive.epoch_seq == 3));

        // a knows no address of b, but b's connection is there.
        a.send_message(b.local_id(), &ping()).await.unwrap();
        let (from, _) = at_b.recv().await.unwrap();
        assert_eq!(from, a.local_id());
    }

    #[tokio::test]
    async fn a_seed_without_a_peer_id_is_connected_to_to_find_out() {
        let (a, a_address) = memory_node();
        let (b, _) = memory_node();
        let mut bare = a_address.clone();
        bare.pop();
        assert_eq!(
            b.resolve_seed(&bare.to_string()).await.unwrap(),
            a.local_id()
        );
        // And b now knows where a is.
        assert_eq!(b.locations(a.local_id()), [a_address.to_string()]);
    }

    #[tokio::test]
    async fn listen_addresses_carry_the_peer_id() {
        let (a, a_address) = memory_node();
        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
        while a.local_addresses().is_empty() && tokio::time::Instant::now() < deadline {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
        assert_eq!(a.local_addresses(), [a_address.to_string()]);
        assert_eq!(
            a.locate(&a_address.to_string()).as_deref(),
            Some(a.local_id())
        );
    }

    #[tokio::test]
    async fn unknown_peers_are_refused() {
        let (b, _) = memory_node();
        let stranger = Keypair::generate_ed25519()
            .public()
            .to_peer_id()
            .to_string();
        assert!(b.send_message(&stranger, &ping()).await.is_err());
        assert!(b.resolve_seed(&stranger).await.is_err());
    }
}