macltok

use std::sync::atomic::{AtomicU64, Ordering};

use matrix_sdk::{media::MediaEventContent, ruma::events::room::MediaSource};

use matrix_sdk_ui::timeline::{MsgLikeKind, TimelineItem as SdkTimelineItem, TimelineItemContent};

use ruma::{

    OwnedMxcUri,

    events::{

        room::{

            ImageInfo, ThumbnailInfo,

            message::{FileInfo, MessageType, VideoInfo},

        },

        sticker::{StickerEventContent, StickerMediaSource},

    },

};

use super::link_preview::sender_previews;

use crate::{MediaDescriptor, MediaId, MediaKind};

static NEXT_MEDIA_SCOPE: AtomicU64 = AtomicU64::new(1);

/// Per-timeline-session nonce namespacing recycled SDK unique IDs.

///

/// SDK 0.18 recycles [`SdkTimelineItem::unique_id`] across the local-to-remote

/// echo replacement, while transaction IDs disappear. Scoping the recycled ID

/// by a process-wide nonce keeps media IDs stable across replacement without

/// colliding across rooms, threads, or reopened sessions.

#[derive(Clone, Debug)]

pub(crate) struct MediaScope {

    nonce: u64,

}

impl MediaScope {

    pub(crate) fn new() -> Self {

        Self {

            nonce: NEXT_MEDIA_SCOPE.fetch_add(1, Ordering::Relaxed),

        }

    }

    pub(crate) fn owns(&self, id: &MediaId) -> bool {

        id.as_str()

            .split_once('#')

            .is_some_and(|(nonce, _)| nonce.parse() == Ok(self.nonce))

    }

}

#[derive(Clone)]

pub(crate) struct MediaRegistration {

    pub(crate) id: MediaId,

    pub(crate) kind: MediaKind,

    pub(crate) full: MediaSourceInfo,

    pub(crate) thumbnail: Option<MediaSourceInfo>,

}

#[derive(Clone)]

pub(crate) struct MediaSourceInfo {

    pub(crate) source: MediaSource,

    pub(crate) mimetype: Option<String>,

    pub(crate) declared_size: Option<u64>,

}

pub(super) fn media_registrations_for_item(

    item: &SdkTimelineItem,

    scope: &MediaScope,

) -> Vec<MediaRegistration> {

    let Some(event) = item.as_event() else {

        return Vec::new();

    };

    let identity = media_identity(item, scope);

    let TimelineItemContent::MsgLike(content) = event.content() else {

        return Vec::new();

    };

    match &content.kind {

        MsgLikeKind::Message(message) => media_registration(&identity, message.msgtype())

            .into_iter()

            .chain(

                sender_previews(Some(&identity), message.msgtype())

                    .into_iter()

                    .filter_map(|(_, registration)| registration),

            )

            .collect(),

        MsgLikeKind::Sticker(sticker) => media_sticker_registration(&identity, sticker.content())

            .into_iter()

            .collect(),

        _ => Vec::new(),

    }

}

pub(super) fn media_identity(item: &SdkTimelineItem, scope: &MediaScope) -> String {

    format!("{}#{}", scope.nonce, item.unique_id().0)

}

pub(super) fn media_from_message(identity: &str, message: &MessageType) -> Option<MediaDescriptor> {

    let descriptor =

        |kind, source: &MediaSource, filename: &str, caption: Option<&str>| MediaDescriptor {

            id: media_id(identity, "media"),

            kind,

            mxc: plain_mxc(source),

            filename: filename.to_owned(),

            caption: caption.map(str::to_owned),

            mimetype: None,

            size_bytes: None,

            width: None,

            height: None,

            duration_ms: None,

            blurhash: None,

            has_thumbnail: false,

        };

    let millis = |duration: std::time::Duration| u64::try_from(duration.as_millis()).ok();

    Some(match message {

        MessageType::Image(content) => {

            let info = content.info.as_deref();

            MediaDescriptor {

                mimetype: info.and_then(|info| info.mimetype.clone()),

                size_bytes: info.and_then(|info| info.size).map(u64::from),

                width: info.and_then(|info| info.width).map(u64::from),

                height: info.and_then(|info| info.height).map(u64::from),

                blurhash: info.and_then(|info| info.blurhash.clone()),

                has_thumbnail: info.is_some_and(|info| info.thumbnail_source.is_some()),

                ..descriptor(

                    MediaKind::Image,

                    &content.source,

                    content.filename(),

                    content.caption(),

                )

            }

        }

        MessageType::Video(content) => {

            let info = content.info.as_deref();

            MediaDescriptor {

                mimetype: info.and_then(|info| info.mimetype.clone()),

                size_bytes: info.and_then(|info| info.size).map(u64::from),

                width: info.and_then(|info| info.width).map(u64::from),

                height: info.and_then(|info| info.height).map(u64::from),

                duration_ms: info.and_then(|info| info.duration).and_then(millis),

                blurhash: info.and_then(|info| info.blurhash.clone()),

                has_thumbnail: info.is_some_and(|info| info.thumbnail_source.is_some()),

                ..descriptor(

                    MediaKind::Video,

                    &content.source,

                    content.filename(),

                    content.caption(),

                )

            }

        }

        MessageType::Audio(content) => {

            let info = content.info.as_deref();

            MediaDescriptor {

                mimetype: info.and_then(|info| info.mimetype.clone()),

                size_bytes: info.and_then(|info| info.size).map(u64::from),

                duration_ms: info.and_then(|info| info.duration).and_then(millis),

                ..descriptor(

                    MediaKind::Audio,

                    &content.source,

                    content.filename(),

                    content.caption(),

                )

            }

        }

        MessageType::File(content) => {

            let info = content.info.as_deref();

            MediaDescriptor {

                mimetype: info.and_then(|info| info.mimetype.clone()),

                size_bytes: info.and_then(|info| info.size).map(u64::from),

                has_thumbnail: info.is_some_and(|info| info.thumbnail_source.is_some()),

                ..descriptor(

                    MediaKind::File,

                    &content.source,

                    content.filename(),

                    content.caption(),

                )

            }

        }

        _ => return None,

    })

}

pub(super) fn media_descriptor_for_sticker(

    identity: &str,

    content: &StickerEventContent,

) -> MediaDescriptor {

    MediaDescriptor {

        id: media_id(identity, "sticker"),

        kind: MediaKind::Sticker,

        mxc: match &content.source {

            StickerMediaSource::Plain(uri) => Some(uri.clone()),

            _ => None,

        },

        filename: content.body.clone(),

        caption: None,

        mimetype: content.info.mimetype.clone(),

        size_bytes: content.info.size.map(u64::from),

        width: content.info.width.map(u64::from),

        height: content.info.height.map(u64::from),

        duration_ms: None,

        blurhash: content.info.blurhash.clone(),

        has_thumbnail: content.info.thumbnail_source.is_some(),

    }

}

fn plain_mxc(source: &MediaSource) -> Option<OwnedMxcUri> {

    match source {

        MediaSource::Plain(uri) => Some(uri.clone()),

        MediaSource::Encrypted(_) => None,

    }

}

pub(super) fn media_id(identity: &str, suffix: &str) -> MediaId {

    MediaId::new(format!("{identity}#{suffix}"))

}

fn media_sticker_registration(

    identity: &str,

    content: &StickerEventContent,

) -> Option<MediaRegistration> {

    Some(MediaRegistration {

        id: media_id(identity, "sticker"),

        kind: MediaKind::Sticker,

        full: MediaSourceInfo {

            source: content.source()?,

            mimetype: content.info.mimetype.clone(),

            declared_size: content.info.size.map(u64::from),

        },

        thumbnail: thumbnail_media_source(

            content.info.thumbnail_source.clone(),

            content.info.thumbnail_info.as_deref(),

        ),

    })

}

fn media_registration(identity: &str, message: &MessageType) -> Option<MediaRegistration> {

    match message {

        MessageType::Image(content) => media_registration_with_info(

            identity,

            MediaKind::Image,

            content.source(),

            content.info.as_deref(),

        ),

        MessageType::Video(content) => media_registration_with_info(

            identity,

            MediaKind::Video,

            content.source(),

            content.info.as_deref(),

        ),

        MessageType::File(content) => media_registration_with_info(

            identity,

            MediaKind::File,

            content.source(),

            content.info.as_deref(),

        ),

        MessageType::Audio(content) => media_registration_with_parts(

            identity,

            MediaKind::Audio,

            content.source(),

            content.info.as_ref().and_then(|info| info.mimetype.clone()),

            content

                .info

                .as_ref()

                .and_then(|info| info.size)

                .map(u64::from),

            None,

        ),

        _ => None,

    }

}

fn media_registration_with_info<T: MediaInfo>(

    identity: &str,

    kind: MediaKind,

    source: Option<MediaSource>,

    info: Option<&T>,

) -> Option<MediaRegistration> {

    media_registration_with_parts(

        identity,

        kind,

        source,

        info.and_then(MediaInfo::mimetype),

        info.and_then(MediaInfo::declared_size),

        info.and_then(MediaInfo::thumbnail),

    )

}

fn media_registration_with_parts(

    identity: &str,

    kind: MediaKind,

    source: Option<MediaSource>,

    mimetype: Option<String>,

    declared_size: Option<u64>,

    thumbnail: Option<MediaSourceInfo>,

) -> Option<MediaRegistration> {

    Some(MediaRegistration {

        id: media_id(identity, "media"),

        kind,

        full: MediaSourceInfo {

            source: source?,

            mimetype,

            declared_size,

        },

        thumbnail,

    })

}

trait MediaInfo {

    fn mimetype(&self) -> Option<String>;

    fn declared_size(&self) -> Option<u64>;

    fn thumbnail(&self) -> Option<MediaSourceInfo>;

}

macro_rules! impl_media_info {

    ($type:ty) => {

        impl MediaInfo for $type {

            fn mimetype(&self) -> Option<String> {

                self.mimetype.clone()

            }

            fn declared_size(&self) -> Option<u64> {

                self.size.map(u64::from)

            }

            fn thumbnail(&self) -> Option<MediaSourceInfo> {

                thumbnail_media_source(

                    self.thumbnail_source.clone(),

                    self.thumbnail_info.as_deref(),

                )

            }

        }

    };

}

impl_media_info!(ImageInfo);

impl_media_info!(VideoInfo);

impl_media_info!(FileInfo);

fn thumbnail_media_source(

    source: Option<MediaSource>,

    info: Option<&ThumbnailInfo>,

) -> Option<MediaSourceInfo> {

    source.map(|source| MediaSourceInfo {

        source,

        mimetype: info.and_then(|info| info.mimetype.clone()),

        declared_size: info.and_then(|info| info.size).map(u64::from),

    })

}

#[cfg(test)]

mod tests {

    use matrix_sdk::ruma::{

        events::room::{ImageInfo, message::ImageMessageEventContent},

        mxc_uri,

    };

    use super::*;

    #[test]

    fn image_messages_keep_media_metadata() {

        let mut info = ImageInfo::new();

        info.mimetype = Some("image/png".to_owned());

        info.size = Some(65_536_u32.into());

        info.width = Some(640_u32.into());

        info.height = Some(480_u32.into());

        let content = ImageMessageEventContent::plain(

            "screen.png".to_owned(),

            mxc_uri!("mxc://example.org/screen").to_owned(),

        )

        .info(Box::new(info));

        let media = media_from_message("$event:example.org", &MessageType::Image(content)).unwrap();

        assert_eq!(media.id.as_str(), "$event:example.org#media");

        assert_eq!(media.kind, MediaKind::Image);

        assert_eq!(media.filename, "screen.png");

        assert_eq!(media.caption, None);

        assert_eq!(media.mimetype.as_deref(), Some("image/png"));

        assert_eq!(media.size_bytes, Some(65_536));

        assert_eq!((media.width, media.height), (Some(640), Some(480)));

    }

}