p2p_chat/app/

client.rs

//! This module contains the primary entry point for running the application in client mode.
use crate::cli::commands::{Node, UiNotification};
use crate::crypto::{Identity, StorageEncryption};
use crate::network::NetworkLayer;
use crate::storage::{
    MessageHistory, SeenTracker, SledFriendsStore, SledKnownMailboxesStore, SledOutboxStore,
    SledSeenTracker,
};
use crate::sync::{SyncEngine, SyncStores};
use crate::types::Message;
use crate::ui::run_tui;
use anyhow::Result;
use libp2p::Multiaddr;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, Mutex};
use tracing::{debug, error, info};

/// Runs the application in client mode.
///
/// This function initializes and runs all the components required for the client
/// to operate, including storage, networking, the synchronization engine,
/// the terminal UI, and the web server.
///
/// # Arguments
///
/// * `identity` - The user's identity.
/// * `db` - The database instance.
/// * `encryption` - The encryption key for the storage, if enabled.
/// * `port` - The port to listen on for P2P connections.
/// * `web_port` - The port for the Web UI.
///
/// # Errors
///
/// This function will return an error if any of the components fail to initialize
/// or run.
pub async fn run(
    identity: Arc<Identity>,
    db: sled::Db,
    encryption: Option<StorageEncryption>,
    port: u16,
    web_port: u16,
) -> Result<()> {
    println!("💬 Starting client mode");

    // Initialize storage components.
    let friends = Arc::new(SledFriendsStore::new(db.clone(), encryption.clone())?);
    let history = Arc::new(MessageHistory::new(db.clone(), encryption.clone())?);
    let outbox = Arc::new(SledOutboxStore::new(db.clone(), encryption.clone())?);
    let seen = Arc::new(SledSeenTracker::new(db.clone())?);
    let known_mailboxes = Arc::new(SledKnownMailboxesStore::new(
        db.clone(),
        encryption.clone(),
    )?);

    let listen_addr = Multiaddr::from_str(&format!("/ip4/0.0.0.0/tcp/{}", port))?;

    let bootstrap_nodes = vec![];

    // Initialize the network layer.
    let (mut network_layer, network_handle) =
        NetworkLayer::new(identity.clone(), listen_addr, false, bootstrap_nodes)?;

    // Create channels for communication between components.
    let (incoming_tx, mut incoming_rx) = mpsc::unbounded_channel::<Message>();
    let (ui_notify_tx, ui_notify_rx) = mpsc::unbounded_channel::<UiNotification>();
    let (web_notify_tx, web_notify_rx) = mpsc::unbounded_channel::<UiNotification>();
    let (network_notify_tx, mut network_notify_rx) = mpsc::unbounded_channel::<UiNotification>();

    let sync_stores = SyncStores::new(
        friends.clone(),
        outbox.clone(),
        history.clone(),
        seen.clone(),
        known_mailboxes.clone(),
    );

    // Initialize the synchronization engine.
    let (sync_engine_instance, sync_event_tx, mut sync_event_rx) = SyncEngine::new_with_network(
        Duration::from_secs(30),
        identity.clone(),
        sync_stores,
        network_handle.clone(),
        ui_notify_tx.clone(),
        Some(web_notify_tx.clone()),
    )?;
    let sync_engine = Arc::new(Mutex::new(sync_engine_instance));

    network_layer.set_sync_event_sender(sync_event_tx.clone());

    // Create the main application node context.
    let node = Arc::new(Node {
        identity,
        friends: friends.clone(),
        history: history.clone(),
        outbox: outbox.clone(),
        network: network_handle,
        ui_notify_tx,
        sync_engine: sync_engine.clone(),
    });

    println!("Client initialized. Starting network and TUI...\n");

    // Spawn the synchronization engine task.
    let sync_engine_clone = sync_engine.clone();
    tokio::spawn(async move {
        let interval_duration = {
            let engine = sync_engine_clone.lock().await;
            engine.interval
        };
        let mut interval_timer = tokio::time::interval(interval_duration);

        info!("Starting sync engine with interval {:?}", interval_duration);

        // Perform an initial discovery and sync cycle.
        {
            let mut engine = sync_engine_clone.lock().await;
            if let Err(e) = engine.initial_discovery().await {
                error!("Initial mailbox discovery failed: {}", e);
            }
            if let Err(e) = engine.sync_cycle().await {
                error!("Initial sync cycle failed: {}", e);
            }
        }

        // Main sync loop.
        loop {
            tokio::select! {
                _ = interval_timer.tick() => {
                    let mut engine = sync_engine_clone.lock().await;
                    if let Err(e) = engine.sync_cycle().await {
                        error!("Sync cycle failed: {}", e);
                    }
                }
                event = sync_event_rx.recv() => {
                    if let Some(event) = event {
                        let mut engine = sync_engine_clone.lock().await;
                        if let Err(e) = engine.handle_event(event).await {
                            error!("Failed to handle sync event: {}", e);
                        }
                    } else {
                        info!("Sync event channel closed, stopping engine.");
                        break;
                    }
                }
            }
        }
    });

    network_layer.set_sync_event_sender(sync_event_tx);
    network_layer.set_ui_notify_sender(network_notify_tx.clone());

    // Spawn the network layer task.
    tokio::spawn(async move {
        if let Err(e) = network_layer.run(incoming_tx).await {
            error!("Network layer error: {}", e);
        }
    });

    // Handle network notifications (e.g., delivery confirmations).
    let node_for_network = node.clone();
    let web_notify_tx_for_network = web_notify_tx.clone();
    tokio::spawn(async move {
        while let Some(notification) = network_notify_rx.recv().await {
            match notification {
                UiNotification::DeliveryStatusUpdate { message_id, new_status } => {
                    // Update the delivery status in the database.
                    if let Err(e) = node_for_network.history.update_delivery_status(&message_id, new_status).await {
                        error!("Failed to update delivery status in database: {}", e);
                    }

                    // Forward the notification to the web UI.
                    let _ = web_notify_tx_for_network.send(UiNotification::DeliveryStatusUpdate {
                        message_id,
                        new_status,
                    });
                }
                // Forward other notifications as-is.
                other => {
                    let _ = web_notify_tx_for_network.send(other);
                }
            }
        }
    });

    // Handle incoming messages from the network.
    let node_clone = node.clone();
    let seen_clone = seen.clone();
    let web_notify_tx_clone = web_notify_tx.clone();
    tokio::spawn(async move {
        while let Some(message) = incoming_rx.recv().await {
            let already_seen = match seen_clone.is_seen(&message.id).await {
                Ok(flag) => flag,
                Err(e) => {
                    error!("Failed to check if message {} was seen: {}", message.id, e);
                    false
                }
            };

            if already_seen {
                debug!("Received duplicate message {}, ignoring", message.id);
                continue;
            }

            if let Err(e) = node_clone.history.store_message(message.clone()).await {
                error!("Failed to store incoming message {}: {}", message.id, e);
                continue;
            }

            if let Err(e) = seen_clone.mark_seen(message.id).await {
                error!("Failed to mark message {} as seen: {}", message.id, e);
            }

            // Send a delivery confirmation back to the sender.
            let confirmation = crate::types::DeliveryConfirmation {
                original_message_id: message.id,
                timestamp: chrono::Utc::now().timestamp_millis(),
            };

            let confirmation_request = crate::types::ChatRequest::DeliveryConfirmation {
                confirmation,
            };

            let network_clone = node_clone.network.clone();
            let sender = message.sender;
            tokio::spawn(async move {
                if let Err(e) = network_clone.send_chat_request(sender, confirmation_request).await
                {
                    debug!("Failed to send delivery confirmation: {}", e);
                }
            });

            // Notify the UI and web server of the new message.
            let _ = node_clone
                .ui_notify_tx
                .send(UiNotification::NewMessage(message.clone()));
            let _ = web_notify_tx_clone.send(UiNotification::NewMessage(message));
        }
    });

    // Start the web server.
    let node_for_web = node.clone();
    tokio::spawn(async move {
        if let Err(e) = crate::web::start_server(node_for_web, web_port, web_notify_rx).await {
            error!("Web server error: {}", e);
        }
    });

    // Run the terminal UI.
    run_tui(node, ui_notify_rx, web_port).await
}