archy/core/archipelago/src/federation/storage.rs

//! Federation persistent storage: node list and invite management on disk.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;
use tokio::fs;

use super::types::{FederatedNode, FederationInvite, NodeStateSnapshot, TrustLevel};

pub(crate) const FEDERATION_DIR: &str = "federation";
pub(crate) const NODES_FILE: &str = "nodes.json";
pub(crate) const INVITES_FILE: &str = "invites.json";

/// Top-level file structures.
#[derive(Debug, Default, Serialize, Deserialize)]
pub(crate) struct NodesFile {
    pub(crate) nodes: Vec<FederatedNode>,
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub(crate) struct InvitesFile {
    pub(crate) outgoing: Vec<FederationInvite>,
    pub(crate) incoming: Vec<FederationInvite>,
}

/// Ensure federation directory exists.
pub(crate) async fn ensure_dir(data_dir: &Path) -> Result<std::path::PathBuf> {
    let dir = data_dir.join(FEDERATION_DIR);
    fs::create_dir_all(&dir)
        .await
        .context("Failed to create federation directory")?;
    Ok(dir)
}

// ──────────────────────────── Node Management ────────────────────────────

pub async fn load_nodes(data_dir: &Path) -> Result<Vec<FederatedNode>> {
    let dir = data_dir.join(FEDERATION_DIR);
    let path = dir.join(NODES_FILE);
    if !path.exists() {
        return Ok(Vec::new());
    }
    let content = fs::read_to_string(&path)
        .await
        .context("Failed to read federation nodes")?;
    let file: NodesFile = serde_json::from_str(&content).unwrap_or_default();
    Ok(file.nodes)
}

pub async fn save_nodes(data_dir: &Path, nodes: &[FederatedNode]) -> Result<()> {
    let dir = ensure_dir(data_dir).await?;
    let file = NodesFile {
        nodes: nodes.to_vec(),
    };
    let content = serde_json::to_string_pretty(&file).context("Failed to serialize nodes")?;
    fs::write(dir.join(NODES_FILE), content)
        .await
        .context("Failed to write federation nodes")?;
    Ok(())
}

pub async fn add_node(data_dir: &Path, node: FederatedNode) -> Result<Vec<FederatedNode>> {
    let mut nodes = load_nodes(data_dir).await?;
    let exists = nodes.iter().any(|n| n.did == node.did);
    if exists {
        anyhow::bail!("Node with DID {} is already federated", node.did);
    }
    nodes.push(node);
    save_nodes(data_dir, &nodes).await?;
    Ok(nodes)
}

pub async fn remove_node(data_dir: &Path, did: &str) -> Result<Vec<FederatedNode>> {
    let mut nodes = load_nodes(data_dir).await?;
    let before = nodes.len();
    nodes.retain(|n| n.did != did);
    if nodes.len() == before {
        anyhow::bail!("No federated node with DID {}", did);
    }
    save_nodes(data_dir, &nodes).await?;
    Ok(nodes)
}

pub async fn set_trust_level(
    data_dir: &Path,
    did: &str,
    trust: TrustLevel,
) -> Result<Vec<FederatedNode>> {
    let mut nodes = load_nodes(data_dir).await?;
    let node = nodes
        .iter_mut()
        .find(|n| n.did == did)
        .ok_or_else(|| anyhow::anyhow!("No federated node with DID {}", did))?;
    node.trust_level = trust;
    save_nodes(data_dir, &nodes).await?;
    Ok(nodes)
}

/// Update a federated node's metadata (onion, pubkey, name, last_seen).
pub async fn update_node(data_dir: &Path, updated: &FederatedNode) -> Result<()> {
    let mut nodes = load_nodes(data_dir).await?;
    if let Some(node) = nodes.iter_mut().find(|n| n.did == updated.did) {
        if !updated.onion.is_empty() {
            node.onion = updated.onion.clone();
        }
        if !updated.pubkey.is_empty() {
            node.pubkey = updated.pubkey.clone();
        }
        if updated.name.is_some() {
            node.name = updated.name.clone();
        }
        if updated.last_seen.is_some() {
            node.last_seen = updated.last_seen.clone();
        }
        save_nodes(data_dir, &nodes).await?;
    }
    Ok(())
}

pub async fn update_node_state(data_dir: &Path, did: &str, state: NodeStateSnapshot) -> Result<()> {
    let mut nodes = load_nodes(data_dir).await?;
    if let Some(node) = nodes.iter_mut().find(|n| n.did == did) {
        node.last_seen = Some(state.timestamp.clone());
        // Update node name from sync if provided (peer announced their name)
        if let Some(ref name) = state.node_name {
            if !name.is_empty() {
                node.name = Some(name.clone());
            }
        }
        node.last_state = Some(state);
        save_nodes(data_dir, &nodes).await?;
    }
    Ok(())
}

// ──────────────────────────── Invite Storage ────────────────────────────

pub(crate) async fn load_invites(data_dir: &Path) -> Result<InvitesFile> {
    let dir = data_dir.join(FEDERATION_DIR);
    let path = dir.join(INVITES_FILE);
    if !path.exists() {
        return Ok(InvitesFile::default());
    }
    let content = fs::read_to_string(&path)
        .await
        .context("Failed to read invites")?;
    let file: InvitesFile = serde_json::from_str(&content).unwrap_or_default();
    Ok(file)
}

pub(crate) async fn save_invites(data_dir: &Path, invites: &InvitesFile) -> Result<()> {
    let dir = ensure_dir(data_dir).await?;
    let content = serde_json::to_string_pretty(invites).context("Failed to serialize invites")?;
    fs::write(dir.join(INVITES_FILE), content)
        .await
        .context("Failed to write invites")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::federation::types::AppStatus;

    fn make_node(did: &str, onion: &str) -> FederatedNode {
        FederatedNode {
            did: did.to_string(),
            pubkey: "aabbccdd".to_string(),
            onion: onion.to_string(),
            name: None,
            trust_level: TrustLevel::Trusted,
            added_at: "2026-01-01T00:00:00Z".to_string(),
            last_seen: None,
            last_state: None,
            fips_npub: None,
        }
    }

    #[tokio::test]
    async fn test_load_nodes_empty_when_no_file() {
        let dir = tempfile::tempdir().unwrap();
        let nodes = load_nodes(dir.path()).await.unwrap();
        assert!(nodes.is_empty());
    }

    #[tokio::test]
    async fn test_save_and_load_nodes_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let nodes = vec![
            make_node("did:key:z1", "a.onion"),
            make_node("did:key:z2", "b.onion"),
        ];
        save_nodes(dir.path(), &nodes).await.unwrap();
        let loaded = load_nodes(dir.path()).await.unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded[0].did, "did:key:z1");
    }

    #[tokio::test]
    async fn test_add_node_deduplicates_by_did() {
        let dir = tempfile::tempdir().unwrap();
        add_node(dir.path(), make_node("did:key:z1", "a.onion"))
            .await
            .unwrap();
        let result = add_node(dir.path(), make_node("did:key:z1", "b.onion")).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_remove_node_by_did() {
        let dir = tempfile::tempdir().unwrap();
        add_node(dir.path(), make_node("did:key:z1", "a.onion"))
            .await
            .unwrap();
        add_node(dir.path(), make_node("did:key:z2", "b.onion"))
            .await
            .unwrap();
        let result = remove_node(dir.path(), "did:key:z1").await.unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].did, "did:key:z2");
    }

    #[tokio::test]
    async fn test_remove_nonexistent_node_errors() {
        let dir = tempfile::tempdir().unwrap();
        let result = remove_node(dir.path(), "did:key:nonexistent").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_set_trust_level() {
        let dir = tempfile::tempdir().unwrap();
        add_node(dir.path(), make_node("did:key:z1", "a.onion"))
            .await
            .unwrap();
        let nodes = set_trust_level(dir.path(), "did:key:z1", TrustLevel::Observer)
            .await
            .unwrap();
        assert_eq!(nodes[0].trust_level, TrustLevel::Observer);
    }

    #[tokio::test]
    async fn test_update_node_state() {
        let dir = tempfile::tempdir().unwrap();
        add_node(dir.path(), make_node("did:key:z1", "a.onion"))
            .await
            .unwrap();

        let state = NodeStateSnapshot {
            timestamp: "2026-03-10T12:00:00Z".to_string(),
            node_name: None,
            apps: vec![AppStatus {
                id: "bitcoin".to_string(),
                status: "running".to_string(),
                version: Some("27.0".to_string()),
            }],
            cpu_usage_percent: Some(45.2),
            mem_used_bytes: Some(4_000_000_000),
            mem_total_bytes: Some(8_000_000_000),
            disk_used_bytes: None,
            disk_total_bytes: None,
            uptime_secs: Some(86400),
            tor_active: Some(true),
            nostr_npub: None,
        };

        update_node_state(dir.path(), "did:key:z1", state)
            .await
            .unwrap();

        let nodes = load_nodes(dir.path()).await.unwrap();
        assert!(nodes[0].last_seen.is_some());
        let ls = nodes[0].last_state.as_ref().unwrap();
        assert_eq!(ls.apps.len(), 1);
        assert_eq!(ls.cpu_usage_percent, Some(45.2));
    }
}