simploxide-client 0.14.0

//! Image previews generation

use base64::prelude::*;
#[cfg(feature = "native_crypto")]
use simploxide_api_types::CryptoFile;
use tokio::io::{AsyncReadExt as _, AsyncSeekExt as _};

use crate::util;

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

const DEFAULT_PREVIEW: &str = "data:image/jpg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/\
2wBDABALDA4MChAODQ4SERATGCgaGBYWGDEjJR0oOjM9PDkzODdASFxOQERXRTc4UG1RV19iZ2hnPk1xeXBkeFxlZ2P/\
2wBDARESEhgVGC8aGi9jQjhCY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2NjY2P/wAARCABVAIADASIAAhEBAxEB/\
8QAFgABAQEAAAAAAAAAAAAAAAAAAAEE/8QAFBABAAAAAAAAAAAAAAAAAAAAAP/EABgBAQEBAQEAAAAAAAAAAAAAAAMCAQUE/8QAFhEBAQEAAAAAAAAAAAAAAAAAAAER/\
9oADAMBAAIRAxEAPwDaKF17qgo3UVBRWjqCjdFUFFaOoKN0VQUVo6go3R0FHi13ago3R1BRWjoijdHUFFSiqCjZR1BRUo6go3RVQHh13aAK1FAVuiqCipR1BRUo6go2UV\
QUVKOoKK0dAHg13aCjdHUFFaOoKKlHUFFSiqCipR1BRsoqgoqUdBR4Nd2oKNlHUUFSjoAqUdAFSjoAqUVAFSjoAqUdAHPd2gCoOqAqDoA2DoAuCoAqDoAqCoAqDr//2Q==";

const MAX_PREVIEW_BYTES: usize = 9350;
#[cfg(feature = "multimedia")]
const MAX_FILE_SIZE: usize = 20 * 1024 * 1024;

/// Thumbnail for [`Image`](crate::messages::Image), [`Video`](crate::messages::Video), and
/// [`Link`](crate::messages::Link) messages. Also used as bot profile pictures. The source is stored
/// lazily and resolved when [`resolve`](Self::resolve) or [`try_resolve`](Self::try_resolve) is
/// called(either manually or automatically by message builders). Any error falls back to a default
/// ~600 bytes in size JPEG placeholder.
#[derive(Clone)]
pub struct ImagePreview {
    source: PreviewSource,
    #[cfg(feature = "multimedia")]
    transcoder: Transcoder,
}

impl Default for ImagePreview {
    fn default() -> Self {
        Self {
            source: PreviewSource::Default,
            #[cfg(feature = "multimedia")]
            transcoder: Transcoder::thumbnail(),
        }
    }
}

impl std::fmt::Debug for ImagePreview {
    #[cfg(not(feature = "multimedia"))]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ImagePreview")
            .field("source", &self.kind())
            .finish()
    }

    #[cfg(feature = "multimedia")]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ImagePreview")
            .field("source", &self.kind())
            .field("transcoder", &self.transcoder)
            .finish()
    }
}

impl ImagePreview {
    /// Thumbnail from raw JPEG bytes. Fails on resolve if the encoded data URI exceeds
    /// [`MAX_PREVIEW_BYTES`].
    pub fn from_bytes(bytes: impl Into<Vec<u8>>) -> Self {
        Self {
            source: PreviewSource::Bytes(bytes.into()),
            #[cfg(feature = "multimedia")]
            transcoder: Transcoder::thumbnail(),
        }
    }

    /// Thumbnail from a pre-assembled `data:image/jpg;base64,{base64_contents} URI string.
    pub fn raw(uri: impl Into<String>) -> Self {
        Self {
            source: PreviewSource::DataUri(uri.into()),
            #[cfg(feature = "multimedia")]
            transcoder: Transcoder::thumbnail(),
        }
    }

    /// Thumbnail loaded from a file; the file is read lazily when resolved.
    pub fn from_file(path: impl AsRef<Path>) -> Self {
        Self {
            source: PreviewSource::File(path.as_ref().to_path_buf()),
            #[cfg(feature = "multimedia")]
            transcoder: Transcoder::thumbnail(),
        }
    }

    pub fn kind(&self) -> PreviewKind {
        match self.source {
            PreviewSource::Default => PreviewKind::Default,
            PreviewSource::Bytes(_) => PreviewKind::Bytes,
            PreviewSource::DataUri(_) => PreviewKind::Raw,
            PreviewSource::File(_) => PreviewKind::File,
            #[cfg(feature = "native_crypto")]
            PreviewSource::CryptoFile(_) => PreviewKind::CryptoFile,
        }
    }

    #[cfg(feature = "native_crypto")]
    /// Thumbnail loaded from an encrypted file; decrypted lazily when resolved.
    pub fn from_crypto_file(file: CryptoFile) -> Self {
        Self {
            source: PreviewSource::CryptoFile(file),
            #[cfg(feature = "multimedia")]
            transcoder: Transcoder::thumbnail(),
        }
    }

    #[cfg(feature = "multimedia")]
    /// Attach a custom [`Transcoder`] to transcode the source as a JPEG thumbnail on resolve.
    /// Transcoder transcodes images of any widespread format to JPEGs.
    ///
    /// Has no effect on `default` and `raw` sources, they always passed as is.
    pub fn with_transcoder(mut self, transcoder: Transcoder) -> Self {
        self.set_transcoder(transcoder);
        self
    }

    #[cfg(feature = "multimedia")]
    pub fn set_transcoder(&mut self, transcoder: Transcoder) {
        self.transcoder = transcoder;
    }

    /// Like [`Self::try_resolve`] but falls back to the default placeholder preview on error.
    pub async fn resolve(self) -> String {
        match self.try_resolve().await {
            Ok(s) => s,
            Err(e) => {
                log::warn!("Falling back to default preview due to an error: {e}");
                default()
            }
        }
    }

    #[cfg(not(feature = "multimedia"))]
    /// Returns the preview as a `data:image/jpg;base64,{base64_contents}` URI. The source is
    /// assumed to be a valid JPEG(encoding is not validated) when multimedia feature is disabled or it is
    /// lazily transcoded into JPEG when multimedia is enabled. Fails if the source cannot be read
    /// or the encoded URI exceeds [`MAX_PREVIEW_BYTES`] bytes.
    pub async fn try_resolve(self) -> Result<String, PreviewError> {
        match self.source {
            PreviewSource::Default => Ok(default()),
            PreviewSource::Bytes(b) => try_encode_jpg_to_uri(&b),
            PreviewSource::DataUri(s) => validate_uri_preview(s),
            PreviewSource::File(path) => {
                let bytes = read_plain_file(&path, MAX_PREVIEW_BYTES).await?;
                try_encode_jpg_to_uri(&bytes)
            }
            #[cfg(feature = "native_crypto")]
            PreviewSource::CryptoFile(file) => {
                let bytes = read_crypto_file(file, MAX_PREVIEW_BYTES).await?;
                try_encode_jpg_to_uri(&bytes)
            }
        }
    }

    #[cfg(feature = "multimedia")]
    /// Returns the preview as a `data:image/jpg;base64,{base64_contents}` URI. The source is
    /// assumed to be a valid JPEG(encoding is not validated) when multimedia feature is disabled or it is
    /// lazily transcoded into JPEG when multimedia is enabled. Fails if the source cannot be read
    /// or the encoded URI exceeds [`MAX_PREVIEW_BYTES`] bytes.
    pub async fn try_resolve(self) -> Result<String, PreviewError> {
        let bytes = match self.source {
            PreviewSource::Default => return Ok(default()),
            PreviewSource::Bytes(b) => b,
            PreviewSource::DataUri(s) => return validate_uri_preview(s),
            PreviewSource::File(path) => read_plain_file(&path, MAX_FILE_SIZE).await?,
            #[cfg(feature = "native_crypto")]
            PreviewSource::CryptoFile(file) => read_crypto_file(file, MAX_FILE_SIZE).await?,
        };

        if self.transcoder.is_enabled() {
            let jpg_bytes =
                tokio::task::spawn_blocking(move || self.transcoder.transcode_to_jpg(bytes))
                    .await??;
            Ok(encode_to_uri(&jpg_bytes))
        } else {
            try_encode_jpg_to_uri(&bytes)
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PreviewKind {
    Default,
    Bytes,
    Raw,
    File,
    #[cfg(feature = "native_crypto")]
    CryptoFile,
}

#[cfg(feature = "multimedia")]
pub mod transcoder {
    use image::{
        ImageFormat, ImageReader,
        codecs::jpeg::JpegEncoder,
        imageops::{self, FilterType},
    };
    use std::io::Cursor;

    use super::PreviewError;

    const THUMBNAIL_MAX_BYTES: usize = 10450;
    const AVATAR_MAX_BYTES: usize = 9350;

    const THUMBNAIL_SIZES: [u16; 4] = [256, 192, 128, 96];
    const AVATAR_SIZES: [u16; 4] = [160, 128, 96, 64];
    const PREVIEW_QUALITIES: [u8; 4] = [85, 75, 60, 45];

    #[derive(Debug, Clone, Copy)]
    pub struct Transcoder {
        sizes: [u16; 4],
        max_bytes: usize,
    }

    impl Default for Transcoder {
        fn default() -> Self {
            Self::thumbnail()
        }
    }

    impl Transcoder {
        pub const fn disabled() -> Self {
            Self {
                sizes: [0; 4],
                max_bytes: 0,
            }
        }

        pub const fn thumbnail() -> Self {
            Self {
                sizes: THUMBNAIL_SIZES,
                max_bytes: THUMBNAIL_MAX_BYTES,
            }
        }

        pub const fn avatar() -> Self {
            Self {
                sizes: AVATAR_SIZES,
                max_bytes: AVATAR_MAX_BYTES,
            }
        }

        pub const fn is_enabled(&self) -> bool {
            !self.is_disabled()
        }

        pub const fn is_disabled(&self) -> bool {
            self.max_bytes == 0 || self.sizes[0] == 0
        }

        pub const fn with_sizes(mut self, sizes: [u16; 4]) -> Self {
            self.sizes = sizes;
            self
        }

        pub const fn with_max_bytes(mut self, max_bytes: usize) -> Self {
            self.max_bytes = max_bytes;
            self
        }

        pub fn transcode_to_jpg(self, mut bytes: Vec<u8>) -> Result<Vec<u8>, PreviewError> {
            if !self.is_enabled() {
                return Ok(bytes);
            }

            let reader = ImageReader::new(Cursor::new(&bytes)).with_guessed_format()?;

            if reader.format() == Some(ImageFormat::Jpeg) && bytes.len() <= self.max_bytes {
                return Ok(bytes);
            }

            let img = reader.decode()?;

            let (orig_w, orig_h) = (img.width(), img.height());
            let max_orig = orig_w.max(orig_h);

            let mut last_effective = u32::MAX;
            for &size in &self.sizes {
                // Never upscale: images smaller than this slot use their natural dimensions.
                let effective = (size as u32).min(max_orig);
                if effective == last_effective {
                    continue;
                }
                last_effective = effective;

                let rgb = if effective < max_orig {
                    let resized = img.resize(effective, effective, FilterType::Lanczos3);
                    imageops::unsharpen(&resized.to_rgb8(), 0.5, 0)
                } else {
                    img.to_rgb8()
                };

                for &quality in &PREVIEW_QUALITIES {
                    bytes.clear();
                    JpegEncoder::new_with_quality(&mut bytes, quality).encode_image(&rgb)?;

                    if bytes.len() <= self.max_bytes {
                        return Ok(bytes);
                    }
                }
            }

            Err(PreviewError::TooLarge)
        }
    }
}

#[cfg(feature = "multimedia")]
pub use transcoder::Transcoder;

const URI_HEADER: &str = "data:image/jpg;base64,";

pub fn default() -> String {
    DEFAULT_PREVIEW.to_owned()
}

/// Returns the default preview on [`PreviewError`]
pub fn encode_jpg_to_uri(bytes: &[u8]) -> String {
    match try_encode_jpg_to_uri(bytes) {
        Ok(s) => s,
        Err(e) => {
            log::warn!("{e}");
            default()
        }
    }
}

pub fn try_encode_jpg_to_uri(bytes: &[u8]) -> Result<String, PreviewError> {
    if bytes.len() > MAX_PREVIEW_BYTES {
        return Err(PreviewError::TooLarge);
    }
    Ok(encode_to_uri(bytes))
}

fn encode_to_uri(bytes: &[u8]) -> String {
    let mut encoded = String::with_capacity(bytes.len() * 4 / 3 + URI_HEADER.len() + 4);
    encoded.push_str(URI_HEADER);
    BASE64_STANDARD.encode_string(bytes, &mut encoded);
    encoded
}

pub fn try_decode_jpg_from_uri(uri_str: &str) -> Result<Vec<u8>, UriDecodeError> {
    let Some(s) = uri_str.strip_prefix(URI_HEADER) else {
        return Err(UriDecodeError::NotAUri);
    };

    BASE64_STANDARD.decode(s).map_err(UriDecodeError::Base64)
}

#[derive(Debug)]
pub enum PreviewError {
    TooLarge,
    BadUri(UriDecodeError),
    Io(std::io::Error),
    #[cfg(feature = "multimedia")]
    Transcoding(image::ImageError),
    #[cfg(feature = "multimedia")]
    Tokio(tokio::task::JoinError),
}

impl From<std::io::Error> for PreviewError {
    fn from(err: std::io::Error) -> Self {
        Self::Io(err)
    }
}

#[cfg(feature = "multimedia")]
impl From<image::ImageError> for PreviewError {
    fn from(err: image::ImageError) -> Self {
        Self::Transcoding(err)
    }
}

#[cfg(feature = "multimedia")]
impl From<tokio::task::JoinError> for PreviewError {
    fn from(err: tokio::task::JoinError) -> Self {
        Self::Tokio(err)
    }
}

impl std::fmt::Display for PreviewError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::TooLarge => {
                write!(
                    f,
                    "preview size exceeds the max possible size({MAX_PREVIEW_BYTES} bytes)"
                )
            }
            Self::BadUri(e) => write!(f, "{e}"),
            Self::Io(error) => write!(f, "Cannot process preview file: {error}"),
            #[cfg(feature = "multimedia")]
            Self::Transcoding(error) => write!(f, "Cannot transcode preview: {error}"),
            #[cfg(feature = "multimedia")]
            Self::Tokio(error) => write!(f, "Failed to join the transcoding task: {error}"),
        }
    }
}

impl std::error::Error for PreviewError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::TooLarge => None,
            Self::BadUri(error) => Some(error),
            Self::Io(error) => Some(error),
            #[cfg(feature = "multimedia")]
            Self::Transcoding(error) => Some(error),
            #[cfg(feature = "multimedia")]
            Self::Tokio(error) => Some(error),
        }
    }
}

#[derive(Debug)]
pub enum UriDecodeError {
    NotAUri,
    Base64(base64::DecodeError),
}

impl std::fmt::Display for UriDecodeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotAUri => write!(f, "not a URI string"),
            Self::Base64(e) => write!(f, "{e}"),
        }
    }
}

impl std::error::Error for UriDecodeError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        if let Self::Base64(e) = self {
            Some(e)
        } else {
            None
        }
    }
}

#[derive(Clone)]
enum PreviewSource {
    Default,
    Bytes(Vec<u8>),
    DataUri(String),
    File(PathBuf),
    #[cfg(feature = "native_crypto")]
    CryptoFile(CryptoFile),
}

async fn read_plain_file(path: &PathBuf, size_limit: usize) -> std::io::Result<Vec<u8>> {
    let mut f = tokio::fs::File::open(&path).await?;
    let size_hint = f.seek(SeekFrom::End(0)).await?;
    f.seek(SeekFrom::Start(0)).await?;
    let size_hint: usize = util::cast_file_size(size_hint)?;

    if size_hint > size_limit {
        return Err(util::file_is_too_large(format!(
            "Size exceeds {size_limit} bytes"
        )));
    }

    let mut buf = Vec::with_capacity(size_hint);
    f.read_to_end(&mut buf).await?;

    Ok(buf)
}

#[cfg(feature = "native_crypto")]
async fn read_crypto_file(file: CryptoFile, size_limit: usize) -> std::io::Result<Vec<u8>> {
    let mut f = crate::crypto::fs::TokioMaybeCryptoFile::from_crypto_file(file).await?;
    let size_hint = f.size_hint().await?;

    if size_hint > size_limit {
        return Err(util::file_is_too_large(format!(
            "Size exceeds {size_limit} bytes"
        )));
    }

    let mut buf = Vec::with_capacity(size_hint);
    f.read_to_end(&mut buf).await?;

    Ok(buf)
}

fn validate_uri_preview(uri: String) -> Result<String, PreviewError> {
    let Some(s) = uri.strip_prefix(URI_HEADER) else {
        return Err(PreviewError::BadUri(UriDecodeError::NotAUri));
    };

    if s.len() > MAX_PREVIEW_BYTES * 4 / 3 {
        return Err(PreviewError::TooLarge);
    }

    Ok(uri)
}