//! Blob storage for release assets. A thin wrapper over object_store's local filesystem backend,
//! so an S3 backend can replace it later without touching the pages.
//!
//! Keys are `assets/<asset id>`. Writes stream through object_store's multipart upload, which
//! assembles the object in a temporary file and renames it into place on `finish`. Every chunk
//! is hashed as it passes, so the sha256 is known without a second read. Blobs are immutable:
//! an asset is never rewritten, only added or deleted.

use std::path::Path;
use std::sync::Arc;

use bytes::Bytes;
use forge_core::{Error, Result};
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
use object_store::local::LocalFileSystem;
use object_store::path::Path as ObjectPath;
use object_store::{ObjectStore, ObjectStoreExt, WriteMultipart};
use sha2::{Digest, Sha256};

/// Bytes buffered per upload part before it is handed to the store.
const PART_BYTES: usize = 8 * 1024 * 1024;
/// Parts uploaded at once while a stream is being written.
const PARTS_IN_FLIGHT: usize = 4;

/// The object key of a release asset.
pub fn asset_key(asset_id: &str) -> String {
    format!("assets/{asset_id}")
}

fn storage_err<E>(e: E) -> Error
where
    E: std::error::Error + Send + Sync + 'static,
{
    Error::Io(std::io::Error::other(e))
}

/// The blob store rooted at `<data_dir>/blobs`. Cheap to clone.
#[derive(Clone)]
pub struct BlobStore {
    store: Arc<dyn ObjectStore>,
}

/// A blob opened for reading: its size and its bytes.
pub struct BlobRead {
    pub size: u64,
    pub stream: BoxStream<'static, std::io::Result<Bytes>>,
}

impl BlobStore {
    /// Opens the store at `root`, creating the directory if needed.
    pub fn open(root: &Path) -> Result<Self> {
        std::fs::create_dir_all(root)?;
        let fs = LocalFileSystem::new_with_prefix(root).map_err(storage_err)?;
        Ok(Self {
            store: Arc::new(fs),
        })
    }

    /// Starts a write of `key`. The bytes are not visible under `key` until `finish`.
    pub async fn begin(&self, key: &str, max_bytes: u64) -> Result<BlobWriter> {
        let upload = self
            .store
            .put_multipart(&ObjectPath::from(key))
            .await
            .map_err(storage_err)?;
        Ok(BlobWriter {
            writer: WriteMultipart::new_with_chunk_size(upload, PART_BYTES),
            hasher: Sha256::new(),
            size: 0,
            max_bytes,
        })
    }

    /// Opens `key` for reading, or `None` if there is no such blob.
    pub async fn read(&self, key: &str) -> Result<Option<BlobRead>> {
        match self.store.get(&ObjectPath::from(key)).await {
            Ok(result) => {
                let size = result.meta.size;
                let stream = result.into_stream().map_err(std::io::Error::other).boxed();
                Ok(Some(BlobRead { size, stream }))
            }
            Err(object_store::Error::NotFound { .. }) => Ok(None),
            Err(e) => Err(storage_err(e)),
        }
    }

    /// Moves `from` to `to`, replacing `to`. Used to give a finished upload its final key.
    pub async fn rename(&self, from: &str, to: &str) -> Result<()> {
        self.store
            .rename(&ObjectPath::from(from), &ObjectPath::from(to))
            .await
            .map_err(storage_err)
    }

    /// Deletes `key`. A blob that is already gone is not an error.
    pub async fn delete(&self, key: &str) -> Result<()> {
        match self.store.delete(&ObjectPath::from(key)).await {
            Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
            Err(e) => Err(storage_err(e)),
        }
    }
}

/// An upload in progress. Dropping it without `finish` leaves no object behind only if `abort`
/// is called, so every error path in the caller must call `abort`.
pub struct BlobWriter {
    writer: WriteMultipart,
    hasher: Sha256,
    size: u64,
    max_bytes: u64,
}

impl BlobWriter {
    /// Appends a chunk. Going past the size limit is `Invalid` and writes nothing more.
    pub async fn push(&mut self, chunk: &[u8]) -> Result<()> {
        let next = self.size.saturating_add(chunk.len() as u64);
        if next > self.max_bytes {
            return Err(Error::Invalid(
                "the file is larger than the upload limit".into(),
            ));
        }
        self.size = next;
        self.hasher.update(chunk);
        self.writer
            .wait_for_capacity(PARTS_IN_FLIGHT)
            .await
            .map_err(storage_err)?;
        self.writer.write(chunk);
        Ok(())
    }

    /// Completes the object. Returns its size and lowercase hex sha256.
    pub async fn finish(self) -> Result<(u64, String)> {
        let digest = hex::encode(self.hasher.finalize());
        self.writer.finish().await.map_err(storage_err)?;
        Ok((self.size, digest))
    }

    /// Discards the partial object. Errors are ignored: the object was never visible.
    pub async fn abort(self) {
        let _ = self.writer.abort().await;
    }
}