velocimeter/src/main.rs

use std::{
    fs::ReadDir,
    path::{Path, PathBuf},
    process::ExitCode,
};

use clap::{Parser, Subcommand};
use device::DeviceList;
use local::LocalServer;
use tracing::{debug, error, info, warn};

mod db;
mod device;
mod local;
mod web;

pub const CONFIG_DIR_NAME: &str = env!("CARGO_PKG_NAME");

#[derive(Parser)]
struct Args {
    #[command(subcommand)]
    cmd: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Lists saved devices
    ListSaved {
        /// Filter for device name
        name: Option<String>,
    },
    /// Deletes a saved device
    RmSaved {
        /// Name of the device to delete
        name: String,
    },
    /// Performs the sync
    Run {
        /// Don't show the QR code in the console
        #[arg(short, long)]
        no_qr_code: bool,
        /// Use saved device for transfer
        #[arg(short, long)]
        device: Option<String>,
        /// Directory to sync
        sync_dir: PathBuf,
    },
}

/// list-devices entrypoint
fn list_devices(list: &DeviceList, name: Option<String>) {
    if let Some(name) = name {
        if let Some(device) = list.find(name) {
            println!("{}", device);
        }
    } else {
        for device in list.iter() {
            println!("{}", device);
        }
    }
}

/// Wrapper function for appending children of a `ReadDir`
async fn append_children(paths: &mut Vec<PathBuf>, dir: ReadDir) -> std::io::Result<()> {
    let mut children = dir
        .map(|f| f.map(|f| f.path()))
        .collect::<std::io::Result<Vec<PathBuf>>>()?;
    paths.append(&mut children);
    Ok(())
}

/// Performs the actual sync process with the device.
async fn sync_dir(dir: impl AsRef<Path>, server: &mut LocalServer) -> bool {
    // To avoid recursing, which seems to be a mess with how we're using async,
    // we continue to add children to our vec as we run into them. So to start,
    // let's get all children of the "root" directory.
    let iter = match dir.as_ref().read_dir() {
        Ok(iter) => iter,
        Err(err) => {
            error!("couldn't read {}: {}", dir.as_ref().display(), err);
            return false;
        }
    };

    let mut paths: Vec<PathBuf> = Vec::new();
    let mut count = 0u32;

    // If we can't get all children of the root dir, assume something's wrong.
    if let Err(err) = append_children(&mut paths, iter).await {
        error!(
            "couldn't read paths from {}: {}",
            dir.as_ref().display(),
            err
        );
        return false;
    }

    // Loop to find all valid files
    while let Some(path) = paths.pop() {
        if path.is_dir() {
            match path.read_dir() {
                Ok(iter) => {
                    if let Err(err) = append_children(&mut paths, iter).await {
                        warn!("couldn't get paths from {}: {}", path.display(), err);
                    }
                }
                Err(err) => {
                    warn!("couldn't read {}: {}", path.display(), err);
                    continue;
                }
            };
        } else if path.exists() && server.should_upload(&path) {
            debug!("Adding {} to queue", path.display());
            // Add the download to the queue
            if let Err(err) = server.queue_upload(&path).await {
                warn!("couldn't send {}: {err}", path.display());
            } else {
                count += 1;
            }
        }
    }

    info!(
        "Processing {} {}.",
        count,
        if count == 1 { "song" } else { "songs" }
    );

    // Wait for the queue to complete, or for an error to occur
    if let Err(err) = server.wait_on_queue().await {
        error!("Error processing uploads: {err}");
        false
    } else {
        true
    }
}

async fn start_sync(
    list: &mut DeviceList,
    qr_code: bool,
    dir: PathBuf,
    device: Option<String>,
) -> ExitCode {
    if !dir.is_dir() {
        error!("can't open {} as a directory", dir.display());
        return ExitCode::FAILURE;
    }

    let mut web_conn = match web::connect().await {
        Ok(conn) => conn,
        Err(err) => {
            error!("unable to connect with Doppler web service: {err}");
            return ExitCode::FAILURE;
        }
    };

    if let Some(device) = device {
        if let Some(device) = list.find(&device) {
            if let Err(err) = web_conn.request_device(device).await {
                error!("requesting device {device} failed: {err}");
                return ExitCode::FAILURE;
            }
        } else {
            error!("device {device} not found in saved list");
            return ExitCode::FAILURE;
        }
    } else {
        if qr_code {
            let qrcode = qrencode::QrCode::new(web_conn.code()).unwrap();
            let encoded = qrcode.render::<char>().module_dimensions(2, 1).build();
            println!("{}", encoded);
        }

        println!("Use code {} to connect your device.", web_conn.code());
    }

    let (dev_id, dev_uri) = match web_conn.wait_for_device(list).await {
        Ok(uri) => uri,
        Err(err) => {
            error!("error getting device URI: {err}");
            return ExitCode::FAILURE;
        }
    };

    let cache = db::FileCache::open_for_device(dev_id).await;

    if let Err(err) = list.commit().await {
        warn!("can't save device: {err}");
    }

    debug!("Got device URI {}", &dev_uri);

    let mut local_server = match local::LocalServer::new(dev_uri.to_string(), cache).await {
        Ok(serv) => serv,
        Err(err) => {
            error!("couldn't connect to {dev_uri}: {err}");
            return ExitCode::FAILURE;
        }
    };

    if sync_dir(dir, &mut local_server).await {
        ExitCode::SUCCESS
    } else {
        ExitCode::FAILURE
    }
}

#[tokio::main]
async fn main() -> ExitCode {
    let mut device_list = match DeviceList::load().await {
        Ok(list) => list,
        Err(err) => {
            eprintln!("error loading saved device list: {err}");
            return ExitCode::FAILURE;
        }
    };

    match Args::parse().cmd {
        Command::ListSaved { name } => {
            list_devices(&device_list, name);
        }
        Command::RmSaved { name } => {
            device_list.drop_by_name(name);
        }
        Command::Run {
            no_qr_code,
            sync_dir,
            device,
        } => {
            tracing_subscriber::fmt().init();
            return start_sync(&mut device_list, !no_qr_code, sync_dir, device).await;
        }
    }

    ExitCode::SUCCESS
}