sneakerweb 1.2.0

use anyhow::{Context, Result, anyhow};
use clap::Args;
use smol::{
    fs::{File, OpenOptions},
    io::{AssertAsync, AsyncBufReadExt, AsyncRead, BufReader},
};
use std::{
    collections::BTreeSet,
    ffi::OsStr,
    io,
    path::{Path, PathBuf},
    pin::Pin,
};
use ufotofu::{
    Consumer, ConsumerExt, IntoConsumer, IntoProducer,
    consumer::compat::writer::writer_to_bulk_consumer, queues::new_unbounded_elastic,
};
use willow25::{
    drop_format::{DropEncoder, EncodeDropError, ExportDropError, export_drop},
    entry::NamespaceId,
    groupings::{Area, Keylike},
    prelude::AuthorisedEntry,
    storage::{PersistentStore, Store},
};

use crate::util::{
    SNEAKERWEB_NAMESPACE_ID_BYTES, parse_domain, sneakerweb_dir, warn_deprecated_domain_encoding,
};

#[derive(Args)]
pub struct ExportArgs {
    /// The filesystem path to write the .snk file to.
    pub dest: PathBuf,

    /// Include the specified domain in the exported .snk file. Can be passed multiple times.
    ///
    /// If this option is not specified, all domains are included.
    #[arg(short, long, group = "domains")]
    pub domain: Vec<String>,

    /// A path to a file listing domains to include in the exported .snk file.
    ///
    /// If no such file is specified, all domains are included. Pass '-' to read from stdin.
    #[arg(short, long, group = "domains")]
    pub exportlist: Option<PathBuf>,
    /// A path to the collection from which to export.
    ///
    /// If no such path is specified, sneakerweb domains are exported from the default
    /// collection (specified by the `DEFAULT_SNEAKERWEB_COLLECTION` environment variable,
    /// or ~/.sneakerweb if no default is set).
    #[arg(short, long)]
    pub collection: Option<PathBuf>,
}

pub async fn export_sneak(args: &ExportArgs) -> Result<PathBuf> {
    let mut dest = args.dest.clone();
    if dest.is_dir() {
        dest = dest.join("exported");
    }
    if dest.extension() != Some(OsStr::new("snk")) {
        dest.add_extension("snk");
    }
    let namespace = NamespaceId::from_bytes(&SNEAKERWEB_NAMESPACE_ID_BYTES);
    let sneakerweb_fs_path = sneakerweb_dir(args.collection.as_ref()).await?;
    let mut store = PersistentStore::new(sneakerweb_fs_path).await?;

    let destination = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(&dest)
        .await?;

    let drop_consumer = writer_to_bulk_consumer(destination, new_unbounded_elastic());

    let mut encoder = DropEncoder::new(drop_consumer);

    let included = match (args.domain.as_slice(), &args.exportlist) {
        (&[], Some(collection)) => {
            let mut included = BTreeSet::new();

            let mut collection_reader: BufReader<Pin<Box<dyn AsyncRead>>> =
                if collection == Path::new("-") {
                    BufReader::new(Box::pin(AssertAsync::new(io::stdin())))
                } else {
                    BufReader::new(Box::pin(
                        File::open(collection)
                            .await
                            .context(format!("could not open {}", collection.to_string_lossy()))?,
                    ))
                };

            let mut line = String::new();

            while let Ok(n) = collection_reader.read_line(&mut line).await
                && n > 0
            {
                // Otherwise, blank lines (including trailing newlines) will be treated as invalid
                // domains.
                if line.trim().is_empty() {
                    continue;
                }

                let (domain_id, _) = parse_domain(&line).context(format!(
                    "invalid domain '{}' found in collection '{}'",
                    line.trim(),
                    collection.to_string_lossy()
                ))?;
                warn_deprecated_domain_encoding(&line, &domain_id).await;
                included.insert(domain_id);
                line.clear();
            }
            Some(included)
        }
        (domains, None) if !domains.is_empty() => {
            let mut domain_ids = BTreeSet::new();
            for domain in domains {
                let (domain_id, _) = parse_domain(domain)?;
                warn_deprecated_domain_encoding(domain, &domain_id).await;
                domain_ids.insert(domain_id);
            }
            Some(domain_ids)
        }
        (&[], None) => None,
        _ => unreachable!(),
    };

    let mut entries: Vec<AuthorisedEntry> = vec![];

    store
        .get_area(
            &namespace,
            &Area::full(),
            &mut (&mut entries)
                .into_consumer()
                .to_filter(async |entry| match &included {
                    Some(collection) => collection.contains(entry.subspace_id()),
                    None => true,
                }),
        )
        .await?;

    export_drop(&mut store, &mut entries.into_producer(), &mut encoder)
        .await
        .map_err(|err| match err {
            ExportDropError::StoreError(err) => anyhow!(err),
            ExportDropError::ConsumerError(err) => match err {
                EncodeDropError::ConsumerError(err) => anyhow!(err),
                EncodeDropError::ConsumedBytesMismatch => {
                    anyhow!("the encoder received a different number of bytes than expected")
                }
                EncodeDropError::ArchitectureTooSmall => {
                    anyhow!("encountered an entry too large for this device")
                }
            },
            ExportDropError::EntryDeleted => {
                anyhow!("an entry was deleted concurrently with writing the export")
            }
        })?;

    encoder.flush().await.map_err(|err| match err {
        EncodeDropError::ConsumerError(err) => anyhow!(err),
        EncodeDropError::ConsumedBytesMismatch => {
            anyhow!("the encoder received a different number of bytes than expected")
        }
        EncodeDropError::ArchitectureTooSmall => {
            anyhow!("encountered an entry too large for this device")
        }
    })?;

    Ok(dest)
}