sneakerweb 1.0.1

use anyhow::{Context, Error, Result, anyhow, bail};
use clap::{Args, ValueEnum};
use smol::{
    fs::File,
    io::{AsyncBufReadExt, BufReader},
};
use std::{collections::BTreeSet, path::PathBuf};
use ufotofu::{
    ExpectedFinalError, ProduceAtLeastError,
    codec::{Blame, DecodeError},
    prelude::*,
    producer::compat::{fn_mut::ClosureProducer, reader::reader_to_bulk_producer},
    queues::new_unbounded_elastic,
};
use willow25::{
    drop_format::{
        DropDecoder, DropSliceMetadata, ImportDropError, SliceStreamProducer, import_drop,
    },
    entry::NamespaceId,
    groupings::Keylike,
    storage::{PersistentStore, Store},
};

use crate::util::{SNEAKERWEB_NAMESPACE_ID_BYTES, domain_style, get_domain, hmm, sneakerweb_dir};

#[derive(Args)]
pub struct ImportArgs {
    /// The filesystem path of the .snk file to import from.
    pub src: PathBuf,
    /// The filtering mode to use when importing entries from the .snk file.
    pub mode: Option<ImportMode>,
}

#[derive(Clone, ValueEnum, Default)]
pub enum ImportMode {
    /// Import everything from the .snk file which you have not explicitly blocked. This is the default behaviour.
    #[default]
    All,
    /// Import only entries from the .snk file which correspond to domains you have already stored.
    Familiar,
}

pub async fn import_sneak(args: &ImportArgs) -> Result<()> {
    let file = File::open(&args.src).await?;
    let drop_producer = reader_to_bulk_producer(file, new_unbounded_elastic());
    let sneakerweb_fs_path = sneakerweb_dir().await?;
    let namespace_id = NamespaceId::from_bytes(&SNEAKERWEB_NAMESPACE_ID_BYTES);

    let mut store = PersistentStore::new(&sneakerweb_fs_path).await?;

    let allowed = match args.mode.clone().unwrap_or_default() {
        ImportMode::All => None,
        ImportMode::Familiar => {
            let mut allowed = BTreeSet::new();
            store
                .subspaces(&namespace_id, &mut (&mut allowed).into_consumer())
                .await?;
            Some(allowed)
        }
    };

    let mut decoder = DropDecoder::new(drop_producer);

    let blocklist_path = sneakerweb_fs_path.join("blocked");

    let mut blocked = BTreeSet::new();

    if blocklist_path.exists() {
        let mut blocklist = BufReader::new(
            File::open(blocklist_path)
                .await
                .context("a blocklist was found but could not be read")?,
        );
        let mut line = String::new();
        while let Ok(n) = blocklist.read_line(&mut line).await
            && n > 0
        {
            let (domain_id, _) = get_domain(Some(&line), "N/A").context(format!(
                "invalid domain '{}' found in blocklist",
                &line.trim()
            ))?;
            blocked.insert(domain_id);
            line.clear();
        }
    }

    let mut skipped = BTreeSet::new();

    let mut filtering_decoder =
        ClosureProducer::<_, (DropSliceMetadata, SliceStreamProducer<_>), (), Error>::new(
            async || {
                loop {
                    let Left((metadata, mut slice_stream)) =
                        decoder.produce().await.map_err(map_decode_err)?
                    else {
                        return Ok(Right(()));
                    };

                    let mut skip_reason = "blocked";
                    let domain_id = metadata.entry().subspace_id();

                    if !blocked.contains(domain_id) {
                        match &allowed {
                            None => return Ok(Left((metadata, slice_stream))),
                            Some(known) => {
                                if known.contains(domain_id) {
                                    return Ok(Left((metadata, slice_stream)));
                                }
                                skip_reason = "unfamiliar";
                            }
                        }
                    }

                    if !skipped.contains(domain_id) {
                        hmm(&format!(
                            "skipping content from {skip_reason} domain {}",
                            domain_style(&base16::encode_lower(domain_id.as_bytes()))
                        ))
                        .await;
                    }

                    skipped.insert(domain_id.clone());

                    slice_stream
                        .skip(
                            metadata.expected_bytes().try_into().context(
                                "an entry in the .snk file was too large for this device",
                            )?,
                        )
                        .await
                        .map_err(|ProduceAtLeastError { count: _, reason }| match reason {
                            Ok(()) => anyhow!("unexpected end of .snk file"),
                            Err(err) => map_decode_err(err),
                        })?;
                    slice_stream
                        .produce_final()
                        .await
                        .map_err(|err| match err {
                            ExpectedFinalError::Item(_) => {
                                anyhow!("an entry in the .snk file had incorrect metadata")
                            }
                            ExpectedFinalError::Error(err) => map_decode_err(err),
                        })?;
                }
            },
        );

    import_drop(&mut store, &mut filtering_decoder, true)
        .await
        .map_err(|err| match err {
            ImportDropError::SliceStreamBytesMismatch => {
                anyhow!("an entry in the .snk file had incorrect metadata")
            }
            ImportDropError::ImportProducerError(err) => err,
            ImportDropError::SliceProducerError(err) => map_decode_err(err),
            ImportDropError::StoreError(err) => err.into(),
            ImportDropError::VerificationError => {
                anyhow!("an entry in the .snk file could not be verified")
            }
            ImportDropError::EntryDeleted => {
                anyhow!("an entry was deleted concurrently with being imported")
            }
            ImportDropError::ArchitectureTooSmall => {
                anyhow!("an entry in the .snk file was too large for this device")
            }
            ImportDropError::PartialPayloadSliceEncountered => {
                unreachable!("we are not treating incompatible slices as fatal errors")
            }
        })?;

    match store.flush().await {
        Ok(_) => Ok(()),
        Err(_) => bail!("failed to flush changes to the sneakerweb store"),
    }
}

fn map_decode_err<E, P: Into<Error>>(err: DecodeError<E, P, Blame>) -> Error {
    match err {
        DecodeError::UnexpectedEndOfInput(_) => {
            anyhow!("unexpected end of .snk file")
        }
        DecodeError::ProducerError(err) => anyhow!(err),
        DecodeError::Other(blame) => match blame {
            Blame::TheirFault => {
                anyhow!("malformed encoding of entry in .snk file")
            }
            Blame::OurFault => anyhow!("failed to decode entry from .snk file"),
        },
    }
}