macltok

use futures_util::{StreamExt, stream};

use matrix_sdk::{

    Client,

    ruma::{

        api::{client::state::get_state_event_for_key, error::ErrorKind as ApiErrorKind},

        events::{GlobalAccountDataEventType, StateEventType, room::MediaSource},

        serde::{JsonObject, Raw},

    },

};

use ruma::{OwnedEventId, OwnedMxcUri, OwnedRoomId};

use serde_json::{Value as JsonValue, json};

use tokio::sync::mpsc;

use super::{

    commands::{self, Job},

    enqueue, media,

    mutations::{Mutation, MutationJob},

};

use crate::{

    CommandOutput, Error, ErrorKind, ImagePackSource, PackCommand, PackImageInfo, PackUsage, Patch,

    normalise::{ENABLED_ROOMS, ROOM_PACKS, USER_PACK, resolve_pack_id, stored_pack},

    transport::Responder,

};

/// read-modify-writes of one pack, which run in order per pack

pub(super) enum PackEdit {

    Create {

        display_name: String,

    },

    Delete,

    Meta {

        display_name: Patch<String>,

        avatar: Patch<OwnedMxcUri>,

        usage: Patch<Vec<PackUsage>>,

    },

    Put {

        shortcode: String,

        mxc: OwnedMxcUri,

        body: Option<String>,

        usage: Option<Vec<PackUsage>>,

        info: Option<PackImageInfo>,

    },

    Rename {

        from: String,

        to: String,

    },

    Remove {

        shortcode: String,

    },

    /// copies already on our homeserver. `create` is the pack meta for a room pack that doesn't exist yet

    Insert {

        images: Vec<(String, JsonObject)>,

        create: Option<JsonObject>,

    },

}

pub(super) async fn edit_pack(

    client: &Client,

    pack: &ImagePackSource,

    edit: PackEdit,

) -> Result<CommandOutput, Error> {

    let current = fetch_pack(client, pack).await?;

    let mut content = current.clone().unwrap_or_default();

    let mut output = CommandOutput::Done;

    match edit {

        PackEdit::Create { display_name } => {

            if !content.is_empty() {

                return Err(Error::new(

                    ErrorKind::Invalid,

                    "that state key already holds a pack",

                ));

            }

            content = object(json!({ "pack": { "display_name": display_name }, "images": {} }));

        }

        PackEdit::Delete => content.clear(),

        PackEdit::Meta {

            display_name,

            avatar,

            usage,

        } => {

            let meta = child(&mut content, "pack");

            patch(meta, "display_name", display_name, JsonValue::String);

            patch(meta, "avatar_url", avatar, |uri| uri.to_string().into());

            patch(meta, "usage", usage, usage_json);

        }

        PackEdit::Put {

            shortcode,

            mxc,

            body,

            usage,

            info,

        } => {

            let images = child(&mut content, "images");

            let mut image = match images.remove(&shortcode) {

                Some(JsonValue::Object(image)) => image,

                _ => JsonObject::new(),

            };

            let same_media = image.get("url").and_then(JsonValue::as_str) == Some(mxc.as_str());

            image.insert("url".to_owned(), mxc.to_string().into());

            if let Some(body) = body {

                image.insert("body".to_owned(), body.into());

            }

            if let Some(usage) = usage {

                image.insert("usage".to_owned(), usage_json(usage));

            }

            match info {

                Some(info) => {

                    image.insert("info".to_owned(), info_json(info));

                }

                None if !same_media => {

                    image.remove("info");

                }

                None => {}

            }

            images.insert(shortcode, image.into());

        }

        PackEdit::Rename { from, to } => {

            let images = child(&mut content, "images");

            if images.contains_key(&to) {

                return Err(Error::new(ErrorKind::Invalid, "that shortcode is taken"));

            }

            let image = images

                .remove(&from)

                .ok_or_else(|| Error::new(ErrorKind::Invalid, "no image has that shortcode"))?;

            images.insert(to, image);

        }

        PackEdit::Remove { shortcode } => {

            child(&mut content, "images").remove(&shortcode);

        }

        PackEdit::Insert { images, create } => {

            if content.is_empty() {

                match (create, pack) {

                    (Some(meta), _) => {

                        content.insert("pack".to_owned(), meta.into());

                    }

                    (None, ImagePackSource::Room { .. }) => {

                        return Err(Error::new(

                            ErrorKind::Invalid,

                            "there's no pack there to copy into",

                        ));

                    }

                    (None, ImagePackSource::User) => {}

                }

            }

            let target = child(&mut content, "images");

            let shortcodes = images

                .into_iter()

                .map(|(shortcode, image)| {

                    let free = free_shortcode(target, &shortcode);

                    target.insert(free.clone(), image.into());

                    free

                })

                .collect();

            output = CommandOutput::Cloned {

                shortcodes,

                skipped: Vec::new(),

            };

        }

    }

    if current.as_ref() != Some(&content) {

        write_pack(client, pack, content).await?;

    }

    Ok(output)

}

pub(super) async fn set_enabled_everywhere(

    client: &Client,

    room_id: OwnedRoomId,

    state_key: String,

    enabled: bool,

) -> Result<(), Error> {

    let account = client.account();

    let current = account

        .fetch_account_data(GlobalAccountDataEventType::from(ENABLED_ROOMS))

        .await

        .map_err(|source| Error::from_sdk("reading enabled packs failed", source))?

        .map(|raw| parse(raw.json().get()))

        .transpose()?;

    let mut content = current.clone().unwrap_or_default();

    let rooms = child(&mut content, "rooms");

    if enabled {

        child(rooms, room_id.as_str()).insert(state_key, JsonObject::new().into());

    } else {

        let state_keys = child(rooms, room_id.as_str());

        state_keys.remove(&state_key);

        if state_keys.is_empty() {

            rooms.remove(room_id.as_str());

        }

    }

    if current.as_ref() == Some(&content) {

        return Ok(());

    }

    account

        .set_account_data_raw(ENABLED_ROOMS.into(), raw(&content))

        .await

        .map(drop)

        .map_err(|source| Error::from_sdk("writing enabled packs failed", source))

}

/// the commands that download, upload or send, so they stay off the per-pack queue and only

/// take their turn for the final write

pub(super) async fn run_detached(

    client: Client,

    mutations: mpsc::Sender<MutationJob>,

    max_bytes: usize,

    command: PackCommand,

) -> Result<CommandOutput, Error> {

    match command {

        PackCommand::SendSticker {

            room_id,

            pack_id,

            shortcode,

        } => send_sticker(&client, room_id, &pack_id, &shortcode)

            .await

            .map(|()| CommandOutput::Done),

        PackCommand::Rehost(mxc) => rehost(&client, mxc, None, max_bytes)

            .await

            .map(CommandOutput::Mxc),

        PackCommand::CloneImage {

            from_pack,

            shortcode,

            to,

            usage,

        } => {

            let mut image = pack_image(&client, &from_pack, &shortcode).await?;

            apply_usage(&mut image, usage);

            let image = rehost_image(&client, image, max_bytes).await?;

            insert(&mutations, to, vec![(shortcode, image)], None, Vec::new()).await

        }

        PackCommand::CloneSticker {

            room_id,

            event_id,

            to,

            shortcode,

            usage,

        } => {

            let mut image = sticker_image(&client, room_id, event_id).await?;

            let body = sticker_body(&image, &shortcode);

            image.insert("body".to_owned(), body.into());

            apply_usage(&mut image, usage);

            let image = rehost_image(&client, image, max_bytes).await?;

            insert(&mutations, to, vec![(shortcode, image)], None, Vec::new()).await

        }

        PackCommand::ClonePack {

            from_pack,

            to,

            display_name,

            usage,

        } => {

            clone_pack(

                &client,

                &mutations,

                max_bytes,

                &from_pack,

                to,

                display_name,

                usage,

            )

            .await

        }

        _ => Err(Error::new(ErrorKind::Other, "pack command was misrouted")),

    }

}

async fn clone_pack(

    client: &Client,

    mutations: &mpsc::Sender<MutationJob>,

    max_bytes: usize,

    from_pack: &str,

    to: ImagePackSource,

    display_name: Option<String>,

    usage: Option<Vec<PackUsage>>,

) -> Result<CommandOutput, Error> {

    let source = resolve_pack_id(client, from_pack)

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "no pack has that id"))?;

    let content = stored_pack(client, &source)

        .await?

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "that pack is empty"))?;

    let mut meta = match content.get("pack") {

        Some(JsonValue::Object(meta)) => meta.clone(),

        _ => JsonObject::new(),

    };

    let images: Vec<(String, JsonObject)> = match content.get("images") {

        Some(JsonValue::Object(images)) => images

            .iter()

            .filter_map(|(shortcode, image)| {

                let mut image = with_pack_usage(image.as_object()?.clone(), &meta);

                apply_usage(&mut image, usage.clone());

                Some((shortcode.clone(), image))

            })

            .collect(),

        _ => Vec::new(),

    };

    let results: Vec<_> = stream::iter(images)

        .map(|(shortcode, image)| async move {

            match rehost_image(client, image, max_bytes).await {

                Ok(image) => Ok((shortcode, image)),

                Err(error) => {

                    tracing::debug!(%error, shortcode, "skipping an image that wouldn't rehost");

                    Err(shortcode)

                }

            }

        })

        .buffered(4)

        .collect()

        .await;

    let (copies, skipped): (Vec<_>, Vec<_>) = results.into_iter().partition(Result::is_ok);

    let copies = copies.into_iter().filter_map(Result::ok).collect();

    let skipped = skipped.into_iter().filter_map(Result::err).collect();

    let create = match to {

        ImagePackSource::User => None,

        ImagePackSource::Room { .. } => {

            if let Some(display_name) = display_name {

                meta.insert("display_name".to_owned(), display_name.into());

            }

            meta.remove("usage");

            if let Some(avatar) = meta.remove("avatar_url").and_then(|uri| mxc(&uri))

                && let Ok(avatar) = rehost(client, avatar, None, max_bytes).await

            {

                meta.insert("avatar_url".to_owned(), avatar.to_string().into());

            }

            Some(meta)

        }

    };

    insert(mutations, to, copies, create, skipped).await

}

async fn insert(

    mutations: &mpsc::Sender<MutationJob>,

    pack: ImagePackSource,

    images: Vec<(String, JsonObject)>,

    create: Option<JsonObject>,

    skipped: Vec<String>,

) -> Result<CommandOutput, Error> {

    let (reply, answer) = Responder::channel();

    let command = Mutation::Pack {

        pack,

        edit: PackEdit::Insert { images, create },

    };

    enqueue(mutations, Job { command, reply });

    match answer.await? {

        CommandOutput::Cloned { shortcodes, .. } => Ok(CommandOutput::Cloned {

            shortcodes,

            skipped,

        }),

        output => Ok(output),

    }

}

async fn send_sticker(

    client: &Client,

    room_id: OwnedRoomId,

    pack_id: &str,

    shortcode: &str,

) -> Result<(), Error> {

    let room = client.get_room(&room_id).ok_or_else(Error::room_unknown)?;

    let image = pack_image(client, pack_id, shortcode).await?;

    let url = image

        .get("url")

        .and_then(mxc)

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "that image has no usable url"))?;

    let content = json!({

        "body": sticker_body(&image, shortcode),

        "url": url,

        "info": image.get("info").cloned().unwrap_or_else(|| json!({})),

    });

    room.send_queue()

        .send_raw(raw(&object(content)), "m.sticker".to_owned())

        .await

        .map(drop)

        .map_err(|source| {

            Error::with_source(ErrorKind::Store, "queueing the sticker failed", source)

        })

}

/// the image with its pack's usage folded in, so it means the same once it leaves the pack

async fn pack_image(client: &Client, pack_id: &str, shortcode: &str) -> Result<JsonObject, Error> {

    let source = resolve_pack_id(client, pack_id)

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "no pack has that id"))?;

    let content = stored_pack(client, &source).await?.unwrap_or_default();

    let image = content

        .get("images")

        .and_then(|images| images.get(shortcode))

        .and_then(JsonValue::as_object)

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "no image has that shortcode"))?

        .clone();

    let meta = content

        .get("pack")

        .and_then(JsonValue::as_object)

        .cloned()

        .unwrap_or_default();

    Ok(with_pack_usage(image, &meta))

}

async fn sticker_image(

    client: &Client,

    room_id: OwnedRoomId,

    event_id: OwnedEventId,

) -> Result<JsonObject, Error> {

    let room = client.get_room(&room_id).ok_or_else(Error::room_unknown)?;

    let event = room

        .load_or_fetch_event(&event_id, None)

        .await

        .map_err(|source| Error::from_sdk("loading the sticker failed", source))?;

    let raw = event.raw();

    if raw.get_field::<String>("type").ok().flatten().as_deref() != Some("m.sticker") {

        return Err(Error::new(ErrorKind::Invalid, "that event isn't a sticker"));

    }

    let content = raw

        .get_field::<JsonObject>("content")

        .ok()

        .flatten()

        .unwrap_or_default();

    let mut image = JsonObject::new();

    for key in ["url", "body", "info"] {

        if let Some(value) = content.get(key) {

            image.insert(key.to_owned(), value.clone());

        }

    }

    image.insert("usage".to_owned(), usage_json(vec![PackUsage::Sticker]));

    Ok(image)

}

async fn rehost_image(

    client: &Client,

    mut image: JsonObject,

    max_bytes: usize,

) -> Result<JsonObject, Error> {

    let url = image

        .get("url")

        .and_then(mxc)

        .ok_or_else(|| Error::new(ErrorKind::Invalid, "that image has no usable url"))?;

    let mimetype = image

        .get("info")

        .and_then(|info| info.get("mimetype"))

        .and_then(JsonValue::as_str)

        .map(str::to_owned);

    let url = rehost(client, url, mimetype, max_bytes).await?;

    image.insert("url".to_owned(), url.to_string().into());

    Ok(image)

}

pub(super) async fn rehost(

    client: &Client,

    mxc: OwnedMxcUri,

    mimetype: Option<String>,

    max_bytes: usize,

) -> Result<OwnedMxcUri, Error> {

    let homeserver = client.user_id().map(ruma::UserId::server_name);

    if mxc

        .server_name()

        .ok()

        .is_some_and(|server| Some(server) == homeserver)

    {

        return Ok(mxc);

    }

    let media = media::download_file(client, MediaSource::Plain(mxc), mimetype, max_bytes).await?;

    commands::upload_media(client, &media.mimetype, media.bytes.to_vec(), max_bytes).await

}

/// the server's copy, since the store lags a sync behind and an edit made just before would be lost

async fn fetch_pack(client: &Client, pack: &ImagePackSource) -> Result<Option<JsonObject>, Error> {

    match pack {

        ImagePackSource::User => client

            .account()

            .fetch_account_data(GlobalAccountDataEventType::from(USER_PACK))

            .await

            .map_err(|source| Error::from_sdk("reading the pack failed", source))?

            .map(|raw| parse(raw.json().get()))

            .transpose(),

        ImagePackSource::Room { room_id, state_key } => {

            let request = get_state_event_for_key::v3::Request::new(