ma-core 0.10.11

#[cfg(not(target_arch = "wasm32"))]
use web_time::{SystemTime, UNIX_EPOCH};
// Iroh-backed [`MaEndpoint`] implementation.

use async_trait::async_trait;
use iroh::{
    endpoint::{presets, Connection},
    protocol::{AcceptError, ProtocolHandler, Router},
    Endpoint, EndpointAddr, EndpointId, SecretKey,
};
use tracing::{debug, warn};

use crate::endpoint::{MaEndpoint, DEFAULT_INBOX_CAPACITY};
use crate::error::{Error, Result};
use crate::inbox::Inbox;
use crate::iroh::channel::Channel;
use crate::outbox::{Outbox, OutboxWire};
use crate::transport::transport_string;
use crate::{Document, Message};
use std::collections::{BTreeMap, HashMap};
use std::sync::Mutex;
use tokio::sync::Mutex as AsyncMutex;

const DEFAULT_MAX_INBOUND_MESSAGE_SIZE: usize = 1024 * 1024;

/// An iroh-backed ma endpoint.
pub struct IrohEndpoint {
    endpoint: Endpoint,
    protocols: Vec<String>,
    inboxes: BTreeMap<String, Inbox<Message>>,
    router: Option<Router>,
    /// Cached open connections keyed by (`endpoint_id`, protocol).
    ///
    /// Keeping a `Connection` clone alive prevents iroh from sending
    /// `APPLICATION_CLOSE` when a `Channel` is dropped after sending,
    /// which would otherwise race with the receiver's `accept_bi()` loop.
    connection_cache: Mutex<HashMap<(String, String), Connection>>,
    /// Serialises concurrent fresh-connection attempts to the same peer.
    ///
    /// Multiple tasks that simultaneously find the cache empty would otherwise
    /// all try to `connect()` in parallel.  The second task to acquire this
    /// lock re-checks the cache and reuses the connection the first task just
    /// established, so only one real QUIC connection is ever opened per peer.
    connect_lock: AsyncMutex<()>,
}

impl IrohEndpoint {
    /// Create an endpoint from raw 32-byte secret key material.
    pub async fn new(secret_bytes: [u8; 32]) -> Result<Self> {
        let secret = SecretKey::from_bytes(&secret_bytes);
        let endpoint = Endpoint::builder(presets::N0)
            .secret_key(secret)
            .bind()
            .await
            .map_err(|e| Error::Transport(format!("endpoint bind failed: {e}")))?;
        endpoint.online().await;

        debug!(
            endpoint_id = %endpoint.id(),
            "iroh endpoint online"
        );

        Ok(Self {
            endpoint,
            protocols: Vec::new(),
            inboxes: BTreeMap::new(),
            router: None,
            connection_cache: Mutex::new(HashMap::new()),
            connect_lock: AsyncMutex::new(()),
        })
    }

    /// Access the underlying iroh endpoint (for Router setup, gossip, etc.).
    pub fn inner(&self) -> &Endpoint {
        &self.endpoint
    }

    /// Consume self and return the underlying iroh endpoint.
    pub fn into_inner(self) -> Endpoint {
        self.endpoint
    }

    /// The endpoint's typed iroh identifier.
    pub fn endpoint_id(&self) -> EndpointId {
        self.endpoint.id()
    }

    /// Open a persistent write-only [`Channel`] to a remote endpoint.
    ///
    /// # Errors
    /// Returns an error if target parsing, connection, or stream opening fails.
    pub async fn open(&self, target: &str, protocol: &str) -> Result<Channel> {
        let addr = self.resolve_addr(target)?;
        self.open_addr(addr, protocol).await
    }

    async fn open_addr(&self, addr: EndpointAddr, protocol: &str) -> Result<Channel> {
        let connection = self
            .endpoint
            .connect(addr, protocol.as_bytes())
            .await
            .map_err(|e| Error::Transport(format!("connect failed: {e}")))?;
        let (send, _recv) = connection
            .open_bi()
            .await
            .map_err(|e| Error::Transport(format!("open_bi failed: {e}")))?;
        Ok(Channel::new(connection, send))
    }

    /// Like [`open_addr`] but reuses a cached connection when available.
    ///
    /// The cache stores a `Connection` clone.  Because iroh `Connection` is an
    /// `Arc`-backed handle, keeping one clone in the cache prevents the QUIC
    /// connection from being torn down when the `Channel` returned to the
    /// caller is later dropped.  This eliminates the race between the sender
    /// dropping its connection handle and the receiver's `accept_bi()` loop.
    ///
    /// If the cached connection is stale (remote restarted, idle timeout, etc.)
    /// `open_bi()` will return an error; in that case the entry is evicted and
    /// a fresh connection is established and cached in its place.
    async fn open_addr_cached(
        &self,
        endpoint_id: &str,
        addr: EndpointAddr,
        protocol: &str,
    ) -> Result<Channel> {
        let cache_key = (endpoint_id.to_string(), protocol.to_string());

        // Fast path: try to reuse an existing open connection (no async lock needed).
        let cached = self
            .connection_cache
            .lock()
            .unwrap()
            .get(&cache_key)
            .cloned();

        if let Some(conn) = cached {
            match conn.open_bi().await {
                Ok((send, _recv)) => {
                    debug!(endpoint_id, protocol, "reusing cached connection");
                    return Ok(Channel::new(conn, send));
                }
                Err(err) => {
                    debug!(
                        endpoint_id,
                        protocol,
                        error = %err,
                        "cached connection stale, reconnecting"
                    );
                    self.connection_cache.lock().unwrap().remove(&cache_key);
                }
            }
        }

        // Slow path: serialise all concurrent fresh-connection attempts.
        // Tasks that were blocked here will re-check the cache below and
        // reuse the connection established by whoever went first.
        let _connect_guard = self.connect_lock.lock().await;

        // Re-check after acquiring the lock — another task may have just
        // connected and cached the result while we were waiting.
        let cached = self
            .connection_cache
            .lock()
            .unwrap()
            .get(&cache_key)
            .cloned();

        if let Some(conn) = cached {
            match conn.open_bi().await {
                Ok((send, _recv)) => {
                    debug!(
                        endpoint_id,
                        protocol, "reusing cached connection (post-lock)"
                    );
                    return Ok(Channel::new(conn, send));
                }
                Err(err) => {
                    debug!(
                        endpoint_id,
                        protocol,
                        error = %err,
                        "cached connection stale after lock, reconnecting"
                    );
                    self.connection_cache.lock().unwrap().remove(&cache_key);
                }
            }
        }

        // Establish a fresh connection and cache a clone.
        let connection = self
            .endpoint
            .connect(addr, protocol.as_bytes())
            .await
            .map_err(|e| Error::Transport(format!("connect failed: {e}")))?;
        let (send, _recv) = connection
            .open_bi()
            .await
            .map_err(|e| Error::Transport(format!("open_bi failed: {e}")))?;

        self.connection_cache
            .lock()
            .unwrap()
            .insert(cache_key, connection.clone());

        Ok(Channel::new(connection, send))
    }

    /// Shut down the endpoint.
    pub async fn close(&mut self) {
        // Gracefully close all cached connections before shutting down.
        let connections: Vec<Connection> = self
            .connection_cache
            .lock()
            .unwrap()
            .drain()
            .map(|(_, conn)| conn)
            .collect();
        for conn in connections {
            conn.close(0u32.into(), b"done");
        }

        if let Some(router) = self.router.take() {
            let _ = router.shutdown().await;
            return;
        }
        self.endpoint.close().await;
    }

    /// Start the inbound router for all registered services.
    pub fn start_router(&mut self) {
        if self.router.is_some() {
            return;
        }

        let mut builder = Router::builder(self.endpoint.clone());
        for protocol in &self.protocols {
            if let Some(inbox) = self.inboxes.get(protocol) {
                let handler = InboxProtocolHandler::new(protocol.clone(), inbox.clone());
                builder = builder.accept(protocol.as_bytes(), handler);
            }
        }

        self.router = Some(builder.spawn());
    }

    /// Remove a registered service protocol.
    ///
    /// Returns `true` when a service existed and was removed.
    /// If the router is already running, it is reloaded so ALPN handlers
    /// match the updated service set.
    pub fn remove_service(&mut self, protocol: &str) -> bool {
        let normalized = normalize_protocol(protocol);
        let removed = self.inboxes.remove(&normalized).is_some();
        if !removed {
            return false;
        }

        self.protocols.retain(|p| p != &normalized);
        self.reload_router_if_running();
        true
    }

    fn reload_router_if_running(&mut self) {
        if self.router.is_none() {
            return;
        }

        // Dropping `Router` aborts the old accept loop quickly; we then spawn
        // a new one with an updated protocol map.
        self.router.take();
        self.start_router();
    }

    fn resolve_addr(&self, endpoint_id: &str) -> Result<EndpointAddr> {
        let target_id: EndpointId = endpoint_id
            .trim()
            .parse()
            .map_err(|e| Error::Transport(format!("invalid endpoint id: {e}")))?;

        let mut addr = EndpointAddr::new(target_id);

        // Add local relay URL as a routing hint.
        // DNS-based address lookup is not available in wasm_browser, so without
        // a relay hint the connect will time out. Both endpoints use the N0
        // preset whose relays interconnect, so any N0 relay URL is a valid hint.
        if let Some(relay_url) = self.endpoint.addr().relay_urls().next() {
            addr = addr.with_relay_url(relay_url.clone());
        }
        Ok(addr)
    }
}

#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl MaEndpoint for IrohEndpoint {
    fn id(&self) -> String {
        self.endpoint.id().to_string()
    }

    fn service(&mut self, protocol: &str) -> Inbox<Message> {
        let normalized = normalize_protocol(protocol);
        if !self.protocols.contains(&normalized) {
            self.protocols.push(normalized.clone());
        }
        if let Some(existing) = self.inboxes.get(&normalized) {
            return existing.clone();
        }

        let inbox = Inbox::new(DEFAULT_INBOX_CAPACITY);
        self.inboxes.insert(normalized, inbox.clone());
        if self.router.is_some() {
            self.reload_router_if_running();
        } else {
            self.start_router();
        }
        inbox
    }

    fn services(&self) -> Vec<String> {
        let id = self.endpoint.id().to_string();
        self.protocols
            .iter()
            .map(|proto| transport_string(&id, proto))
            .collect()
    }

    async fn connect_outbox(
        &self,
        _doc: &Document,
        endpoint_id: &str,
        did: &str,
        protocol: &str,
    ) -> Result<Outbox> {
        if endpoint_id == self.id() {
            let normalized = normalize_protocol(protocol);
            let inbox = self
                .inboxes
                .get(&normalized)
                .ok_or_else(|| Error::NoInboxTransport(format!("no local inbox for {protocol}")))?;
            return Ok(Outbox::from_transport(
                LoopbackWire {
                    inbox: inbox.clone(),
                },
                did.to_string(),
                protocol.to_string(),
            ));
        }
        let addr = self.resolve_addr(endpoint_id)?;
        let channel = self.open_addr_cached(endpoint_id, addr, protocol).await?;
        Ok(Outbox::from_transport(
            channel,
            did.to_string(),
            protocol.to_string(),
        ))
    }

    async fn send_to(&self, target: &str, protocol: &str, message: &Message) -> Result<()> {
        message.headers().validate()?;
        if target == self.id() {
            let normalized = normalize_protocol(protocol);
            let inbox = self
                .inboxes
                .get(&normalized)
                .ok_or_else(|| Error::NoInboxTransport(format!("no local inbox for {protocol}")))?;
            let expires_at = if message.exp == 0 {
                0
            } else {
                message.exp / 1_000_000_000
            };
            inbox.push(now_secs(), expires_at, message.clone());
            return Ok(());
        }
        let cbor = message.encode()?;
        let addr = self.resolve_addr(target)?;
        let mut channel = self.open_addr_cached(target, addr, protocol).await?;
        channel.send(&cbor).await?;
        channel.close();
        Ok(())
    }

    async fn close(&mut self) {
        IrohEndpoint::close(self).await;
    }
}

/// Delivers a message directly into a local [`Inbox`] without going through
/// the iroh transport layer.
///
/// iroh rejects QUIC connections where the target is the sender's own endpoint
/// ID. For self-addressed messages we bypass the network entirely and push
/// straight into the registered inbox — which is both correct per Hewitt's
/// actor model and more efficient than a loopback network hop.
#[derive(Debug)]
struct LoopbackWire {
    inbox: Inbox<Message>,
}

#[async_trait]
impl OutboxWire for LoopbackWire {
    async fn send_payload(&mut self, payload: &[u8]) -> Result<()> {
        let message = Message::decode(payload)?;
        message.headers().validate()?;
        let expires_at = if message.exp == 0 {
            0
        } else {
            message.exp / 1_000_000_000
        };
        self.inbox.push(now_secs(), expires_at, message);
        Ok(())
    }

    fn close_box(self: Box<Self>) {}
}

#[derive(Debug, Clone)]
struct InboxProtocolHandler {
    protocol: String,
    inbox: Inbox<Message>,
    max_message_size: usize,
}

impl InboxProtocolHandler {
    fn new(protocol: String, inbox: Inbox<Message>) -> Self {
        Self {
            protocol,
            inbox,
            max_message_size: DEFAULT_MAX_INBOUND_MESSAGE_SIZE,
        }
    }
}

impl ProtocolHandler for InboxProtocolHandler {
    async fn accept(&self, connection: Connection) -> std::result::Result<(), AcceptError> {
        loop {
            let (mut send, mut recv) = match connection.accept_bi().await {
                Ok(streams) => streams,
                Err(err) => {
                    debug!(
                        protocol = %self.protocol,
                        remote = %connection.remote_id(),
                        error = %err,
                        "inbound connection closed"
                    );
                    break;
                }
            };

            let payload = match recv.read_to_end(self.max_message_size).await {
                Ok(payload) => payload,
                Err(err) => {
                    warn!(
                        protocol = %self.protocol,
                        remote = %connection.remote_id(),
                        error = %err,
                        "failed to read inbound stream"
                    );
                    let _ = send.finish();
                    continue;
                }
            };

            let _ = send.finish();

            let message = match Message::decode(&payload) {
                Ok(message) => message,
                Err(err) => {
                    warn!(
                        protocol = %self.protocol,
                        remote = %connection.remote_id(),
                        error = %err,
                        "invalid inbound message payload"
                    );
                    continue;
                }
            };

            if let Err(err) = message.headers().validate() {
                warn!(
                    protocol = %self.protocol,
                    remote = %connection.remote_id(),
                    error = %err,
                    "invalid inbound message headers"
                );
                continue;
            }

            let expires_at = if message.exp == 0 {
                0
            } else {
                message.exp / 1_000_000_000
            };

            self.inbox.push(now_secs(), expires_at, message);
        }

        Ok(())
    }
}

fn normalize_protocol(input: &str) -> String {
    let protocol = input.trim();
    if protocol.is_empty() {
        return String::new();
    }

    format!("/{}", protocol.trim_start_matches('/'))
}

fn now_secs() -> u64 {
    #[cfg(target_arch = "wasm32")]
    {
        return (js_sys::Date::now() / 1000.0).floor() as u64;
    }

    #[cfg(not(target_arch = "wasm32"))]
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system clock before UNIX epoch")
        .as_secs()
}

#[allow(
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss
)]
fn message_created_at_secs(created_at: f64) -> u64 {
    if !created_at.is_finite() || created_at <= 0.0 {
        0
    } else if created_at >= u64::MAX as f64 {
        u64::MAX
    } else {
        created_at.floor() as u64
    }
}

#[cfg(test)]
mod tests {
    use crate::{Did, Document};

    fn test_doc() -> Document {
        let did = Did::new_url(
            "k51qzi5uqu5dj9807pbuod1pplf0vxh8m4lfy3ewl9qbm2s8dsf9ugdf9gedhr",
            None::<String>,
        )
        .expect("valid did");
        Document::new(&did, &did)
    }

    // ─── IrohEndpoint service/router lifecycle tests ─────────────────────────

    fn test_secret() -> [u8; 32] {
        let mut bytes = [0u8; 32];
        bytes[0] = 42;
        bytes
    }

    fn test_message() -> crate::Message {
        use crate::{Did, SigningKey};
        let did =
            Did::new_identity("k51qzi5uqu5dkkciu33khkzbcmxtyhn376i1e83tya8kuy7z9euedzyr5nhoew")
                .expect("valid did");
        let did_id = did.id();
        let sk = SigningKey::generate(did).expect("signing key");
        crate::Message::new(
            did_id,
            String::new(),
            crate::service::MESSAGE_TYPE_BROADCAST,
            "application/octet-stream",
            b"test".to_vec(),
            &sk,
        )
        .expect("message")
    }

    // Requires network (iroh endpoint bind); run with `cargo test -- --ignored`.
    #[tokio::test]
    #[ignore = "requires iroh network runtime"]
    async fn service_returns_shared_inbox() {
        use super::IrohEndpoint;
        use crate::endpoint::MaEndpoint;

        let mut endpoint = IrohEndpoint::new(test_secret()).await.unwrap();
        let inbox_a = endpoint.service("/ma/inbox/0.0.1");
        let inbox_b = endpoint.service("/ma/inbox/0.0.1");

        // Both clones point to the same underlying queue.
        inbox_a.push(0, 0, test_message());
        assert_eq!(inbox_b.len(), 1, "cloned inbox should share the same queue");

        endpoint.close().await;
    }

    #[tokio::test]
    #[ignore = "requires iroh network runtime"]
    async fn service_auto_starts_router() {
        use super::IrohEndpoint;
        use crate::endpoint::MaEndpoint;

        let mut endpoint = IrohEndpoint::new(test_secret()).await.unwrap();
        assert!(endpoint.router.is_none(), "router should start stopped");

        endpoint.service("/ma/inbox/0.0.1");

        assert!(
            endpoint.router.is_some(),
            "router should auto-start on first service registration"
        );

        endpoint.close().await;
    }

    #[tokio::test]
    #[ignore = "requires iroh network runtime"]
    async fn remove_service_updates_protocol_list() {
        use super::IrohEndpoint;
        use crate::endpoint::MaEndpoint;

        let mut endpoint = IrohEndpoint::new(test_secret()).await.unwrap();
        let _inbox = endpoint.service("/ma/custom/1.0");
        assert!(endpoint
            .services()
            .iter()
            .any(|s| s.contains("/ma/custom/1.0")));

        let removed = endpoint.remove_service("/ma/custom/1.0");
        assert!(
            removed,
            "remove_service should return true for registered protocol"
        );
        assert!(
            endpoint
                .services()
                .iter()
                .all(|s| !s.contains("/ma/custom/1.0")),
            "protocol should be absent from services after removal"
        );

        endpoint.close().await;
    }

    #[tokio::test]
    #[ignore = "requires iroh network runtime"]
    async fn service_after_start_router_triggers_reload() {
        use super::IrohEndpoint;
        use crate::endpoint::MaEndpoint;

        let mut endpoint = IrohEndpoint::new(test_secret()).await.unwrap();
        endpoint.service("/ma/inbox/0.0.1");
        endpoint.start_router();
        assert!(
            endpoint.router.is_some(),
            "router should be running after start_router"
        );

        // Adding a new service while router is running should transparently reload.
        endpoint.service("/ma/custom/1.0");
        assert!(
            endpoint.router.is_some(),
            "router should still be running after service addition"
        );
        assert!(
            endpoint
                .services()
                .iter()
                .any(|s| s.contains("/ma/custom/1.0")),
            "new service should appear in services() after hot-add"
        );

        endpoint.close().await;
    }
}