pdsfs / src / main.rs

mod client;
mod error;
mod firehose;
mod fs;
mod resolver;

use atrium_api::{client::AtpServiceClient, com, types};
use atrium_common::resolver::Resolver;
use atrium_identity::identity_resolver::ResolvedIdentity;
use atrium_repo::{Repository, blockstore::CarStore};
use atrium_xrpc_client::isahc::IsahcClient;
use fuser::MountOption;
use futures::{StreamExt, stream};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use std::{
    io::{Cursor, Write},
    path::PathBuf,
    sync::Arc,
};

fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let matches = clap::command!()
        .arg(
            clap::Arg::new("handles")
                .index(1)
                .required(true)
                .num_args(1..)
                .help("One or more handles to download and mount"),
        )
        .arg(
            clap::Arg::new("mountpoint")
                .short('m')
                .action(clap::ArgAction::Set)
                .value_parser(clap::value_parser!(PathBuf)),
        )
        .get_matches();
    let handles = matches
        .get_many::<String>("handles")
        .unwrap()
        .cloned()
        .collect::<Vec<_>>();
    let mountpoint = matches
        .get_one::<PathBuf>("mountpoint")
        .map(ToOwned::to_owned)
        .unwrap_or(PathBuf::from("mnt"));
    let _ = std::fs::create_dir_all(&mountpoint);

    let resolver = Arc::new(resolver::id_resolver());
    let bars = Arc::new(MultiProgress::new());
    let repos = rt.block_on(
        stream::iter(handles)
            .then(|handle| {
                let h = handle.clone();
                let r = Arc::clone(&resolver);
                let b = Arc::clone(&bars);
                async move {
                    let id = r.resolve(&h).await?;
                    let bytes = cached_download(&id, &b).await?;
                    let repo = build_repo(bytes).await?;
                    Ok::<_, error::Error>((id.did.clone(), id.pds.clone(), repo))
                }
            })
            .collect::<Vec<_>>(),
    );
    let (success, errors): (Vec<_>, Vec<_>) = repos.into_iter().partition(|r| r.is_ok());
    for e in errors {
        eprintln!("{:?}", e.as_ref().unwrap_err());
    }
    let repos_with_pds: Vec<_> = success
        .into_iter()
        .map(|s| s.unwrap())
        .collect();

    // construct the fs
    let mut fs = fs::PdsFs::new();

    // Extract (did, pds) pairs for WebSocket tasks before consuming repos
    let did_pds_pairs: Vec<_> = repos_with_pds.iter()
        .map(|(did, pds, _)| (did.clone(), pds.clone()))
        .collect();

    // Consume repos_with_pds to add repos to filesystem
    for (did, _, repo) in repos_with_pds {
        rt.block_on(fs.add(did, repo))
    }

    // get shared state for WebSocket tasks
    let (_repos_arc, inodes_arc, sizes_arc, content_cache_arc) = fs.get_shared_state();

    // mount
    let options = vec![
        MountOption::RO,
        MountOption::FSName("pdsfs".to_string()),
        MountOption::AllowOther,
        MountOption::CUSTOM("local".to_string()),
        MountOption::CUSTOM("volname=pdsfs".to_string()),
    ];

    // Create session and get notifier for Finder refresh
    let session = fuser::Session::new(fs, &mountpoint, &options).unwrap();
    let notifier = session.notifier();
    let _bg = session.spawn().unwrap();

    // spawn WebSocket subscription tasks for each DID using the runtime handle
    let rt_handle = rt.handle().clone();
    for (did, pds) in did_pds_pairs {
        let inodes_clone = Arc::clone(&inodes_arc);
        let sizes_clone = Arc::clone(&sizes_arc);
        let content_cache_clone = Arc::clone(&content_cache_arc);
        let notifier_clone = notifier.clone();

        rt_handle.spawn(async move {
            if let Err(e) = firehose::subscribe_to_repo::<atrium_repo::blockstore::CarStore<std::io::Cursor<Vec<u8>>>>(
                did,
                pds,
                inodes_clone,
                sizes_clone,
                content_cache_clone,
                notifier_clone,
            ).await {
                eprintln!("WebSocket error: {:?}", e);
            }
        });
    }

    println!("mounted at {mountpoint:?}");
    print!("hit enter to unmount and exit...");
    std::io::stdout().flush().unwrap();

    // Wait for user input
    let mut input = String::new();
    std::io::stdin().read_line(&mut input).unwrap();

    println!("unmounted {mountpoint:?}");
}

async fn cached_download(
    id: &ResolvedIdentity,
    m: &MultiProgress,
) -> Result<Vec<u8>, error::Error> {
    let mut pb = ProgressBar::new_spinner();
    pb.set_style(
        ProgressStyle::default_spinner()
            .template("{spinner:.green} [{elapsed_precise}] {msg}")
            .unwrap()
            .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]),
    );
    pb.enable_steady_tick(std::time::Duration::from_millis(100));
    pb = m.add(pb);

    // Always download fresh - no caching for now to ensure up-to-date data
    pb.set_message(format!("downloading CAR file for...{}", id.did));
    let bytes = download_car_file(id, &pb).await?;

    pb.finish();
    Ok(bytes)
}

async fn download_car_file(
    id: &ResolvedIdentity,
    pb: &ProgressBar,
) -> Result<Vec<u8>, error::Error> {
    // download the entire car file first before mounting it as a fusefs
    let client = AtpServiceClient::new(IsahcClient::new(&id.pds));
    let did = types::string::Did::new(id.did.clone()).unwrap();

    let bytes = client
        .service
        .com
        .atproto
        .sync
        .get_repo(com::atproto::sync::get_repo::Parameters::from(
            com::atproto::sync::get_repo::ParametersData { did, since: None },
        ))
        .await?;

    pb.finish_with_message(format!("download complete for \t...\t{}", id.did));

    Ok(bytes)
}

async fn build_repo(bytes: Vec<u8>) -> Result<Repository<CarStore<Cursor<Vec<u8>>>>, error::Error> {
    let store = CarStore::open(Cursor::new(bytes)).await?;
    let root = store.roots().next().unwrap();
    let repo = Repository::open(store, root).await?;
    Ok(repo)
}