iroh_ex

#![no_main]
#![allow(unused_imports)]
#![allow(dead_code)]
#![allow(unused_variables)]
#![allow(deprecated)]
#![allow(unused_must_use)]
#![allow(non_local_definitions)]
// #![allow(unexpected_cfgs)]
// #[cfg(not(clippy))]

use chrono::Local;
use iroh::endpoint;
use rustler::{
    Encoder, Env, Error as RustlerError, LocalPid, NifResult, OwnedEnv, ResourceArc, Term,
};

use atty::Stream;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::FmtSubscriber;

use rand::Rng;

use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use tokio::runtime::Runtime;
use tokio::sync::mpsc;
use tokio::sync::RwLock;
use tokio::time::{sleep, timeout, Duration};

use std::fmt;
use std::ptr;
use std::str::FromStr;

use anyhow::{Context, Result};
use iroh::{
    endpoint::Connection,
    protocol::{ProtocolHandler, Router},
    Endpoint, NodeAddr, NodeId, SecretKey,
};

// use quic_rpc::transport::flume::FlumeConnector;

// pub(crate) type BlobsClient = iroh_blobs::rpc::client::blobs::Client<
//     FlumeConnector<iroh_blobs::rpc::proto::Response, iroh_blobs::rpc::proto::Request>,
// >;
// pub(crate) type DocsClient = iroh_docs::rpc::client::docs::Client<
//     FlumeConnector<iroh_docs::rpc::proto::Response, iroh_docs::rpc::proto::Request>,
// >;

// use iroh_gossip::{net::Gossip, ALPN as GossipALPN};
// use iroh_gossip::proto::TopicId;

use iroh_gossip::{
    net::{Event, Gossip, GossipEvent, GossipReceiver, GossipSender},
    proto::TopicId,
    ALPN as GossipALPN,
};

use iroh::discovery::DiscoveryItem;

use serde::{Deserialize, Serialize};

use n0_future::boxed::BoxFuture;
use n0_future::StreamExt;

use rand::distributions::Alphanumeric;

pub static RUNTIME: Lazy<Runtime> =
    Lazy::new(|| tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime"));

const ALPN: &[u8] = b"iroh-example/echo/0";

// const TOPIC_NAME: &str = "ehaaöskdjfasdjföasdjföa";

static TOPIC_NAME: Lazy<String> = Lazy::new(generate_topic_name);

fn generate_topic_name() -> String {
    rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(20)
        .collect()
}

// const topic_bytes = rand::random();
// static topic_bytes: [u8; 32] = rand::random();

// pub static TOPIC_BYTES: Lazy<[u8; 32]> =
//     Lazy::new(|| rand::random<[u8; 32]>().expect("Failed to create topic random bytes"));

#[derive(Debug, Serialize, Deserialize)]
enum Message {
    AboutMe { from: NodeId, name: String },
    Message { from: NodeId, text: String },
}

impl Message {
    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        serde_json::from_slice(bytes).map_err(Into::into)
    }

    pub fn to_vec(&self) -> Vec<u8> {
        serde_json::to_vec(self).expect("serde_json::to_vec is infallible")
    }
}

pub struct NodeRef(pub(crate) Arc<Mutex<NodeState>>);

pub struct NodeState {
    pub pid: LocalPid,
    pub endpoint: Endpoint,
    pub router: Router,
    pub gossip: Gossip,
    pub sender: GossipSender,
}

impl NodeState {
    pub fn new(
        pid: LocalPid,
        endpoint: Endpoint,
        router: Router,
        gossip: Gossip,
        sender: GossipSender,
    ) -> Self {
        NodeState {
            pid,
            endpoint,
            router,
            gossip,
            sender,
        }
    }
}

// impl Drop for NodeState {
//     fn drop(&mut self) {
//         tracing::info!("🚀 Cleaning up NodeState before exit!");
//         RUNTIME.block_on(self.gossip.shutdown());
//         RUNTIME.block_on(self.router.shutdown());
//         RUNTIME.block_on(self.endpoint.close());
//         tracing::info!("✅ NodeState cleanup complete!");
//         // .map_err(|e| RustlerError::Term(Box::new(format!("Endpoint error: {}", e))))?;
//         // RUNTIME.spawn(async move {
//         //     self.gossip.shutdown().await;
//         //     self.router.shutdown();
//         //     self.endpoint.close().await;
//         //     tracing::info!("✅ NodeState cleanup complete!");
//         // });
//     }
// }

struct GossipWrapper {
    gossip: iroh_gossip::net::Gossip,
}

impl Drop for GossipWrapper {
    fn drop(&mut self) {
        tracing::debug!(
            "🚀 GossipWrapper: Gossip {:?} is being dropped!",
            self.gossip
        );

        tracing::debug!(
            "🗑️ GossipWrapper: Gossip {:?} at address: {:p}",
            self.gossip,
            ptr::addr_of!(self)
        );
    }
}

struct EndpointWrapper {
    endpoint: iroh::Endpoint,
}

impl Drop for EndpointWrapper {
    fn drop(&mut self) {
        tracing::debug!(
            "🚀 EndpointWrapper: Endpoint {:?} is being dropped!",
            self.endpoint.node_id().fmt_short()
        );
    }
}

mod atoms {
    rustler::atoms! {
        ok,
        error,

        // errors
        lock_fail,
        not_found,
        offer_error,

        candidate_error,
        iroh_gossip_joined,
        iroh_gossip_neighbor_up,
        iroh_gossip_neighbor_down,
        iroh_gossip_node_discovered,
        iroh_gossip_message_received,
        iroh_gossip_message_unhandled,
    }
}

fn string_to_32_byte_array(s: &str) -> [u8; 32] {
    let mut result = [0u8; 32];
    let bytes = s.as_bytes();
    let len = std::cmp::min(bytes.len(), 32);
    result[..len].copy_from_slice(&bytes[..len]);
    result
}

#[rustler::nif]
pub fn generate_secretkey(env: Env) -> Result<String, RustlerError> {
    // let secret_key = SecretKey::generate(rand::rngs::OsRng);
    //println!("secret key: {secret_key}");
    let secret_key = "blabalba";
    Ok(secret_key.to_string())
}

#[rustler::nif(schedule = "DirtyCpu")]
pub fn create_node_async(env: Env, pid: LocalPid) -> Result<ResourceArc<NodeRef>, RustlerError> {
    // Block inside DirtyCpu to safely wait for the async task
    tokio::task::block_in_place(|| {
        RUNTIME.block_on(async { create_node_async_internal(pid).await })
    })
}

async fn create_node_async_internal(pid: LocalPid) -> Result<ResourceArc<NodeRef>, RustlerError> {
    let endpoint = Endpoint::builder()
        .discovery_local_network()
        // .discovery_n0()
        .bind()
        .await
        .map_err(|e| RustlerError::Term(Box::new(format!("Endpoint error: {}", e))))?;

    println!("Endpoint node id: {:?}", endpoint.node_id());

    let endpoint_clone = endpoint.clone();
    let mut builder = iroh::protocol::Router::builder(endpoint.clone());

    let gossip = Gossip::builder()
        .spawn(endpoint.clone())
        .await
        .map_err(|e| RustlerError::Term(Box::new(format!("Gossip protocol error: {}", e))))?;

    builder = builder.accept(GossipALPN, gossip.clone());
    builder = builder.accept(ALPN, Echo);

    let router = builder
        .spawn()
        .await
        .map_err(|e| RustlerError::Term(Box::new(format!("Router error: {}", e))))?;

    let router_clone = router.clone();

    let node_addr = router
        .endpoint()
        .node_addr()
        .await
        .map_err(|e| RustlerError::Term(Box::new(format!("Node addr error: {}", e))))?;

    let topic = gossip
        .subscribe(
            TopicId::from_bytes(string_to_32_byte_array(&TOPIC_NAME.to_string())),
            vec![],
        )
        .map_err(|e| RustlerError::Term(Box::new(format!("Gossip error: {:?}", e))))?;

    let (sender, _receiver) = topic.split();

    let state = NodeState::new(pid, endpoint_clone, router_clone, gossip, sender);
    let resource = ResourceArc::new(NodeRef(Arc::new(Mutex::new(state))));
    Ok(resource)
}

#[rustler::nif]
pub fn create_node(env: Env, pid: LocalPid) -> Result<ResourceArc<NodeRef>, RustlerError> {
    let endpoint = RUNTIME
        .block_on(
            Endpoint::builder()
                // .discovery_local_network()
                .discovery_n0()
                .bind(),
        )
        .map_err(|e| RustlerError::Term(Box::new(format!("Endpoint error: {}", e))))?;

    println!("Endpoint node id: {:?}", endpoint.node_id());

    let endpoint_clone = endpoint.clone();

    let mut builder = iroh::protocol::Router::builder(endpoint.clone());

    let gossip = RUNTIME
        .block_on(Gossip::builder().spawn(endpoint.clone()))
        .map_err(|e| RustlerError::Term(Box::new(format!("Gossip protocol error: {}", e))))?;

    builder = builder.accept(GossipALPN, gossip.clone());
    builder = builder.accept(ALPN, Echo);

    let router = RUNTIME
        .block_on(builder.spawn())
        .map_err(|e| RustlerError::Term(Box::new(format!("Router error: {}", e))))?;

    let router_clone = router.clone();

    let node_addr = RUNTIME
        .block_on(router.endpoint().node_addr())
        .map_err(|e| RustlerError::Term(Box::new(format!("Node addr error: {}", e))))?;

    let node_ids: Vec<_> = vec![];

    // Subscribe to the topic.
    // Since the `node_ids` list is empty, we will
    // subscribe to the topic, but not attempt to
    // connect to any other nodes.

    let topic = RUNTIME
        .block_on(async {
            gossip.subscribe(
                TopicId::from_bytes(string_to_32_byte_array(&TOPIC_NAME.to_string())),
                node_ids,
            )
        })
        .map_err(|e| RustlerError::Term(Box::new(format!("Gossip error: {:?}", e))))?;

    let (sender, _receiver) = topic.split();

    let state = NodeState::new(pid, endpoint_clone, router_clone, gossip, sender);
    let resource = ResourceArc::new(NodeRef(Arc::new(Mutex::new(state))));

    Ok(resource)
}

#[rustler::nif]
pub fn create_ticket(env: Env, node_ref: ResourceArc<NodeRef>) -> Result<String, RustlerError> {
    println!("Create ticket");

    let resource_arc = node_ref.0.clone();
    let state = resource_arc.lock().unwrap();

    let endpoint = { state.endpoint.clone() };

    let topic = TopicId::from_bytes(string_to_32_byte_array(&TOPIC_NAME.to_string()));

    let node_addr = RUNTIME
        .block_on(endpoint.node_addr())
        .map_err(|e| RustlerError::Term(Box::new(format!("Node addr error: {}", e))))?;

    let ticket = {
        // Get our address information, includes our
        // `NodeId`, our `RelayUrl`, and any direct
        // addresses.
        let me = node_addr;
        let nodes = vec![me];
        Ticket { topic, nodes }
    };

    Ok(ticket.to_string())
}

#[rustler::nif(schedule = "DirtyIo")]
fn gen_node_addr(node_ref: ResourceArc<NodeRef>) -> NifResult<String> {
    let resource_arc = node_ref.0.clone();
    let state = resource_arc.lock().unwrap();

    let endpoint = { state.endpoint.clone() };

    let node_id = endpoint.node_id();

    // let addr = node.local_peer_id().to_string();
    Ok(node_id.fmt_short())
}

#[rustler::nif]
pub fn send_message(
    env: Env,
    node_ref: ResourceArc<NodeRef>,
    message: String,
) -> Result<ResourceArc<NodeRef>, RustlerError> {
    // println!("Message: {:?}", message);

    let resource_arc = node_ref.0.clone();
    let state = resource_arc.lock().unwrap();

    let endpoint = { state.endpoint.clone() };

    let gossip = { state.gossip.clone() };

    let sender = { state.sender.clone() };

    let message = Message::AboutMe {
        from: endpoint.node_id(),
        name: message,
    };

    let result = RUNTIME.block_on(sender.broadcast(message.to_vec().into()));
    if let Err(e) = result {
        tracing::error!("Failed to send message: {:?}", e);
        return Err(RustlerError::Term(Box::new(e.to_string())));
    }

    Ok(node_ref)
}

#[rustler::nif]
pub fn connect_node(
    env: Env,
    node_ref: ResourceArc<NodeRef>,
    ticket: String,
) -> Result<ResourceArc<NodeRef>, RustlerError> {
    let node_ref_clone = node_ref.clone();

    RUNTIME.spawn(async move {
        if let Err(e) = connect_node_async_internal(node_ref_clone, ticket).await {
            tracing::error!("❌ Error in async task: {:?}", e);
        }
    });

    // Return immediately, allowing Elixir's Task to execute in parallel
    Ok(node_ref)
}

async fn connect_node_async_internal(node_ref: ResourceArc<NodeRef>, ticket: String) -> Result<()> {
    let resource_arc = node_ref.0.clone();

    let Ticket { topic, nodes } = Ticket::from_str(&ticket).context("❌ Failed to parse ticket")?;

    // let (pid_test, test) = {
    //     let state = resource_arc.lock().unwrap();
    //     (state.pid.clone, state.endpoint.clone())
    // };

    let pid = {
        let state = resource_arc.lock().unwrap();
        state.pid.clone()
    };

    let endpoint_conn = {
        let state = resource_arc.lock().unwrap();
        state.endpoint.clone()
    };

    let gossip = {
        let state = resource_arc.lock().unwrap();
        state.gossip.clone()
    };

    let node_id = {
        let state = resource_arc.lock().unwrap();
        state.endpoint.node_id()
    };

    let node_id_short: String = {
        let state = resource_arc.lock().unwrap();
        state.endpoint.node_id().fmt_short()
    };

    let endpoint_clone = endpoint_conn.clone();
    // RUNTIME.spawn(log_discovery_stream(endpoint_clone, pid.clone()));

    tracing::info!(
        "connect_node endpoint_Ptr:{:?} topic: {:?} nodes: {:?}",
        &endpoint_conn as *const _,
        topic,
        nodes
    );

    // avoid adding me, myself and I
    let nodes_filtered: Vec<_> = nodes
        .iter()
        .filter(|n| n.node_id != node_id)
        .filter(|n| n.node_id.fmt_short() != node_id_short)
        .collect();

    let node_ids: Vec<_> = nodes_filtered
        .iter()
        .map(|p| p.node_id)
        .collect::<HashSet<_>>()
        .into_iter()
        .collect();

    // let node_ids: Vec<_> = nodes.iter().map(|p| p.node_id).collect();
    if nodes.is_empty() {
        tracing::info!("Empty nodes list {:?}", nodes);
    } else {
        for node in nodes_filtered {
            tracing::info!("Adding node to addr book {:?}", node);
            if let Err(e) = endpoint_clone.add_node_addr(node.clone()) {
                tracing::error!("❌ Failed to add node to address book: {:?}", e);
            }
        }
    }

    let topic = gossip
        .subscribe_and_join(topic, node_ids)
        .await
        .context("❌ Failed to subscribe and join gossip")?;

    tracing::info!("Subscribed to: {:?}", topic);

    let (sender, mut receiver) = topic.split();
    // let (mpsc_event_sender, mpsc_event_receiver) = mpsc::channel::<Event>(100);
    // let mpsc_event_receiver_arc = Arc::new(RwLock::new(mpsc_event_receiver));
    // let mpsc_event_receiver_arc_clone = mpsc_event_receiver_arc.clone();

    RUNTIME.spawn(async move {
        // let mut names = HashMap::new();
        //let receiver = mpsc_event_receiver_arc_clone.read().await;
        // let mut msg_env = OwnedEnv::new();

        while let Ok(Some(event)) = receiver.try_next().await {
            // while let Some(event) = mpsc_event_receiver_arc_clone.write().await.recv().await {
            let mut msg_env = OwnedEnv::new();

            match event {
                Event::Gossip(GossipEvent::Joined(pub_keys)) => {
                    tracing::info!("Joined {:?} {:?}", node_id_short, pub_keys);
                    if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                        (
                            atoms::iroh_gossip_joined(),
                            pub_keys
                                .iter()
                                .map(|pk| pk.fmt_short()) // Apply formatting
                                .collect::<Vec<_>>() // Collect into a Vec<String>
                                .join(","),
                        )
                            .encode(env)
                    }) {
                        tracing::debug!("⚠️ Failed to send joined message: {:?}", e);
                    }
                }

                Event::Gossip(GossipEvent::NeighborUp(pub_key)) => {
                    tracing::info!("NeighborUp {:?}", pub_key);
                    if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                        (
                            atoms::iroh_gossip_neighbor_up(),
                            node_id_short.clone(),
                            pub_key.clone().fmt_short(),
                        )
                            .encode(env)
                    }) {
                        tracing::debug!("⚠️ Failed to send neighborup message: {:?}", e);
                    }
                }

                Event::Gossip(GossipEvent::NeighborDown(pub_key)) => {
                    tracing::info!("NeighborDown {:?}", pub_key);
                    if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                        (
                            atoms::iroh_gossip_neighbor_down(),
                            node_id_short.clone(),
                            pub_key.clone().fmt_short(),
                        )
                            .encode(env)
                    }) {
                        tracing::debug!("⚠️ Failed to send neighbordown message: {:?}", e);
                    }
                }

                Event::Gossip(GossipEvent::Received(msg)) => {
                    // if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                    //     (atoms::iroh_gossip_message_received(), "test3", "test4").encode(env)
                    // }) {
                    //     tracing::debug!("⚠️ Failed to send joined message: {:?}", e);
                    // }

                    match Message::from_bytes(&msg.content) {
                        Ok(message) => {
                            match message {
                                Message::AboutMe { from, name } => {
                                    // names.insert(from, name.clone());
                                    tracing::debug!("💬 FROM: {} MSG: {}", from.fmt_short(), name);

                                    // let mut msg_env = OwnedEnv::new();

                                    // if !pid.is_alive() {
                                    //     tracing::error!(
                                    //         "🚨 Elixir process (pid) is dead! Cannot send message."
                                    //     );
                                    // } else {
                                    //     tracing::info!(
                                    //         "✅ Elixir process (pid) is alive. Sending message..."
                                    //     );
                                    // }

                                    if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                                        (
                                            atoms::iroh_gossip_message_received(),
                                            node_id_short.clone(),
                                            name.clone(),
                                            // "test_static1",
                                            // "test_static2",
                                        )
                                            .encode(env)
                                    }) {
                                        tracing::debug!(
                                            "⚠️ Failed to forward received message: {:?}",
                                            e
                                        );
                                    }
                                }
                                Message::Message { from, text } => {
                                    // let name = names.get(&from).map_or_else(|| from.fmt_short(), String::to_string);
                                    // tracing::info!("📝 {}: {}", name, text);
                                    tracing::info!("📝 {}: {}", from, text);
                                }
                            }
                        }
                        Err(e) => {
                            tracing::error!("❌ Failed to parse message: {:?}", e);
                        }
                    }
                }
                _ => {
                    tracing::debug!("🔍 Ignored unhandled event: {:?}", event);
                    let message = format!("🔍 Ignored unhandled event: {:?}", event);
                    if let Err(e) = msg_env.send_and_clear(&pid, |env| {
                        (atoms::iroh_gossip_message_unhandled(), message).encode(env)
                    }) {
                        tracing::debug!("⚠️ Failed to send unhandled message: {:?}", e);
                    }
                }
            }
        }
    });

    Ok(())
}

#[rustler::nif(schedule = "DirtyIo")]
fn disconnect_node(node_ref: ResourceArc<NodeRef>) -> NifResult<()> {
    // let node = node_ref.lock().unwrap();
    // node.disconnect_all();  // Assuming an API to disconnect all peers
    Ok(())
}

#[rustler::nif(schedule = "DirtyIo")]
fn list_peers(node_ref: ResourceArc<NodeRef>) -> NifResult<Vec<String>> {
    let node = node_ref.0.clone();
    let state = node.lock().unwrap();

    let endpoint = { state.endpoint.clone() };

    //endpoint.discovery_stream();

    let peers: Vec<_> = vec![]; //node.peers().iter().map(|p| p.to_string()).collect();
    Ok(peers)
}

// The protocol definition:
#[derive(Debug, Clone)]
struct Echo;

impl ProtocolHandler for Echo {
    fn accept(&self, connection: Connection) -> BoxFuture<Result<()>> {
        Box::pin(async move {
            let (mut send, mut recv) = connection.accept_bi().await?;

            // Echo any bytes received back directly.
            let bytes_sent = tokio::io::copy(&mut recv, &mut send).await?;

            send.finish()?;
            connection.closed().await;

            Ok(())
        })
    }
}

async fn log_discovery_stream(endpoint: Endpoint, pid: LocalPid) {
    let mut stream = endpoint.discovery_stream();
    let mut msg_env = OwnedEnv::new();

    while let Some(result) = stream.next().await {
        match result {
            Ok(node_addr) => {
                match msg_env.send_and_clear(&pid, |env| {
                    (
                        atoms::iroh_gossip_node_discovered(),
                        endpoint.node_id().fmt_short(),
                        node_addr.node_id().fmt_short(),
                    )
                        .encode(env)
                }) {
                    Ok(_) => tracing::debug!("✅ Sent device discovery message"),
                    Err(e) => {
                        tracing::debug!("⚠️ Failed to send discovery message: {:?}", e)
                    }
                }

                tracing::info!(
                    "🔍 {:?} Discovered Node: {:?}",
                    endpoint.node_id().fmt_short(),
                    node_addr
                );
            }
            Err(lagged) => {
                tracing::warn!(
                    "🚨 {:?} Discovery stream lagged! Some items may have been lost.{:?}",
                    endpoint.node_id().fmt_short(),
                    lagged
                );
            }
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct Ticket {
    topic: TopicId,
    nodes: Vec<NodeAddr>,
}

impl Ticket {
    /// Deserialize from a slice of bytes to a Ticket.
    fn from_bytes(bytes: &[u8]) -> Result<Self> {
        serde_json::from_slice(bytes).map_err(Into::into)
    }

    /// Serialize from a `Ticket` to a `Vec` of bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        serde_json::to_vec(self).expect("serde_json::to_vec is infallible")
    }
}

// The `Display` trait allows us to use the `to_string`
// method on `Ticket`.
impl fmt::Display for Ticket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut text = data_encoding::BASE32_NOPAD.encode(&self.to_bytes()[..]);
        text.make_ascii_lowercase();
        write!(f, "{}", text)
    }
}

// The `FromStr` trait allows us to turn a `str` into
// a `Ticket`
impl FromStr for Ticket {
    type Err = anyhow::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let bytes = data_encoding::BASE32_NOPAD.decode(s.to_ascii_uppercase().as_bytes())?;
        Self::from_bytes(&bytes)
    }
}

#[rustler::nif]
fn add(a: i64, b: i64) -> i64 {
    a + b
}

// Rustler init

fn on_load(env: Env, _info: Term) -> bool {
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("iroh=error,iroh_ex=info"));

    let subscriber = tracing_subscriber::FmtSubscriber::builder()
        .with_env_filter(filter)
        .with_ansi(atty::is(atty::Stream::Stdout))
        .finish();

    tracing::subscriber::set_global_default(subscriber).expect("Failed to set up logging");

    tracing::debug!("Debug message from iroh_ex!");

    println!("Initializing Rust Iroh NIF module ...");
    rustler::resource!(NodeRef, env);
    println!("Rust NIF Iroh module loaded successfully.");
    true
}

rustler::init!("Elixir.IrohEx.Native", load = on_load);