//! Addressing for iroh endpoints.
//!
//! This module contains some common addressing types for iroh.  An endpoint is uniquely
//! identified by the [`EndpointId`] but that does not make it addressable on the network layer.
//! For this the addition of a [`RelayUrl`] and/or direct addresses are required.
//!
//! The primary way of addressing an endpoint is by using the [`EndpointAddr`].

use std::{collections::BTreeSet, fmt, net::SocketAddr};

use data_encoding::HEXLOWER;
use n0_error::stack_error;
use serde::{Deserialize, Serialize};

use crate::{EndpointId, PublicKey, RelayUrl};

/// Network-level addressing information for an iroh endpoint.
///
/// This combines an endpoint's identifier with network-level addressing information of how to
/// contact the endpoint.
///
/// To establish a network connection to an endpoint both the [`EndpointId`] and one or more network
/// paths are needed.  The network paths can come from various sources:
///
/// - An [Address Lookup] service which can provide routing information for a given [`EndpointId`].
///
/// - A [`RelayUrl`] of the endpoint's [home relay], this allows establishing the connection via
///   the Relay server and is very reliable.
///
/// - One or more *IP based addresses* on which the endpoint might be reachable.  Depending on the
///   network location of both endpoints it might not be possible to establish a direct
///   connection without the help of a [Relay server].
///
/// This structure will always contain the required [`EndpointId`] and will contain an optional
/// number of other addressing information.  It is a generic addressing type used whenever a connection
/// to other endpoints needs to be established.
///
/// [Address Lookup]: https://docs.rs/iroh/*/iroh/index.html#address-lookup
/// [home relay]: https://docs.rs/iroh/*/iroh/relay/index.html
/// [Relay server]: https://docs.rs/iroh/*/iroh/index.html#relay-servers
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EndpointAddr {
    /// The endpoint's identifier.
    pub id: EndpointId,
    /// The endpoint's addresses.
    pub addrs: BTreeSet<TransportAddr>,
}

/// Available address types.
#[derive(
    derive_more::Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash,
)]
#[non_exhaustive]
pub enum TransportAddr {
    /// A relay server address.
    #[debug("Relay({_0})")]
    Relay(RelayUrl),
    /// An IP based address.
    Ip(SocketAddr),
    /// Custom transport address
    Custom(CustomAddr),
}

impl TransportAddr {
    /// Whether this is a transport address via a relay server.
    pub fn is_relay(&self) -> bool {
        matches!(self, Self::Relay(_))
    }

    /// Whether this is an IP transport address.
    pub fn is_ip(&self) -> bool {
        matches!(self, Self::Ip(_))
    }

    /// Whether this is a custom transport address.
    pub fn is_custom(&self) -> bool {
        matches!(self, Self::Custom(_))
    }
}

impl fmt::Display for TransportAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Relay(url) => write!(f, "relay:{url}"),
            Self::Ip(addr) => write!(f, "ip:{addr}"),
            Self::Custom(addr) => write!(f, "custom:{addr}"),
        }
    }
}

impl EndpointAddr {
    /// Creates a new [`EndpointAddr`] with no network level addresses.
    ///
    /// This still is usable with e.g. an address lookup service to establish a connection,
    /// depending on the situation.
    pub fn new(id: PublicKey) -> Self {
        EndpointAddr {
            id,
            addrs: Default::default(),
        }
    }

    /// Creates a new [`EndpointAddr`] from its parts.
    pub fn from_parts(id: PublicKey, addrs: impl IntoIterator<Item = TransportAddr>) -> Self {
        Self {
            id,
            addrs: addrs.into_iter().collect(),
        }
    }

    /// Adds a [`RelayUrl`] address.
    pub fn with_relay_url(mut self, relay_url: RelayUrl) -> Self {
        self.addrs.insert(TransportAddr::Relay(relay_url));
        self
    }

    /// Adds an IP based address.
    pub fn with_ip_addr(mut self, addr: SocketAddr) -> Self {
        self.addrs.insert(TransportAddr::Ip(addr));
        self
    }

    /// Adds a list of addresses.
    pub fn with_addrs(mut self, addrs: impl IntoIterator<Item = TransportAddr>) -> Self {
        for addr in addrs.into_iter() {
            self.addrs.insert(addr);
        }
        self
    }

    /// Returns true if only an [`EndpointId`] is present.
    pub fn is_empty(&self) -> bool {
        self.addrs.is_empty()
    }

    /// Returns an iterator over the IP addresses of this endpoint address.
    pub fn ip_addrs(&self) -> impl Iterator<Item = &SocketAddr> {
        self.addrs.iter().filter_map(|addr| match addr {
            TransportAddr::Ip(addr) => Some(addr),
            _ => None,
        })
    }

    /// Returns an iterator over the relay URLs of this endpoint address.
    ///
    ///  In practice this is expected to be zero or one home relay for all known cases currently.
    pub fn relay_urls(&self) -> impl Iterator<Item = &RelayUrl> {
        self.addrs.iter().filter_map(|addr| match addr {
            TransportAddr::Relay(url) => Some(url),
            _ => None,
        })
    }
}

impl From<EndpointId> for EndpointAddr {
    fn from(endpoint_id: EndpointId) -> Self {
        EndpointAddr::new(endpoint_id)
    }
}

/// A custom transport address consisting of a transport id and opaque address data.
///
/// This is a generic address type that allows external crates to implement custom
/// transports for iroh.
///
/// Transport ids are freely chosen u64 numbers. A registry for well-known transport ids
/// is maintained at <https://github.com/n0-computer/iroh/blob/main/TRANSPORTS.md>.
/// The opaque address data is not validated or size-limited in any way.
///
/// # String encoding
///
/// Used by [`Display`] and [`FromStr`] implementations.
/// Format: `<id>_<data>` where `<id>` is the transport id as lowercase hex (no `0x`
/// prefix, no leading zeros) and `<data>` is the address bytes as lowercase hex,
/// separated by `_`.
///
/// # Binary encoding
///
/// Used by [`Self::to_vec`] and [`Self::from_bytes`].
/// Format: 8-byte little-endian `u64` transport id, followed by raw address data bytes.
/// The minimum valid length is 8 bytes (id only with empty data).
///
/// [`Display`]: std::fmt::Display
/// [`FromStr`]: std::str::FromStr
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CustomAddr {
    /// The transport id.
    id: u64,
    /// Opaque address data for this transport.
    data: CustomAddrBytes,
}

impl fmt::Display for CustomAddr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:x}_{}", self.id, HEXLOWER.encode(self.data.as_bytes()))
    }
}

impl std::str::FromStr for CustomAddr {
    type Err = CustomAddrParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let Some((id_str, data_str)) = s.split_once('_') else {
            return Err(CustomAddrParseError::MissingSeparator);
        };
        let Ok(id) = u64::from_str_radix(id_str, 16) else {
            return Err(CustomAddrParseError::InvalidId);
        };
        let Ok(data) = HEXLOWER.decode(data_str.as_bytes()) else {
            return Err(CustomAddrParseError::InvalidData);
        };
        Ok(Self::from_parts(id, &data))
    }
}

/// Error returned when parsing a [`CustomAddr`] from its string encoding fails.
///
/// Parsing a string into a [`CustomAddr`] represents just the first part of
/// validation. Even if the string is well-formed, the resulting [`CustomAddr`] might
/// still have an invalid data size or format for the transport type.
#[stack_error(derive)]
#[allow(missing_docs)]
pub enum CustomAddrParseError {
    /// Missing `_` separator between id and data.
    #[error("missing '_' separator")]
    MissingSeparator,
    /// Invalid hex-encoded id.
    #[error("invalid id")]
    InvalidId,
    /// Invalid hex-encoded data.
    #[error("invalid data")]
    InvalidData,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum CustomAddrBytes {
    Inline { size: u8, data: [u8; 30] },
    Heap(Box<[u8]>),
}

impl fmt::Debug for CustomAddrBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !f.alternate() {
            write!(f, "[{}]", HEXLOWER.encode(self.as_bytes()))
        } else {
            let bytes = self.as_bytes();
            match self {
                Self::Inline { .. } => write!(f, "Inline[{}]", HEXLOWER.encode(bytes)),
                Self::Heap(_) => write!(f, "Heap[{}]", HEXLOWER.encode(bytes)),
            }
        }
    }
}

impl From<(u64, &[u8])> for CustomAddr {
    fn from((id, data): (u64, &[u8])) -> Self {
        Self::from_parts(id, data)
    }
}

impl CustomAddrBytes {
    fn len(&self) -> usize {
        match self {
            Self::Inline { size, .. } => *size as usize,
            Self::Heap(data) => data.len(),
        }
    }

    fn as_bytes(&self) -> &[u8] {
        match self {
            Self::Inline { size, data } => &data[..*size as usize],
            Self::Heap(data) => data,
        }
    }

    fn copy_from_slice(data: &[u8]) -> Self {
        if data.len() <= 30 {
            let mut inline = [0u8; 30];
            inline[..data.len()].copy_from_slice(data);
            Self::Inline {
                size: data.len() as u8,
                data: inline,
            }
        } else {
            Self::Heap(data.to_vec().into_boxed_slice())
        }
    }
}

impl CustomAddr {
    /// Creates a new [`CustomAddr`] from a transport id and raw address data.
    pub fn from_parts(id: u64, data: &[u8]) -> Self {
        Self {
            id,
            data: CustomAddrBytes::copy_from_slice(data),
        }
    }

    /// Returns the transport id.
    ///
    /// You can freely choose this. There is a table of reserved custom transport ids in
    /// <https://github.com/n0-computer/iroh/blob/main/TRANSPORTS.md>, where you could
    /// submit your transport for registration to get a reserved id.
    ///
    /// But this is only relevant if you care for interop.
    pub fn id(&self) -> u64 {
        self.id
    }

    /// Returns the opaque address data for this transport.
    ///
    /// Below a certain size (currently 30 bytes) this is stored inline, otherwise on the heap.
    ///
    /// Note that there are no guarantees about the size of this data. When parsing custom
    /// addresses you must be prepared to handle unexpected sizes here.
    pub fn data(&self) -> &[u8] {
        self.data.as_bytes()
    }

    /// Serializes to the binary encoding.
    ///
    /// See [`CustomAddr`] docs for details on the encoding.
    pub fn to_vec(&self) -> Vec<u8> {
        let mut out = vec![0u8; 8 + self.data.len()];
        out[..8].copy_from_slice(&self.id().to_le_bytes());
        out[8..].copy_from_slice(self.data());
        out
    }

    /// Parses from the binary encoding.
    ///
    /// See [`CustomAddr`] docs for details on the encoding.
    pub fn from_bytes(data: &[u8]) -> Result<Self, &'static str> {
        if data.len() < 8 {
            return Err("data too short");
        }
        let id = u64::from_le_bytes(data[..8].try_into().expect("data length checked above"));
        let data = &data[8..];
        Ok(Self::from_parts(id, data))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Hash)]
    #[non_exhaustive]
    enum NewAddrType {
        /// Relays
        Relay(RelayUrl),
        /// IP based addresses
        Ip(SocketAddr),
        /// New addr type for testing
        Cool(u16),
    }

    #[test]
    fn test_roundtrip_new_addr_type() {
        let old = vec![
            TransportAddr::Ip("127.0.0.1:9".parse().unwrap()),
            TransportAddr::Relay("https://example.com".parse().unwrap()),
        ];
        let old_ser = postcard::to_stdvec(&old).unwrap();
        let old_back: Vec<TransportAddr> = postcard::from_bytes(&old_ser).unwrap();
        assert_eq!(old, old_back);

        let new = vec![
            NewAddrType::Ip("127.0.0.1:9".parse().unwrap()),
            NewAddrType::Relay("https://example.com".parse().unwrap()),
            NewAddrType::Cool(4),
        ];
        let new_ser = postcard::to_stdvec(&new).unwrap();
        let new_back: Vec<NewAddrType> = postcard::from_bytes(&new_ser).unwrap();

        assert_eq!(new, new_back);

        // serialize old into new
        let old_new_back: Vec<NewAddrType> = postcard::from_bytes(&old_ser).unwrap();

        assert_eq!(
            old_new_back,
            vec![
                NewAddrType::Ip("127.0.0.1:9".parse().unwrap()),
                NewAddrType::Relay("https://example.com".parse().unwrap()),
            ]
        );
    }

    #[test]
    fn test_custom_addr_roundtrip() {
        // Small id, small data (e.g., Bluetooth MAC)
        let addr = CustomAddr::from_parts(1, &[0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6]);
        let s = addr.to_string();
        assert_eq!(s, "1_a1b2c3d4e5f6");
        let parsed: CustomAddr = s.parse().unwrap();
        assert_eq!(addr, parsed);

        // Larger id, 32-byte data (e.g., Tor pubkey)
        let addr = CustomAddr::from_parts(42, &[0xab; 32]);
        let s = addr.to_string();
        assert_eq!(
            s,
            "2a_abababababababababababababababababababababababababababababababab"
        );
        let parsed: CustomAddr = s.parse().unwrap();
        assert_eq!(addr, parsed);

        // Zero id, empty data
        let addr = CustomAddr::from_parts(0, &[]);
        let s = addr.to_string();
        assert_eq!(s, "0_");
        let parsed: CustomAddr = s.parse().unwrap();
        assert_eq!(addr, parsed);

        // Large id
        let addr = CustomAddr::from_parts(0xdeadbeef, &[0x01, 0x02]);
        let s = addr.to_string();
        assert_eq!(s, "deadbeef_0102");
        let parsed: CustomAddr = s.parse().unwrap();
        assert_eq!(addr, parsed);
    }

    #[test]
    fn test_custom_addr_parse_errors() {
        // Missing separator
        assert!("abc123".parse::<CustomAddr>().is_err());

        // Invalid id (not hex)
        assert!("xyz_0102".parse::<CustomAddr>().is_err());

        // Invalid data (not hex)
        assert!("1_ghij".parse::<CustomAddr>().is_err());

        // Odd-length hex data
        assert!("1_abc".parse::<CustomAddr>().is_err());
    }
}
