nolgia 0.1.0

//! Hand-written. Uploading files into the library and downloading results.

use std::path::{Path, PathBuf};

use futures_util::StreamExt;
use reqwest::{Url, header::CONTENT_TYPE};
use tokio::io::AsyncWriteExt;
use uuid::Uuid;

use crate::{
    Client, Typed,
    api::types,
    error::{Error, Result},
};

/// Options for [`Client::upload_file`] and [`Client::upload_bytes`].
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct UploadOptions {
    /// The MIME type. Guessed from the file extension when absent.
    pub content_type: Option<String>,
    /// Tags for the asset (at most 10).
    pub tags: Vec<String>,
    /// A display name for the asset.
    pub display_name: Option<String>,
    /// File it into this project instead of the default Library project.
    pub project_id: Option<Uuid>,
}

impl UploadOptions {
    /// Set the MIME type.
    pub fn content_type(mut self, content_type: impl Into<String>) -> Self {
        self.content_type = Some(content_type.into());
        self
    }

    /// Add a tag.
    pub fn tag(mut self, tag: impl Into<String>) -> Self {
        self.tags.push(tag.into());
        self
    }

    /// Set the display name.
    pub fn display_name(mut self, name: impl Into<String>) -> Self {
        self.display_name = Some(name.into());
        self
    }

    /// File it into a project.
    pub fn project_id(mut self, project: Uuid) -> Self {
        self.project_id = Some(project);
        self
    }
}

/// The MIME type the upload endpoint expects for a file name, from its
/// extension. `None` for a type the API does not accept.
pub fn content_type_for(path: impl AsRef<Path>) -> Option<&'static str> {
    let ext = path.as_ref().extension()?.to_str()?.to_ascii_lowercase();
    Some(match ext.as_str() {
        "png" => "image/png",
        "jpg" | "jpeg" => "image/jpeg",
        "webp" => "image/webp",
        "mp4" | "m4v" => "video/mp4",
        "mov" => "video/quicktime",
        "webm" => "video/webm",
        "mp3" => "audio/mpeg",
        "wav" => "audio/wav",
        "ogg" | "oga" => "audio/ogg",
        "m4a" => "audio/mp4",
        "glb" => "model/gltf-binary",
        _ => return None,
    })
}

impl Client {
    /// Upload a file into your library and return the ready asset.
    ///
    /// Uses the signed-upload flow (`POST /assets/uploads`, a `PUT` of the
    /// bytes straight to storage, then `POST /assets/uploads/{id}/complete`),
    /// which takes video up to 50 GiB, audio up to 5 GiB, images up to
    /// 100 MiB and 3D models up to 200 MiB. The file streams from disk.
    pub async fn upload_file(
        &self,
        path: impl AsRef<Path>,
        options: UploadOptions,
    ) -> Result<types::Asset> {
        let path = path.as_ref();
        let io = |source| Error::Io {
            path: path.to_path_buf(),
            source,
        };
        let size = tokio::fs::metadata(path).await.map_err(io)?.len();
        let filename = path
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("upload")
            .to_owned();
        let content_type = match &options.content_type {
            Some(ct) => ct.clone(),
            None => content_type_for(path)
                .ok_or_else(|| {
                    Error::InvalidRequest(format!(
                        "cannot tell the content type of {}; set UploadOptions::content_type",
                        path.display()
                    ))
                })?
                .to_owned(),
        };
        let upload = self
            .create_upload(&filename, &content_type, size, &options)
            .await?;
        let file = tokio::fs::File::open(path).await.map_err(io)?;
        let body = reqwest::Body::wrap_stream(tokio_util_stream(file));
        self.put_upload(&upload, &content_type, size, body).await?;
        self.complete_upload(upload.upload_id).await
    }

    /// Upload bytes from memory into your library and return the ready asset.
    pub async fn upload_bytes(
        &self,
        bytes: impl Into<Vec<u8>>,
        filename: impl Into<String>,
        options: UploadOptions,
    ) -> Result<types::Asset> {
        let bytes = bytes.into();
        let filename = filename.into();
        let content_type = match &options.content_type {
            Some(ct) => ct.clone(),
            None => content_type_for(&filename)
                .ok_or_else(|| {
                    Error::InvalidRequest(format!(
                        "cannot tell the content type of {filename}; set UploadOptions::content_type"
                    ))
                })?
                .to_owned(),
        };
        let size = bytes.len() as u64;
        let upload = self
            .create_upload(&filename, &content_type, size, &options)
            .await?;
        self.put_upload(&upload, &content_type, size, reqwest::Body::from(bytes))
            .await?;
        self.complete_upload(upload.upload_id).await
    }

    async fn create_upload(
        &self,
        filename: &str,
        content_type: &str,
        size: u64,
        options: &UploadOptions,
    ) -> Result<types::AssetUpload> {
        let mut body = serde_json::json!({
            "filename": filename,
            "content_type": content_type,
            "size_bytes": size,
        });
        if !options.tags.is_empty() {
            body["tags"] = serde_json::json!(options.tags);
        }
        if let Some(name) = &options.display_name {
            body["display_name"] = serde_json::json!(name);
        }
        if let Some(project) = options.project_id {
            body["project_id"] = serde_json::json!(project);
        }
        self.send_json(
            self.http().post(self.url("/assets/uploads")).json(&body),
            "create_asset_upload",
        )
        .await
    }

    async fn put_upload(
        &self,
        upload: &types::AssetUpload,
        content_type: &str,
        size: u64,
        body: reqwest::Body,
    ) -> Result<()> {
        // The signed URL is the credential; never send the API key to storage.
        let response = self
            .anon
            .put(&upload.upload_url)
            .header(CONTENT_TYPE, content_type)
            .header(reqwest::header::CONTENT_LENGTH, size)
            .body(body)
            .send()
            .await
            .map_err(Error::transport)?;
        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }
        Ok(())
    }

    async fn complete_upload(&self, upload_id: Uuid) -> Result<types::Asset> {
        self.complete_asset_upload()
            .id(upload_id)
            .send()
            .typed()
            .await
    }

    /// Save what is at `url` (an asset's `signed_url`, a [`Media::url`](crate::Media))
    /// to `path`, returning the bytes written.
    ///
    /// The body streams through `<path>.part`, renamed into place only after
    /// the last byte, so an interrupted download never leaves a truncated
    /// file. The API key is sent only when `url` is on the API's own origin.
    /// A signed URL that has expired answers `403`: fetch the asset again.
    pub async fn download(&self, url: &str, path: impl AsRef<Path>) -> Result<u64> {
        let path = path.as_ref();
        let response = self.fetch(url).await?;
        let part = part_path(path);
        let io = |p: &Path| {
            let p = p.to_path_buf();
            move |source| Error::Io { path: p, source }
        };
        let mut file = tokio::fs::File::create(&part).await.map_err(io(&part))?;
        let mut written = 0u64;
        let mut body = response.bytes_stream();
        let outcome: Result<()> = async {
            while let Some(chunk) = body.next().await {
                let chunk = chunk.map_err(Error::transport)?;
                file.write_all(&chunk).await.map_err(io(&part))?;
                written += chunk.len() as u64;
            }
            file.flush().await.map_err(io(&part))?;
            Ok(())
        }
        .await;
        drop(file);
        if let Err(err) = outcome {
            let _ = tokio::fs::remove_file(&part).await;
            return Err(err);
        }
        tokio::fs::rename(&part, path).await.map_err(io(path))?;
        Ok(written)
    }

    /// [`download`](Self::download) into memory.
    pub async fn download_bytes(&self, url: &str) -> Result<Vec<u8>> {
        let response = self.fetch(url).await?;
        Ok(response.bytes().await.map_err(Error::transport)?.to_vec())
    }

    async fn fetch(&self, url: &str) -> Result<reqwest::Response> {
        let parsed = Url::parse(url)
            .map_err(|_| Error::InvalidRequest(format!("not a URL: {}", redact(url))))?;
        let same_origin = Url::parse(self.base_url())
            .map(|base| base.origin() == parsed.origin())
            .unwrap_or(false);
        if same_origin {
            return self.send(self.http().get(parsed), "download").await;
        }
        let response = self
            .anon
            .get(parsed)
            .send()
            .await
            .map_err(Error::transport)?;
        if !response.status().is_success() {
            return Err(Error::from_response(response).await);
        }
        Ok(response)
    }
}

fn part_path(path: &Path) -> PathBuf {
    let mut name = path
        .file_name()
        .map(|n| n.to_os_string())
        .unwrap_or_default();
    name.push(".part");
    path.with_file_name(name)
}

/// The URL without its query string (signed URLs carry credentials there).
fn redact(url: &str) -> &str {
    url.split('?').next().unwrap_or(url)
}

/// Stream a file in 64 KiB chunks without pulling in tokio-util.
fn tokio_util_stream(
    file: tokio::fs::File,
) -> impl futures_util::Stream<Item = std::io::Result<bytes::Bytes>> + Send + 'static {
    use tokio::io::AsyncReadExt;
    futures_util::stream::try_unfold(file, |mut file| async move {
        let mut buffer = vec![0u8; 64 * 1024];
        let read = file.read(&mut buffer).await?;
        if read == 0 {
            return Ok(None);
        }
        buffer.truncate(read);
        Ok(Some((bytes::Bytes::from(buffer), file)))
    })
}

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

    #[test]
    fn content_types_follow_the_upload_contract() {
        assert_eq!(content_type_for("a/b/clip.MOV"), Some("video/quicktime"));
        assert_eq!(content_type_for("x.jpeg"), Some("image/jpeg"));
        assert_eq!(content_type_for("x.glb"), Some("model/gltf-binary"));
        assert_eq!(content_type_for("x.txt"), None);
    }

    #[test]
    fn part_file_sits_beside_the_target() {
        assert_eq!(
            part_path(Path::new("/tmp/a.png")),
            PathBuf::from("/tmp/a.png.part")
        );
        assert_eq!(redact("https://s/x.png?sig=secret"), "https://s/x.png");
    }
}