macltok

use std::{

    borrow::Cow,

    collections::{HashMap, VecDeque},

    io::Read as _,

    panic::AssertUnwindSafe,

    sync::Arc,

};

use futures_util::{

    FutureExt, StreamExt,

    future::{BoxFuture, Shared, WeakShared},

};

use matrix_sdk::{

    Client, Room,

    media::{MediaFormat, MediaRequestParameters, MediaThumbnailSettings},

    room::reply::{EnforceThread, Reply as SdkReply},

    ruma::{

        api::{

            Metadata, OutgoingRequest,

            auth_scheme::SendAccessToken,

            client::{authenticated_media, media},

        },

        events::room::{

            MediaSource,

            message::{AddMentions, ReplyWithinThread},

        },

    },

};

use matrix_sdk_crypto::{AttachmentDecryptor, MediaEncryptionInfo};

use ruma::MxcUri;

use tokio::sync::Mutex;

use crate::{

    AttachmentUpload, Error, ErrorKind, Media, MediaId, MediaKind, MediaSize, TimelineId,

    normalise::{MediaRegistration, MediaScope},

    upload::prepare_upload,

};

/// the uploader's thumbnail is used when there is one, this is only what we ask the server for

const TIMELINE_THUMBNAIL: u32 = 800;

type FlightFuture = BoxFuture<'static, Result<Media, Error>>;

type Flight = Shared<FlightFuture>;

/// owns everything needed to turn a media id or an mxc uri into bytes, so the

/// handle can answer media queries without a trip through the engine

pub(crate) struct MediaService {

    client: Client,

    registry: Mutex<MediaRegistry>,

    flights: Mutex<HashMap<String, WeakShared<FlightFuture>>>,

    max_bytes: usize,

}

impl MediaService {

    pub(crate) fn new(client: Client, max_bytes: usize, max_sources: usize) -> Self {

        Self {

            client,

            registry: Mutex::new(MediaRegistry::new(max_sources)),

            flights: Mutex::new(HashMap::new()),

            max_bytes,

        }

    }

    pub(crate) async fn register(&self, media: Vec<MediaRegistration>) {

        if !media.is_empty() {

            self.registry.lock().await.insert_all(media);

        }

    }

    pub(crate) async fn release(&self, scope: &MediaScope) {

        self.registry.lock().await.release(scope);

    }

    pub(crate) async fn clear(&self) {

        self.registry.lock().await.clear();

    }

    pub(crate) async fn timeline_media(

        &self,

        media_id: &MediaId,

        size: MediaSize,

    ) -> Result<Media, Error> {

        let media = self

            .registry

            .lock()

            .await

            .get(media_id)

            .cloned()

            .ok_or_else(|| {

                Error::new(

                    ErrorKind::MediaUnknown,

                    "media is not registered in a retained timeline",

                )

            })?;

        let (source, format, mimetype, declared_size) = match size {

            MediaSize::Full => (

                media.full.source,

                MediaFormat::File,

                media.full.mimetype,

                media.full.declared_size,

            ),

            MediaSize::Thumbnail { .. } => match media.thumbnail {

                Some(thumbnail) => (

                    thumbnail.source,

                    MediaFormat::File,

                    thumbnail.mimetype,

                    thumbnail.declared_size,

                ),

                // the server can only resize what it can read

                None if matches!(media.kind, MediaKind::Image | MediaKind::Video)

                    && matches!(media.full.source, MediaSource::Plain(_)) =>

                {

                    (

                        media.full.source,

                        thumbnail_format(TIMELINE_THUMBNAIL, TIMELINE_THUMBNAIL),

                        media.full.mimetype,

                        None,

                    )

                }

                None => {

                    return Err(Error::new(

                        ErrorKind::MediaUnavailable,

                        "the requested media size is unavailable",

                    ));

                }

            },

        };

        if declared_size.is_some_and(|size| size > self.max_bytes as u64) {

            return Err(too_large("media exceeds the configured byte limit"));

        }

        self.fetch(MediaRequestParameters { source, format }, mimetype)

            .await

    }

    pub(crate) async fn mxc_media(&self, uri: &MxcUri, size: MediaSize) -> Result<Media, Error> {

        if !uri.is_valid() {

            return Err(Error::new(ErrorKind::Invalid, "mxc uri is malformed"));

        }

        let format = match size {

            MediaSize::Full => MediaFormat::File,

            MediaSize::Thumbnail { width, height } => thumbnail_format(width, height),

        };

        self.fetch(

            MediaRequestParameters {

                source: MediaSource::Plain(uri.to_owned()),

                format,

            },

            None,

        )

        .await

    }

    async fn fetch(

        &self,

        request: MediaRequestParameters,

        mimetype: Option<String>,

    ) -> Result<Media, Error> {

        let key = flight_key(&request);

        let flight = {

            let mut flights = self.flights.lock().await;

            flights.retain(|_, flight| flight.upgrade().is_some());

            match flights.get(&key).and_then(WeakShared::upgrade) {

                Some(flight) => flight,

                None => {

                    let flight: Flight = AssertUnwindSafe(fetch(

                        self.client.clone(),

                        request,

                        mimetype,

                        self.max_bytes,

                    ))

                    .catch_unwind()

                    .map(|result| {

                        result.unwrap_or_else(|_| {

                            Err(Error::new(ErrorKind::Other, "media fetch panicked"))

                        })

                    })

                    .boxed()

                    .shared();

                    if let Some(weak) = flight.downgrade() {

                        flights.insert(key, weak);

                    }

                    flight

                }

            }

        };

        flight.await

    }

}

fn thumbnail_format(width: u32, height: u32) -> MediaFormat {

    MediaFormat::Thumbnail(MediaThumbnailSettings::new(width.into(), height.into()))

}

fn flight_key(request: &MediaRequestParameters) -> String {

    let uri = match &request.source {

        MediaSource::Plain(uri) => uri.as_str(),

        MediaSource::Encrypted(file) => file.url.as_str(),

    };

    match &request.format {

        MediaFormat::File => uri.to_owned(),

        MediaFormat::Thumbnail(settings) => format!(

            "{uri}|{}|{}x{}|{}",

            settings.method, settings.width, settings.height, settings.animated

        ),

    }

}

fn too_large(message: &'static str) -> Error {

    Error::new(ErrorKind::MediaTooLarge, message)

}

async fn fetch(

    client: Client,

    request: MediaRequestParameters,

    mimetype: Option<String>,

    max_bytes: usize,

) -> Result<Media, Error> {

    if let Some(cached) = cached_media(&client, &request).await? {

        if cached.len() > max_bytes {

            if let Ok(store) = client.media_store().lock().await {

                let _ = store.remove_media_content(&request).await;

            }

            return Err(too_large("cached media exceeds the configured byte limit"));

        }

        return Ok(media(cached, mimetype));

    }

    let (bytes, served_mimetype) = download_bounded(&client, &request, max_bytes).await?;

    // the sdk caches plaintext under the encrypted key, so we do too or nobody can read anything

    let bytes = maybe_decrypt(bytes, &request, max_bytes)?;

    if let Ok(store) = client.media_store().lock().await {

        let _ = store

            .add_media_content(

                &request,

                bytes.clone(),

                matrix_sdk::media::store::IgnoreMediaRetentionPolicy::No,

            )

            .await;

    }

    Ok(media(bytes, mimetype.or(served_mimetype)))

}

fn media(bytes: Vec<u8>, mimetype: Option<String>) -> Media {

    let mimetype = mimetype

        .or_else(|| sniff_image(&bytes).map(str::to_owned))

        .unwrap_or_else(|| "application/octet-stream".to_owned());

    Media {

        bytes: Arc::from(bytes),

        mimetype,

    }

}

/// the cache forgets content types, and avatars never declare one. images are the common case

fn sniff_image(bytes: &[u8]) -> Option<&'static str> {

    match bytes {

        [0x89, b'P', b'N', b'G', ..] => Some("image/png"),

        [0xff, 0xd8, 0xff, ..] => Some("image/jpeg"),

        [b'G', b'I', b'F', b'8', ..] => Some("image/gif"),

        [

            b'R',

            b'I',

            b'F',

            b'F',

            _,

            _,

            _,

            _,

            b'W',

            b'E',

            b'B',

            b'P',

            ..,

        ] => Some("image/webp"),

        _ => None,

    }

}

async fn cached_media(

    client: &Client,

    request: &MediaRequestParameters,

) -> Result<Option<Vec<u8>>, Error> {

    fn cache_failed(source: impl std::error::Error + Send + Sync + 'static) -> Error {

        Error::with_source(ErrorKind::Store, "reading media cache failed", source)

    }

    client

        .media_store()

        .lock()

        .await

        .map_err(cache_failed)?

        .get_media_content(request)

        .await

        .map_err(cache_failed)

}

fn maybe_decrypt(

    bytes: Vec<u8>,

    request: &MediaRequestParameters,

    max_bytes: usize,

) -> Result<Vec<u8>, Error> {

    let MediaSource::Encrypted(file) = &request.source else {

        return Ok(bytes);

    };

    fn decrypt_failed(source: impl std::error::Error + Send + Sync + 'static) -> Error {

        Error::with_source(

            ErrorKind::MediaUnavailable,

            "media decryption failed",

            source,

        )

    }

    let info: MediaEncryptionInfo = (**file).clone().into();

    let mut cursor = std::io::Cursor::new(bytes);

    let mut decryptor = AttachmentDecryptor::new(&mut cursor, info).map_err(decrypt_failed)?;

    let mut decrypted = Vec::new();

    let mut chunk = [0_u8; 8192];

    loop {

        let read = decryptor.read(&mut chunk).map_err(decrypt_failed)?;

        if read == 0 {

            return Ok(decrypted);

        }

        if read > max_bytes.saturating_sub(decrypted.len()) {

            return Err(too_large(

                "decrypted media exceeds the configured byte limit",

            ));

        }

        decrypted.extend_from_slice(&chunk[..read]);

    }

}

/// the sdk's own download has no size cap and no timeout, so we do it by hand

async fn download_bounded(

    client: &Client,

    request: &MediaRequestParameters,

    max_bytes: usize,

) -> Result<(Vec<u8>, Option<String>), Error> {

    let uri = match &request.source {

        MediaSource::Plain(uri) => uri,

        MediaSource::Encrypted(file) => &file.url,

    };

    if uri

        .server_name()

        .is_ok_and(|server| server.as_str() == "send-queue.localhost")

    {

        return Err(Error::new(

            ErrorKind::MediaUnavailable,

            "local media is not in the media cache",

        ));

    }

    let network_failed =

        |source| Error::with_source(ErrorKind::Network, "media download failed", source);

    let mut refreshed = false;

    let response = loop {

        let http_request = media_http_request(client, request).await?;

        let (parts, _) = http_request.into_parts();

        let mut builder = client.http_client().get(parts.uri.to_string());

        for (name, value) in &parts.headers {

            builder = builder.header(name.as_str(), value.as_bytes());

        }

        let response = builder.send().await.map_err(network_failed)?;

        if response.status() != matrix_sdk::reqwest::StatusCode::UNAUTHORIZED {

            break response;

        }

        if refreshed

            || client

                .session_tokens()

                .is_none_or(|tokens| tokens.refresh_token.is_none())

        {

            return Err(Error::new(

                ErrorKind::Unauthenticated,

                "media download was not authorized",

            ));

        }

        client.refresh_access_token().await.map_err(|source| {

            Error::from_refresh_token("refreshing media access token failed", source)

        })?;

        refreshed = true;

    };

    let status = response.status();

    if !status.is_success() {

        return Err(Error::new(

            ErrorKind::MediaUnavailable,

            format!("media download failed with status {status}"),

        ));

    }

    if response

        .content_length()

        .is_some_and(|length| length > max_bytes as u64)

    {

        return Err(too_large("media exceeds the configured byte limit"));

    }

    let mimetype = response

        .headers()

        .get(http::header::CONTENT_TYPE)

        .and_then(|value| value.to_str().ok())

        .map(str::to_owned);

    let mut stream = response.bytes_stream();

    let mut bytes = Vec::new();

    while let Some(chunk) = stream.next().await {

        let chunk = chunk.map_err(network_failed)?;

        if chunk.len() > max_bytes.saturating_sub(bytes.len()) {

            return Err(too_large(

                "downloaded media exceeds the configured byte limit",

            ));

        }

        bytes.extend_from_slice(&chunk);

    }

    Ok((bytes, mimetype))

}

async fn media_http_request(

    client: &Client,

    request: &MediaRequestParameters,

) -> Result<http::Request<Vec<u8>>, Error> {

    let versions = client.supported_versions().await.map_err(|source| {

        Error::with_source(

            ErrorKind::Network,

            "reading supported server versions failed",

            source,

        )

    })?;

    let authenticated =

        authenticated_media::get_content::v1::Request::PATH_BUILDER.is_supported(&versions);

    let base = client.homeserver().to_string();

    let token = client.access_token();

    let invalid_uri =

        |source| Error::with_source(ErrorKind::Invalid, "media URI is invalid", source);

    let (uri, thumbnail) = match (&request.source, &request.format) {

        (MediaSource::Encrypted(file), _) => (&file.url, None),

        (MediaSource::Plain(uri), MediaFormat::File) => (uri, None),

        (MediaSource::Plain(uri), MediaFormat::Thumbnail(settings)) => (uri, Some(settings)),

    };

    // each endpoint has its own auth input type, so no generic helper can take them all

    macro_rules! into_http {

        ($endpoint:expr, $token:expr) => {

            $endpoint

                .try_into_http_request::<Vec<u8>>(&base, $token, Cow::Borrowed(&versions))

                .map_err(|source| {

                    Error::with_source(ErrorKind::Other, "media request is invalid", source)

                })?

        };

    }

    let authed = SendAccessToken::IfRequired(token.as_deref().unwrap_or(""));

    #[allow(deprecated)]

    let mut http_request = match (thumbnail, authenticated) {

        (None, true) => into_http!(

            authenticated_media::get_content::v1::Request::from_uri(uri).map_err(invalid_uri)?,

            authed

        ),

        (None, false) => into_http!(

            media::get_content::v3::Request::from_url(uri).map_err(invalid_uri)?,

            SendAccessToken::None

        ),

        (Some(settings), true) => {

            let mut endpoint = authenticated_media::get_content_thumbnail::v1::Request::from_uri(

                uri,

                settings.width,

                settings.height,

            )

            .map_err(invalid_uri)?;

            endpoint.method = Some(settings.method.clone());

            endpoint.animated = Some(settings.animated);

            into_http!(endpoint, authed)

        }

        (Some(settings), false) => {

            let mut endpoint = media::get_content_thumbnail::v3::Request::from_url(

                uri,

                settings.width,

                settings.height,

            )

            .map_err(invalid_uri)?;

            endpoint.method = Some(settings.method.clone());

            endpoint.animated = Some(settings.animated);

            into_http!(endpoint, SendAccessToken::None)

        }

    };

    // legacy endpoints don't ask for a token but plenty of servers want one anyway

    if let Some(token) = token.as_deref() {

        let value = format!("Bearer {token}").parse().map_err(|_| {

            Error::new(ErrorKind::Unauthenticated, "stored access token is invalid")

        })?;

        http_request

            .headers_mut()

            .insert(http::header::AUTHORIZATION, value);

    }

    Ok(http_request)