//! iroh QUIC transport implementation.
//!
//! Wraps an [`iroh::Endpoint`] as a [`MeshTransport`], using the same
//! length-prefixed framing as the existing [`P2pNode`](crate::P2pNode).

use anyhow::{bail, Result};
use iroh::{Endpoint, EndpointAddr, PublicKey, SecretKey};

use crate::transport::{MeshTransport, TransportAddr, TransportInfo, TransportPacket};

/// ALPN protocol identifier for the transport-abstracted mesh layer.
/// Distinct from `P2P_ALPN` to avoid conflicts with the existing P2pNode.
const MESH_ALPN: &[u8] = b"quicprochat/mesh/1";

/// iroh QUIC mesh transport.
///
/// Provides encrypted, NAT-traversing connections via iroh relay infrastructure.
pub struct IrohTransport {
    endpoint: Endpoint,
}

impl IrohTransport {
    /// Create a new iroh transport, binding a fresh endpoint.
    ///
    /// If `secret_key` is `None`, a random identity is generated.
    pub async fn new(secret_key: Option<SecretKey>) -> Result<Self> {
        let mut builder = Endpoint::builder();
        if let Some(sk) = secret_key {
            builder = builder.secret_key(sk);
        }
        builder = builder.alpns(vec![MESH_ALPN.to_vec()]);

        let endpoint = builder.bind().await?;

        tracing::info!(
            node_id = %endpoint.id().fmt_short(),
            "IrohTransport started"
        );

        Ok(Self { endpoint })
    }

    /// Create an `IrohTransport` from an already-bound endpoint.
    ///
    /// The caller must ensure the endpoint was configured with `MESH_ALPN`.
    pub fn from_endpoint(endpoint: Endpoint) -> Self {
        Self { endpoint }
    }

    /// Return this transport's iroh public key.
    pub fn public_key(&self) -> PublicKey {
        self.endpoint.id()
    }

    /// Return the endpoint address for sharing with peers.
    pub fn endpoint_addr(&self) -> EndpointAddr {
        self.endpoint.addr()
    }

    /// Convert a `TransportAddr::Iroh` into an `EndpointAddr`.
    fn to_endpoint_addr(addr: &TransportAddr) -> Result<EndpointAddr> {
        match addr {
            TransportAddr::Iroh(id) => {
                let key_bytes: [u8; 32] = id
                    .as_slice()
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("iroh addr must be 32 bytes, got {}", id.len()))?;
                let pk = PublicKey::from_bytes(&key_bytes)?;
                Ok(EndpointAddr::from(pk))
            }
            other => bail!("IrohTransport cannot send to {other}"),
        }
    }
}

#[async_trait::async_trait]
impl MeshTransport for IrohTransport {
    fn info(&self) -> TransportInfo {
        TransportInfo {
            name: "iroh-quic".to_string(),
            mtu: 65535,
            bitrate: 100_000_000,
            bidirectional: true,
        }
    }

    async fn send(&self, dest: &TransportAddr, data: &[u8]) -> Result<()> {
        let addr = Self::to_endpoint_addr(dest)?;
        let conn = self.endpoint.connect(addr, MESH_ALPN).await?;

        let mut send = conn.open_uni().await.map_err(|e| anyhow::anyhow!("{e}"))?;

        // Length-prefixed framing: [u32 BE length][payload].
        let len = (data.len() as u32).to_be_bytes();
        send.write_all(&len)
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))?;
        send.write_all(data)
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))?;
        send.finish().map_err(|e| anyhow::anyhow!("{e}"))?;
        send.stopped().await.map_err(|e| anyhow::anyhow!("{e}"))?;

        tracing::debug!(
            peer = %conn.remote_id().fmt_short(),
            bytes = data.len(),
            "IrohTransport: message sent"
        );

        Ok(())
    }

    async fn recv(&self) -> Result<TransportPacket> {
        let incoming = self
            .endpoint
            .accept()
            .await
            .ok_or_else(|| anyhow::anyhow!("no more incoming connections"))?;

        let conn = incoming.await.map_err(|e| anyhow::anyhow!("{e}"))?;
        let sender = conn.remote_id();

        let mut recv = conn
            .accept_uni()
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))?;

        // Read length-prefixed payload.
        let mut len_buf = [0u8; 4];
        recv.read_exact(&mut len_buf)
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))?;
        let len = u32::from_be_bytes(len_buf) as usize;

        if len > 5 * 1024 * 1024 {
            bail!("payload too large: {len} bytes");
        }

        let mut payload = vec![0u8; len];
        recv.read_exact(&mut payload)
            .await
            .map_err(|e| anyhow::anyhow!("{e}"))?;

        tracing::debug!(
            peer = %sender.fmt_short(),
            bytes = len,
            "IrohTransport: message received"
        );

        Ok(TransportPacket {
            from: TransportAddr::Iroh(sender.as_bytes().to_vec()),
            data: payload,
        })
    }

    async fn close(&self) -> Result<()> {
        self.endpoint.close().await;
        Ok(())
    }
}
