myst-client 0.1.1

pub mod gossip;

use anyhow::Error as AnyhowError;
use async_openai::types::{CreateChatCompletionRequest, CreateChatCompletionStreamResponse};
use base32::Alphabet;
use earendil_crypt::{HavenEndpoint, HavenFingerprint};
use futures::{Stream, StreamExt};
use hyper_util::rt::TokioIo;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::{
    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
    pin::Pin,
    str::FromStr,
};
use thiserror::Error;
use tokio_socks::tcp::Socks5Stream;
use tonic::transport::Channel;

use crate::{
    proto::myst_service_client::MystServiceClient,
    proto::{GossipRequest, ProxyRequest},
};

use self::gossip::NetworkState;

#[derive(Clone)]
pub struct RpcClient {
    inner: MystServiceClient<Channel>,
}

impl RpcClient {
    pub async fn new(
        node_id: NodeId,
        proxy_addr: Option<SocketAddr>,
    ) -> Result<Self, TransportError> {
        eprintln!("Creating new RPC client with node ID: {node_id}");

        let channel = match (node_id, proxy_addr) {
            (NodeId::Direct(addr), _) => Channel::from_shared(format!("http://{}", addr))
                .unwrap()
                .connect()
                .await
                .unwrap(),
            (NodeId::Earendil(haven_endpoint), Some(proxy_addr)) => {
                let haven_url = format!(
                    "{}.haven:{}",
                    haven_endpoint.fingerprint, haven_endpoint.port
                );

                Channel::from_shared(format!("http://{}", haven_url))
                    .map_err(|e| TransportError::Earendil(e.to_string()))?
                    .connect_with_connector(tower::service_fn(move |_| {
                        let proxy = proxy_addr;
                        let target = haven_url.clone();

                        async move {
                            // NOTE: This creates a new SOCKS5 connection for each request
                            // maybe we can cache/reuse the connection?
                            let socks_stream =
                                Socks5Stream::connect(proxy, target).await.map_err(|e| {
                                    std::io::Error::new(std::io::ErrorKind::Other, e.to_string())
                                })?;

                            Ok::<_, std::io::Error>(TokioIo::new(socks_stream.into_inner()))
                        }
                    }))
                    .await
                    .map_err(|e| TransportError::Earendil(e.to_string()))?
            }
            (NodeId::Tor(_), Some(_proxy_addr)) => {
                return Err(TransportError::Network(
                    "Tor transport not implemented".into(),
                ));
            }
            (NodeId::Nym(_), Some(_proxy_addr)) => {
                return Err(TransportError::Network(
                    "Nym transport not implemented".into(),
                ));
            }
            _ => return Err(TransportError::Network("unknown transport error".into())),
        };

        Ok(Self {
            inner: MystServiceClient::new(channel),
        })
    }

    pub async fn gossip(&mut self, state: NetworkState) -> anyhow::Result<()> {
        let bytes = serde_json::to_vec(&state)?;
        self.inner
            .gossip(GossipRequest {
                network_state: bytes,
            })
            .await?;
        Ok(())
    }

    pub async fn compute_text(
        &mut self,
        request: CreateChatCompletionRequest,
        node_id: Option<NodeId>,
    ) -> anyhow::Result<
        Pin<
            Box<
                dyn Stream<Item = Result<CreateChatCompletionStreamResponse, anyhow::Error>> + Send,
            >,
        >,
    > {
        let proxy_request = ProxyRequest {
            node_id: node_id.clone().map(|id| id.to_string()),
            request: serde_json::to_vec(&request)?,
        };

        eprintln!("sending proxy stream request to node: {:?}", &node_id);
        let response = self.inner.proxy_stream(proxy_request).await?;
        let stream = response.into_inner();

        // Transform tonic::Streaming into a standard Stream
        let transformed_stream = stream.map(|result| {
            result
                .map_err(anyhow::Error::from)
                .and_then(|proxy_response| {
                    serde_json::from_slice(&proxy_response.response).map_err(anyhow::Error::from)
                })
        });

        Ok(Box::pin(transformed_stream))
    }

    pub async fn generate_image(
        &mut self,
        prompt: String,
        node_id: Option<NodeId>,
    ) -> anyhow::Result<String> {
        let proxy_request = ProxyRequest {
            node_id: node_id.map(|id| id.to_string()),
            request: serde_json::to_vec(&prompt)?,
        };

        let response = self.inner.proxy_unary(proxy_request).await?;
        let bytes = response.into_inner().response;
        Ok(String::from_utf8(bytes)?) // Base64 encoded image
    }
}

#[derive(Error, Debug)]
pub enum TransportError {
    #[error("earendil connection error: {0}")]
    Earendil(String),
    #[error("TCP connection error: {0}")]
    Direct(String),
    #[error("Unknown error: {0}")]
    Network(String),
    #[error("Socks5 error: {0}")]
    Socks5(tokio_socks::Error),
}

impl From<AnyhowError> for TransportError {
    fn from(err: AnyhowError) -> Self {
        TransportError::Network(err.to_string())
    }
}

impl From<tokio_socks::Error> for TransportError {
    fn from(err: tokio_socks::Error) -> Self {
        TransportError::Socks5(err)
    }
}

/// ID for uniquely identifying a node in the network.
/// This wraps the underlying transport address to keep things simple.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum NodeId {
    Direct(SocketAddr),
    Earendil(HavenEndpoint),
    Tor(String),
    Nym(String),
}

impl Serialize for NodeId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.to_string())
    }
}

impl<'de> Deserialize<'de> for NodeId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        s.parse().map_err(serde::de::Error::custom)
    }
}

impl std::fmt::Display for NodeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NodeId::Direct(addr) => write!(f, "{}", addr),
            NodeId::Earendil(endpoint) => write!(f, "{}:{}", endpoint.fingerprint, endpoint.port),
            NodeId::Tor(addr) => write!(f, "{}", addr),
            NodeId::Nym(addr) => write!(f, "{}", addr),
        }
    }
}

impl FromStr for NodeId {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, anyhow::Error> {
        if s.is_empty() {
            return Err(anyhow::anyhow!("Empty string"));
        }

        // Try to parse as socket address first
        if let Ok(addr) = s.parse::<SocketAddr>() {
            return Ok(NodeId::Direct(addr));
        }

        // Try to parse as Earendil endpoint (fingerprint:port)
        if let Some((fingerprint, port)) = s.split_once(':') {
            if let (Ok(fingerprint), Ok(port)) =
                (HavenFingerprint::from_str(fingerprint), port.parse::<u16>())
            {
                return Ok(NodeId::Earendil(HavenEndpoint { fingerprint, port }));
            }
        }

        // Could add Tor/Nym parsing here if needed
        Err(anyhow::anyhow!("Invalid node ID format"))
    }
}