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,

    TimelineId,

    normalise::{

        ENABLED_ROOMS, ROOM_PACKS, USER_PACK, effective_usage, 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> {

    match &edit {

        PackEdit::Meta {

            usage: Patch::Set(usage),

            ..

        } => validate_usage(Some(usage))?,

        PackEdit::Put { usage, .. } => validate_usage(usage.as_ref())?,

        _ => {}

    }

    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> {

    if let PackCommand::CloneImage { usage, .. }

    | PackCommand::CloneSticker { usage, .. }

    | PackCommand::ClonePack { usage, .. } = &command

    {

        validate_usage(usage.as_ref())?;

    }

    match command {

        PackCommand::SendSticker {

            timeline_id,

            pack_id,

            shortcode,

        } => send_sticker(&client, &timeline_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,

    timeline_id: &TimelineId,

    pack_id: &str,

    shortcode: &str,

) -> Result<(), Error> {

    if matches!(timeline_id, TimelineId::Pinned(_)) {

        return Err(Error::new(

            ErrorKind::Invalid,

            "the pinned timeline can't send or paginate",

        ));

    }

    let room = client

        .get_room(timeline_id.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 mut content = json!({

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

        "url": url,

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

    });

    // the same fallback the send queue gives a threaded attachment with no reply

    if let Some(root) = timeline_id.thread_root() {

        content["m.relates_to"] = json!({

            "rel_type": "m.thread",

            "event_id": root,

            "is_falling_back": true,

            "m.in_reply_to": { "event_id": root },

        });

    }

    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(

                room_id.clone(),

                StateEventType::from(ROOM_PACKS),

                state_key.clone(),

            );

            match client.send(request).await {

                Ok(response) => parse(response.event_or_content.get()).map(Some),

                Err(error) if error.client_api_error_kind() == Some(&ApiErrorKind::NotFound) => {

                    Ok(None)

                }

                Err(source) => Err(Error::from_http("reading the pack failed", source)),

            }

        }

    }

}

async fn write_pack(

    client: &Client,

    pack: &ImagePackSource,

    content: JsonObject,

) -> Result<(), Error> {

    let written = match pack {

        ImagePackSource::User => client

            .account()

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

            .await

            .map(drop),

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

            .get_room(room_id)

            .ok_or_else(Error::room_unknown)?

            .send_state_event_raw(ROOM_PACKS, state_key, JsonValue::Object(content))

            .await

            .map(drop),

    };

    written.map_err(|source| Error::from_sdk("writing the pack failed", source))

}

/// values this crate doesn't know ride along untouched, they may mean something to someone

fn with_pack_usage(mut image: JsonObject, meta: &JsonObject) -> JsonObject {

    if known_usage(image.get("usage")).is_some() {

        return image;

    }

    let mut usage = match meta.get("usage") {

        Some(JsonValue::Array(values)) if known_usage(meta.get("usage")).is_some() => {

            values.clone()

        }

        _ => usage_values(effective_usage(None, None)),

    };

    if let Some(JsonValue::Array(own)) = image.get("usage") {

        for value in own {

            if !usage.contains(value) {

                usage.push(value.clone());

            }

        }

    }

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

    image

}

fn known_usage(value: Option<&JsonValue>) -> Option<Vec<PackUsage>> {

    let values = value?.as_array()?;

    let usage = [

        ("emoticon", PackUsage::Emoticon),

        ("sticker", PackUsage::Sticker),

    ]

    .into_iter()

    .filter_map(|(name, usage)| {

        values

            .iter()

            .any(|value| value.as_str() == Some(name))

            .then_some(usage)

    })

    .collect::<Vec<_>>();

    (!usage.is_empty()).then_some(usage)

}

/// `Some` of nothing would read back as inheriting, which `None` and `Patch::Clear` already say

fn validate_usage(usage: Option<&Vec<PackUsage>>) -> Result<(), Error> {

    if usage.is_some_and(Vec::is_empty) {

        return Err(Error::new(

            ErrorKind::Invalid,

            "usage needs at least one value, leave it unset to inherit",

        ));

    }

    Ok(())

}

fn apply_usage(image: &mut JsonObject, usage: Option<Vec<PackUsage>>) {

    if let Some(usage) = usage {

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

    }

}

fn sticker_body(image: &JsonObject, shortcode: &str) -> String {

    image

        .get("body")

        .and_then(JsonValue::as_str)

        .filter(|body| !body.trim().is_empty())

        .unwrap_or(shortcode)

        .to_owned()

}

fn free_shortcode(images: &JsonObject, wanted: &str) -> String {

    if !images.contains_key(wanted) {

        return wanted.to_owned();

    }

    (2..)

        .map(|n| format!("{wanted}-{n}"))

        .find(|candidate| !images.contains_key(candidate))

        .expect("an unbounded range finds a free name")

}

/// the object at `key`, made or replaced if it was missing or malformed

fn child<'a>(parent: &'a mut JsonObject, key: &str) -> &'a mut JsonObject {

    let slot = parent

        .entry(key.to_owned())

        .or_insert_with(|| JsonObject::new().into());

    if !slot.is_object() {

        *slot = JsonObject::new().into();

    }

    slot.as_object_mut().expect("just made it an object")

}

fn patch<T>(

    meta: &mut JsonObject,

    key: &str,

    change: Patch<T>,

    to_json: impl FnOnce(T) -> JsonValue,

) {

    match change {

        Patch::Keep => {}

        Patch::Clear => {

            meta.remove(key);

        }

        Patch::Set(value) => {

            meta.insert(key.to_owned(), to_json(value));

        }

    }

}

fn usage_json(usage: Vec<PackUsage>) -> JsonValue {

    usage_values(usage).into()

}

fn usage_values(usage: Vec<PackUsage>) -> Vec<JsonValue> {

    usage

        .into_iter()

        .map(|usage| match usage {

            PackUsage::Emoticon => "emoticon".into(),

            PackUsage::Sticker => "sticker".into(),

        })

        .collect()

}

fn info_json(info: PackImageInfo) -> JsonValue {

    let mut out = JsonObject::new();

    if let Some(mimetype) = info.mimetype {

        out.insert("mimetype".to_owned(), mimetype.into());

    }

    for (key, value) in [

        ("size", info.size_bytes),

        ("w", info.width),

        ("h", info.height),

    ] {

        if let Some(value) = value {

            out.insert(key.to_owned(), value.into());

        }

    }

    out.into()

}

fn mxc(value: &JsonValue) -> Option<OwnedMxcUri> {

    let uri = OwnedMxcUri::from(value.as_str()?);

    uri.is_valid().then_some(uri)

}

fn object(value: JsonValue) -> JsonObject {

    match value {

        JsonValue::Object(object) => object,

        _ => JsonObject::new(),

    }

}

fn parse(json: &str) -> Result<JsonObject, Error> {

    serde_json::from_str(json)

        .map_err(|source| Error::with_source(ErrorKind::Server, "the pack is malformed", source))

}

fn raw<T>(content: &JsonObject) -> Raw<T> {

    Raw::new(content)

        .expect("a json object always serialises")

        .cast_unchecked()

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn pack_usage_inherits_known_values_then_defaults_and_keeps_unknown_ones() {

        let pack = object(json!({ "usage": ["sticker", "unknown"] }));

        let inherited = with_pack_usage(object(json!({ "usage": [] })), &pack);

        assert_eq!(inherited.get("usage"), Some(&json!(["sticker", "unknown"])));

        let pack = object(json!({ "usage": ["emoticon"] }));

        let inherited = with_pack_usage(object(json!({ "usage": ["unknown"] })), &pack);

        assert_eq!(

            inherited.get("usage"),

            Some(&json!(["emoticon", "unknown"]))

        );

        let pack = object(json!({ "usage": ["sticker"] }));

        let inherited = with_pack_usage(object(json!({ "usage": ["emoticon", "unknown"] })), &pack);

        assert_eq!(

            inherited.get("usage"),

            Some(&json!(["emoticon", "unknown"]))

        );

        let pack = object(json!({ "usage": [] }));

        let inherited = with_pack_usage(object(json!({})), &pack);

        assert_eq!(

            inherited.get("usage"),

            Some(&json!(["emoticon", "sticker"]))

        );

        let pack = object(json!({ "usage": ["unknown"] }));

        let inherited = with_pack_usage(object(json!({})), &pack);

        assert_eq!(

            inherited.get("usage"),

            Some(&json!(["emoticon", "sticker"]))

        );

    }

    #[test]

    fn sticker_body_falls_back_when_missing_or_blank() {

        assert_eq!(sticker_body(&object(json!({})), "wave"), "wave");

        assert_eq!(sticker_body(&object(json!({ "body": "" })), "wave"), "wave");

        assert_eq!(

            sticker_body(&object(json!({ "body": " \t\n " })), "wave"),

            "wave"

        );

        assert_eq!(

            sticker_body(&object(json!({ "body": "  hello  " })), "wave"),

            "  hello  "

        );

    }

}