//! Blob store operations on top of the iroh-blobs persistent fs store.

use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};
use std::time::Duration;

use anyhow::Context;
use iroh_blobs::api::blobs::{AddPathOptions, BlobStatus, ImportMode};
use iroh_blobs::format::collection::Collection;
use iroh_blobs::hashseq::HashSeq;
use iroh_blobs::store::fs::options::{Options, PathOptions};
use iroh_blobs::store::fs::FsStore;
use iroh_blobs::store::{GcConfig, ProtectCb, ProtectOutcome};
use iroh_blobs::{BlobFormat, Hash, HashAndFormat};
use tracing::{debug, info, warn};
use varde_proto::MaterializeMode;

use crate::meta::StoredFormat;

/// Snapshot of the gc roots (the pinned hashes with their formats),
/// consulted by the store's gc before every run.
pub type GcRootsFn = Arc<dyn Fn() -> Vec<HashAndFormat> + Send + Sync>;

/// Operation errors that map onto protocol error codes.
#[derive(Debug, thiserror::Error)]
pub enum OpError {
    #[error("{0}")]
    NotFound(String),
    #[error("{0}")]
    InvalidArgument(String),
    #[error(transparent)]
    Io(#[from] std::io::Error),
    #[error(transparent)]
    Internal(#[from] anyhow::Error),
    #[error("{0}")]
    Unimplemented(String),
}

/// Wrap an iroh-blobs API error as an internal error.
fn internal(e: impl std::error::Error + Send + Sync + 'static) -> OpError {
    OpError::Internal(anyhow::Error::new(e))
}

/// The daemon's blob store: an iroh-blobs fs store rooted at
/// `<store_dir>/blobs`.
#[derive(Debug, Clone)]
pub struct BlobStore {
    store: FsStore,
    data_dir: PathBuf,
}

impl BlobStore {
    /// Open the store. `gc_roots` supplies the pinned hashes the built-in
    /// garbage collector must keep (temp tags and stored tags are
    /// protected by the store itself); their hashseq children are
    /// expanded here before every run.
    pub async fn open(
        store_dir: &Path,
        gc_interval: Duration,
        gc_roots: GcRootsFn,
    ) -> anyhow::Result<BlobStore> {
        let blobs_dir = store_dir.join("blobs");
        // The protect callback needs the store API to expand hashseq
        // roots, but it is installed before the store exists — hence the
        // cell, filled right after load.
        let store_cell: Arc<OnceLock<iroh_blobs::api::Store>> = Arc::new(OnceLock::new());
        let protect: ProtectCb = {
            let cell = store_cell.clone();
            Arc::new(move |live: &mut HashSet<Hash>| {
                let cell = cell.clone();
                let roots = gc_roots();
                // The callback future must be Sync; the store API's
                // futures are not, so the expansion runs on its own task.
                Box::pin(async move {
                    let expanded = tokio::spawn(async move {
                        let mut hashes: Vec<Hash> = Vec::new();
                        for HashAndFormat { hash, format } in roots {
                            hashes.push(hash);
                            if format.is_raw() {
                                continue;
                            }
                            let Some(store) = cell.get() else { continue };
                            match children_of(store, hash).await {
                                Ok(children) => hashes.extend(children),
                                Err(e) => {
                                    warn!(hash = %hash.to_hex(), error = %e, "gc protect: expanding hashseq");
                                    return None;
                                }
                            }
                        }
                        Some(hashes)
                    })
                    .await;
                    match expanded {
                        Ok(Some(hashes)) => {
                            live.extend(hashes);
                            ProtectOutcome::Continue
                        }
                        // Confused about liveness: skip this run rather
                        // than sweep blobs that should be protected.
                        Ok(None) => ProtectOutcome::Abort,
                        Err(e) => {
                            warn!(error = %e, "gc protect task failed");
                            ProtectOutcome::Abort
                        }
                    }
                })
            })
        };
        let options = Options {
            path: PathOptions::new(&blobs_dir),
            inline: Default::default(),
            batch: Default::default(),
            gc: Some(GcConfig {
                interval: gc_interval,
                add_protected: Some(protect),
            }),
        };
        let store = FsStore::load_with_opts(blobs_dir.join("blobs.db"), options)
            .await
            .with_context(|| format!("opening blob store at {}", blobs_dir.display()))?;
        let _ = store_cell.set((*store).clone());
        Ok(BlobStore {
            store,
            data_dir: blobs_dir.join("data"),
        })
    }

    /// Access to the underlying iroh-blobs store API (used by the
    /// transfer layer). `FsStore` derefs to the API handle.
    pub fn api(&self) -> &iroh_blobs::api::Store {
        &self.store
    }

    /// Import a file or directory. Returns the root hash, total imported
    /// bytes, and the resulting format (`Raw` for files, `HashSeq` for
    /// directories).
    pub async fn add_path(
        &self,
        path: &Path,
        recursive: bool,
    ) -> Result<(Hash, u64, StoredFormat), OpError> {
        if !path.is_absolute() {
            return Err(OpError::InvalidArgument(format!(
                "path must be absolute: {}",
                path.display()
            )));
        }
        let meta = std::fs::metadata(path)?;
        if meta.is_file() {
            let (tag, size) = self.import_one(path.to_owned()).await?;
            info!(hash = %tag.hash().to_hex(), size, "imported file");
            return Ok((tag.hash(), size, StoredFormat::Raw));
        }
        if !meta.is_dir() {
            return Err(OpError::InvalidArgument(format!(
                "not a file or directory: {}",
                path.display()
            )));
        }
        if !recursive {
            return Err(OpError::InvalidArgument(format!(
                "{} is a directory; pass recursive=true",
                path.display()
            )));
        }

        let mut files = Vec::new();
        collect_files(path, path, &mut files)?;
        if files.is_empty() {
            return Err(OpError::InvalidArgument(format!(
                "directory {} contains no files",
                path.display()
            )));
        }

        // Temp tags keep the children alive until the collection root is
        // stored (and this method's caller records/pins the root).
        let mut children = Vec::new();
        let mut total = 0u64;
        for (name, file_path) in files {
            let (tag, size) = self.import_one(file_path).await?;
            total += size;
            children.push((name, tag));
        }
        let collection: Collection = children
            .iter()
            .map(|(name, tag)| (name.clone(), tag.hash()))
            .collect();
        let root_tag = collection
            .store(&self.store)
            .await
            .context("storing collection")?;
        let root = root_tag.hash();
        info!(hash = %root.to_hex(), files = children.len(), total, "imported directory");
        Ok((root, total, StoredFormat::HashSeq))
    }

    async fn import_one(
        &self,
        path: PathBuf,
    ) -> Result<(iroh_blobs::api::TempTag, u64), OpError> {
        let tag = self
            .store
            .add_path_with_opts(AddPathOptions {
                path,
                format: BlobFormat::Raw,
                mode: ImportMode::Copy,
            })
            .temp_tag()
            .await
            .map_err(internal)?;
        let size = match self.store.blobs().status(tag.hash()).await.map_err(internal)? {
            BlobStatus::Complete { size } => size,
            _ => {
                return Err(OpError::Internal(anyhow::anyhow!(
                    "freshly imported blob {} is not complete",
                    tag.hash().to_hex()
                )))
            }
        };
        Ok((tag, size))
    }

    /// Export `hash` to `dest`. Returns bytes written and whether at least
    /// one file was reflinked.
    pub async fn materialize(
        &self,
        hash: Hash,
        format: StoredFormat,
        dest: &Path,
        mode: MaterializeMode,
    ) -> Result<(u64, bool), OpError> {
        if !dest.is_absolute() {
            return Err(OpError::InvalidArgument(format!(
                "dest must be absolute: {}",
                dest.display()
            )));
        }
        let allow_reflink = matches!(mode, MaterializeMode::ReflinkOrCopy);
        match format {
            StoredFormat::Raw => self.export_blob(hash, dest, allow_reflink).await,
            StoredFormat::HashSeq => {
                let collection = Collection::load(hash, self.api())
                    .await
                    .map_err(|e| OpError::NotFound(format!("loading collection: {e}")))?;
                let mut total = 0u64;
                let mut any_reflink = false;
                tokio::fs::create_dir_all(dest).await?;
                for (name, child) in collection.iter() {
                    let rel = sanitize_collection_name(name)?;
                    let (bytes, reflinked) = self
                        .export_blob(*child, &dest.join(rel), allow_reflink)
                        .await?;
                    total += bytes;
                    any_reflink |= reflinked;
                }
                Ok((total, any_reflink))
            }
        }
    }

    /// Export one blob. Tries `FICLONE` via reflink first (zero-copy on
    /// btrfs/XFS when store and dest share a filesystem), falls back to
    /// the store's export (a copy). Never hardlinks: store files must
    /// stay immutable.
    async fn export_blob(
        &self,
        hash: Hash,
        dest: &Path,
        allow_reflink: bool,
    ) -> Result<(u64, bool), OpError> {
        let size = match self.store.blobs().status(hash).await.map_err(internal)? {
            BlobStatus::Complete { size } => size,
            BlobStatus::Partial { .. } => {
                return Err(OpError::NotFound(format!(
                    "{} is only partially present",
                    hash.to_hex()
                )))
            }
            BlobStatus::NotFound => {
                return Err(OpError::NotFound(format!(
                    "{} is not in the store",
                    hash.to_hex()
                )))
            }
        };
        if let Some(parent) = dest.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        // Materialize replaces the destination, it never appends.
        match tokio::fs::remove_file(dest).await {
            Ok(()) => {}
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
            Err(e) => return Err(e.into()),
        }

        if allow_reflink {
            // Complete blobs above the inline threshold live as plain
            // files at data/<hex>.data; those we can clone directly.
            // Inline (small) blobs have no backing file and are copied.
            let src = self.data_dir.join(format!("{}.data", hash.to_hex()));
            let src_ok = std::fs::metadata(&src)
                .map(|m| m.len() == size)
                .unwrap_or(false);
            if src_ok {
                match reflink_copy::reflink(&src, dest) {
                    Ok(()) => {
                        debug!(hash = %hash.to_hex(), dest = %dest.display(), "reflinked");
                        return Ok((size, true));
                    }
                    Err(e) => {
                        debug!(error = %e, "reflink failed, falling back to copy");
                    }
                }
            }
        }

        let bytes = self
            .store
            .blobs()
            .export(hash, dest)
            .await
            .map_err(internal)?;
        Ok((bytes, false))
    }

    /// Presence of `hash`: (have_bytes, total_bytes, complete).
    pub async fn presence(&self, hash: Hash) -> Result<(u64, Option<u64>, bool), OpError> {
        match self.store.blobs().status(hash).await.map_err(internal)? {
            BlobStatus::NotFound => Ok((0, None, false)),
            BlobStatus::Partial { size } => {
                // Chunk-accurate accounting from the store's bitfield:
                // report the bytes actually present and verified.
                let bitfield = self.store.blobs().observe(hash).await.map_err(internal)?;
                let total = bitfield.validated_size().or(size);
                Ok((bitfield.total_bytes().min(total.unwrap_or(u64::MAX)), total, false))
            }
            BlobStatus::Complete { size } => Ok((size, Some(size), true)),
        }
    }

    /// The child hashes of a complete HashSeq root (empty for absent or
    /// partial roots).
    pub async fn hashseq_children(&self, root: Hash) -> Result<Vec<Hash>, OpError> {
        children_of(self.api(), root).await
    }

    /// Flush and shut down the store actor.
    pub async fn shutdown(&self) {
        if let Err(e) = self.store.shutdown().await {
            debug!(error = %e, "store shutdown");
        }
    }
}

/// The child hashes of a complete HashSeq root (empty for absent or
/// partial roots).
async fn children_of(store: &iroh_blobs::api::Store, root: Hash) -> Result<Vec<Hash>, OpError> {
    match store.blobs().status(root).await.map_err(internal)? {
        BlobStatus::Complete { .. } => {}
        _ => return Ok(Vec::new()),
    }
    let bytes = store.blobs().get_bytes(root).await.map_err(internal)?;
    let seq = HashSeq::try_from(bytes)
        .map_err(|e| OpError::Internal(anyhow::anyhow!("invalid hashseq: {e}")))?;
    Ok(seq.iter().collect())
}

/// Recursively collect regular files under `dir` as (relative-name, path),
/// sorted for deterministic collection hashes. Symlinks are followed for
/// files; symlinked directories are rejected to avoid cycles.
fn collect_files(root: &Path, dir: &Path, out: &mut Vec<(String, PathBuf)>) -> Result<(), OpError> {
    let mut entries: Vec<_> =
        std::fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
    entries.sort_by_key(|e| e.file_name());
    for entry in entries {
        let path = entry.path();
        let file_type = entry.file_type()?;
        if file_type.is_dir() {
            collect_files(root, &path, out)?;
        } else if file_type.is_symlink() && std::fs::metadata(&path)?.is_dir() {
            return Err(OpError::InvalidArgument(format!(
                "symlinked directory not supported: {}",
                path.display()
            )));
        } else {
            let rel = path
                .strip_prefix(root)
                .expect("walked path is under root")
                .components()
                .map(|c| c.as_os_str().to_string_lossy())
                .collect::<Vec<_>>()
                .join("/");
            out.push((rel, path));
        }
    }
    Ok(())
}

/// Turn a collection entry name into a safe relative path.
fn sanitize_collection_name(name: &str) -> Result<PathBuf, OpError> {
    let mut path = PathBuf::new();
    for comp in name.split('/') {
        if comp.is_empty() || comp == "." || comp == ".." {
            return Err(OpError::InvalidArgument(format!(
                "unsafe name in collection: {name:?}"
            )));
        }
        path.push(comp);
    }
    Ok(path)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn collection_names_are_sanitized() {
        assert!(sanitize_collection_name("a/b.txt").is_ok());
        assert!(sanitize_collection_name("../etc/passwd").is_err());
        assert!(sanitize_collection_name("/abs").is_err());
        assert!(sanitize_collection_name("a//b").is_err());
        assert!(sanitize_collection_name("a/./b").is_err());
    }
}