magnetar/src/model/processing/note.rs

use std::sync::Arc;

use crate::model::data::id::BaseId;
use crate::model::data::note::{NoteAttachmentSource, NoteBaseSource, NoteDetailSource};
use crate::model::processing::emoji::EmojiModel;
use crate::model::processing::user::UserModel;
use crate::model::processing::{get_mm_token_emoji, PackError, PackResult};
use crate::model::{PackType, PackingContext};
use compact_str::CompactString;
use either::Either;
use futures_util::future::{try_join_all, BoxFuture};
use futures_util::{FutureExt, StreamExt, TryFutureExt, TryStreamExt};
use magnetar_calckey_model::ck::sea_orm_active_enums::NoteVisibilityEnum;
use magnetar_calckey_model::emoji::EmojiTag;
use magnetar_calckey_model::note_model::{
    sub_interaction_reaction, sub_interaction_renote, NoteData, NoteResolveOptions,
    NoteVisibilityFilterFactory,
};
use magnetar_calckey_model::poll::PollResolver;
use magnetar_calckey_model::sea_orm::prelude::Expr;
use magnetar_calckey_model::sea_orm::sea_query::{Alias, IntoIden, PgFunc, Query, SimpleExpr};
use magnetar_calckey_model::sea_orm::{ActiveEnum, ColumnTrait, IntoSimpleExpr};
use magnetar_calckey_model::{ck, CalckeyDbError};
use magnetar_common::util::{parse_reaction, RawReaction};
use magnetar_sdk::mmm::Token;
use magnetar_sdk::types::drive::PackDriveFileBase;
use magnetar_sdk::types::emoji::EmojiContext;
use magnetar_sdk::types::note::{
    NoteAttachmentExt, NoteBase, NoteDetailExt, NoteSelfContextExt, PackNoteBase,
    PackNoteMaybeAttachments, PackNoteMaybeFull, PackPollBase, PollBase, Reaction, ReactionPair,
    ReactionShortcode, ReactionUnicode, ReactionUnknown,
};
use magnetar_sdk::types::user::UserRelationship;
use magnetar_sdk::types::{Id, MmXml};
use magnetar_sdk::{mmm, Optional, Packed, Required};
use serde::Deserialize;
use tokio::try_join;

use super::drive::DriveModel;
use super::user::UserShapedData;

#[derive(Debug, Clone)]
pub struct NoteVisibilityFilterSimple(Option<String>);

impl NoteVisibilityFilterFactory for NoteVisibilityFilterSimple {
    fn with_note_and_user_tables(&self, note_tbl: Option<Alias>) -> SimpleExpr {
        let note_tbl_name =
            note_tbl.map_or_else(|| ck::note::Entity.into_iden(), |a| a.into_iden());

        let note_visibility = Expr::col((note_tbl_name.clone(), ck::note::Column::Visibility));
        let note_mentions = Expr::col((note_tbl_name.clone(), ck::note::Column::Mentions));
        let note_reply_user_id = Expr::col((note_tbl_name.clone(), ck::note::Column::ReplyUserId));
        let note_visible_user_ids =
            Expr::col((note_tbl_name.clone(), ck::note::Column::VisibleUserIds));
        let note_user_id = Expr::col((note_tbl_name, ck::note::Column::UserId));

        let is_public = note_visibility
            .clone()
            .eq(NoteVisibilityEnum::Public.as_enum())
            .or(note_visibility
                .clone()
                .eq(NoteVisibilityEnum::Home.as_enum()));

        let Some(user_id_str) = &self.0 else {
            return is_public;
        };

        let self_user_id = SimpleExpr::Constant(user_id_str.into());

        let is_self = note_user_id.clone().eq(self_user_id.clone());

        let is_visible_specified = {
            let either_specified_or_followers = note_visibility
                .clone()
                .eq(NoteVisibilityEnum::Specified.as_enum())
                .or(note_visibility
                    .clone()
                    .eq(NoteVisibilityEnum::Followers.as_enum()))
                .into_simple_expr();

            let mentioned_or_specified = self_user_id
                .clone()
                .eq(PgFunc::any(note_mentions.into_simple_expr()))
                .or(self_user_id.eq(PgFunc::any(note_visible_user_ids)));

            either_specified_or_followers.and(mentioned_or_specified)
        };

        let is_visible_followers = {
            note_visibility
                .eq(NoteVisibilityEnum::Followers.as_enum())
                .and(
                    note_user_id.in_subquery(
                        Query::select()
                            .column(ck::following::Column::FolloweeId)
                            .from(ck::following::Entity)
                            .cond_where(ck::following::Column::FollowerId.eq(user_id_str))
                            .to_owned(),
                    ),
                )
                .or(note_reply_user_id.eq(user_id_str))
        };

        is_self
            .or(is_public)
            .or(is_visible_followers)
            .or(is_visible_specified)
    }
}

pub struct NoteVisibilityFilterModel;

impl NoteVisibilityFilterModel {
    pub async fn is_note_visible(
        &self,
        ctx: &PackingContext,
        user: Option<&ck::user::Model>,
        note: &ck::note::Model,
    ) -> Result<bool, CalckeyDbError> {
        if user.is_some_and(|user| user.id == note.user_id) {
            return Ok(true);
        }

        if matches!(
            note.visibility,
            NoteVisibilityEnum::Public | NoteVisibilityEnum::Home
        ) {
            return Ok(true);
        }

        if matches!(
            note.visibility,
            NoteVisibilityEnum::Followers | NoteVisibilityEnum::Specified
        ) {
            let Some(user) = user else {
                return Ok(false);
            };

            if note.mentions.contains(&user.id) || note.visible_user_ids.contains(&user.id) {
                return Ok(true);
            }

            if matches!(note.visibility, NoteVisibilityEnum::Specified) {
                return Ok(false);
            }

            let following = ctx
                .is_relationship_between(
                    Either::Right(user),
                    Either::Left(&note.user_id),
                    UserRelationship::Follow,
                )
                .await?;

            // The second condition generally will not happen in the API,
            // however it allows some AP processing, with activities
            // between two foreign objects
            return Ok(following || user.host.is_some() && note.user_host.is_some());
        }

        Ok(false)
    }

    pub fn new_note_visibility_filter(&self, user: Option<&str>) -> NoteVisibilityFilterSimple {
        NoteVisibilityFilterSimple(user.map(str::to_string))
    }
}

pub trait NoteShapedData<'a>: Send + Sync {
    fn note(&self) -> &'a ck::note::Model;
    fn interaction_user_renote(&self) -> Option<&'a sub_interaction_renote::Model>;
    fn interaction_user_reaction(&self) -> Option<&'a sub_interaction_reaction::Model>;
    fn user(&self) -> Arc<dyn UserShapedData<'a> + 'a>;
    fn reply(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>>;
    fn renote(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>>;
}

pub struct NoteBorrowedData<'a> {
    pub note: &'a ck::note::Model,
    pub interaction_user_renote: Option<&'a sub_interaction_renote::Model>,
    pub interaction_user_reaction: Option<&'a sub_interaction_reaction::Model>,
    pub user: Arc<dyn UserShapedData<'a> + 'a>,
    pub reply: Option<Arc<dyn NoteShapedData<'a> + 'a>>,
    pub renote: Option<Arc<dyn NoteShapedData<'a> + 'a>>,
}

impl<'a> NoteShapedData<'a> for NoteBorrowedData<'a> {
    fn note(&self) -> &'a ck::note::Model {
        self.note
    }

    fn interaction_user_renote(&self) -> Option<&'a sub_interaction_renote::Model> {
        self.interaction_user_renote
    }

    fn interaction_user_reaction(&self) -> Option<&'a sub_interaction_reaction::Model> {
        self.interaction_user_reaction
    }

    fn user(&self) -> Arc<dyn UserShapedData<'a> + 'a> {
        self.user.clone()
    }

    fn reply(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>> {
        self.reply.as_ref().cloned()
    }

    fn renote(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>> {
        self.renote.as_ref().cloned()
    }
}

impl<'a> NoteShapedData<'a> for &'a NoteData {
    fn note(&self) -> &'a ck::note::Model {
        &self.note
    }

    fn interaction_user_renote(&self) -> Option<&'a sub_interaction_renote::Model> {
        self.interaction_user_renote.as_ref()
    }

    fn interaction_user_reaction(&self) -> Option<&'a sub_interaction_reaction::Model> {
        self.interaction_user_reaction.as_ref()
    }

    fn user(&self) -> Arc<dyn UserShapedData<'a> + 'a> {
        Arc::new(&self.user)
    }

    fn reply(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>> {
        self.reply
            .as_deref()
            .map(|x| Arc::new(x) as Arc<dyn NoteShapedData<'a> + 'a>)
    }

    fn renote(&self) -> Option<Arc<dyn NoteShapedData<'a> + 'a>> {
        self.renote
            .as_deref()
            .map(|x| Arc::new(x) as Arc<dyn NoteShapedData<'a> + 'a>)
    }
}

struct SpeechTransformNyan;

impl SpeechTransformNyan {
    fn new() -> Self {
        SpeechTransformNyan
    }

    fn transform(&self, text: &mut CompactString) {
        // TODO
    }
}

pub struct NoteModel {
    pub attachments: bool,
    pub with_context: bool,
}

impl NoteModel {
    pub fn tokenize_note_text(&self, note: &ck::note::Model) -> Option<Token> {
        note.text
            .as_deref()
            .map(|text| mmm::Context::default().parse_full(text))
    }

    pub fn tokenize_note_cw(&self, note: &ck::note::Model) -> Option<Token> {
        note.cw
            .as_deref()
            .map(|text| mmm::Context::default().parse_ui(text))
    }

    pub async fn extract_base(
        &self,
        ctx: &PackingContext,
        note_data: &dyn NoteShapedData<'_>,
    ) -> PackResult<PackNoteBase> {
        let note = note_data.note();

        let Required(ref user) = UserModel
            .base_from_existing(ctx, note_data.user().as_ref())
            .await?
            .user;

        let cw_tok = self.tokenize_note_cw(note);
        let mut text_tok = self.tokenize_note_text(note);

        let mut emoji_extracted = Vec::new();

        if let Some(cw_tok) = &cw_tok {
            emoji_extracted.extend_from_slice(&get_mm_token_emoji(cw_tok));
        }

        if let Some(text_tok) = &mut text_tok {
            emoji_extracted.extend_from_slice(&get_mm_token_emoji(text_tok));

            if note_data.user().user().is_cat && note_data.user().user().speak_as_cat {
                let transformer = SpeechTransformNyan::new();
                text_tok.walk_speech_transform(&|text| transformer.transform(text));
            }
        }

        let emoji_model = EmojiModel;

        let shortcodes = emoji_model.deduplicate_emoji(ctx, emoji_extracted);
        // Parse the JSON into an ordered map and turn it into a Vec of pairs, parsing the reaction codes
        // Failed reaction parses -> Left, Successful ones -> Right
        let reactions_raw =
            serde_json::Map::<String, serde_json::Value>::deserialize(&note.reactions)?
                .into_iter()
                .map(|(ref code, count)| {
                    let reaction = parse_reaction(code)
                        .map_or_else(|| Either::Left(code.to_string()), Either::Right);

                    (
                        reaction,
                        count,
                        note_data
                            .interaction_user_reaction()
                            .as_ref()
                            .and_then(|r| r.reaction_name.as_deref())
                            .map(|r| r == code),
                    )
                })
                .map(|(code, count, self_reacted)| {
                    Ok((code, usize::deserialize(count)?, self_reacted))
                })
                .filter(|v| !v.as_ref().is_ok_and(|(_, count, _)| *count == 0))
                .collect::<Result<Vec<_>, serde_json::Error>>()?;
        // Pick out all successfully-parsed shortcode emojis
        let reactions_to_resolve = reactions_raw
            .iter()
            .map(|(code, _, _)| code)
            .map(Either::as_ref)
            .filter_map(Either::right)
            .filter_map(|c| match c {
                RawReaction::Shortcode { shortcode, host } => Some(EmojiTag {
                    name: shortcode,
                    host: host.as_deref(),
                }),
                _ => None,
            })
            .collect::<Vec<_>>();

        let reaction_fetch = ctx
            .service
            .emoji_cache
            .get_many_tagged(&reactions_to_resolve)
            .map_err(PackError::from);
        let emoji_fetch = emoji_model.fetch_many_emojis(
            ctx,
            &shortcodes,
            note_data.user().user().host.as_deref(),
        );

        let (reactions_fetched, emojis) = try_join!(reaction_fetch, emoji_fetch)?;

        // Left reactions and the Right ones that didn't resolve to any emoji are turned back into Unknown
        let reactions = &reactions_raw
            .into_iter()
            .map(|(raw, count, self_reaction)| {
                let reaction = raw.either(
                    |raw| Reaction::Unknown(ReactionUnknown { raw }),
                    |raw| match raw {
                        RawReaction::Unicode(text) => Reaction::Unicode(ReactionUnicode(text)),
                        RawReaction::Shortcode { shortcode, host } => reactions_fetched
                            .iter()
                            .find(|e| e.host == host && e.name == shortcode)
                            .map_or_else(
                                || {
                                    Reaction::Unknown(ReactionUnknown {
                                        raw: format!(
                                            ":{shortcode}{}:",
                                            host.as_deref()
                                                .map(|h| format!("@{h}"))
                                                .unwrap_or_default()
                                        ),
                                    })
                                },
                                |e| {
                                    Reaction::Shortcode(ReactionShortcode {
                                        name: shortcode.clone(),
                                        host: host.clone(),
                                        url: e.public_url.clone(),
                                    })
                                },
                            ),
                    },
                );

                match self_reaction {
                    Some(self_reaction) => {
                        ReactionPair::WithContext(reaction, count, self_reaction)
                    }
                    None => ReactionPair::WithoutContext(reaction, count),
                }
            })
            .collect::<Vec<_>>();

        let emoji_context = &EmojiContext(emojis);

        let note_base = NoteBase::extract(
            ctx,
            NoteBaseSource {
                note,
                cw_mm: cw_tok
                    .as_ref()
                    .map(mmm::to_xml_string)
                    .and_then(Result::ok)
                    .map(MmXml)
                    .as_ref(),
                text_mm: text_tok
                    .as_ref()
                    .map(mmm::to_xml_string)
                    .and_then(Result::ok)
                    .map(MmXml)
                    .as_ref(),
                reactions,
                user,
                emoji_context,
            },
        );

        Ok(PackNoteBase::pack_from((
            Required(Id::from(note.id.clone())),
            Required(note_base),
        )))
    }

    fn extract_interaction(
        &self,
        ctx: &PackingContext,
        note: &dyn NoteShapedData<'_>,
    ) -> PackResult<Option<NoteSelfContextExt>> {
        Ok(note
            .interaction_user_renote()
            .map(|renote_info| NoteSelfContextExt::extract(ctx, renote_info)))
    }

    async fn extract_attachments(
        &self,
        ctx: &PackingContext,
        drive_model: &DriveModel,
        note: &ck::note::Model,
    ) -> PackResult<Option<Vec<PackDriveFileBase>>> {
        if self.attachments {
            let futures = try_join_all(
                note.file_ids
                    .iter()
                    .map(|id| drive_model.get_cached_base(ctx, id)),
            );

            let att = futures.await?.into_iter().flatten().collect::<Vec<_>>();

            return Ok(Some(att));
        }

        Ok(None)
    }

    async fn extract_poll(
        &self,
        ctx: &PackingContext,
        note: &ck::note::Model,
    ) -> PackResult<Option<PackPollBase>> {
        if !note.has_poll {
            return Ok(None);
        }

        let poll_resolver = PollResolver::new(ctx.service.db.clone());

        let Some(poll) = poll_resolver.get_poll(&note.id).await? else {
            return Ok(None);
        };

        let votes = match ctx.self_user.as_deref() {
            Some(u) => Some(poll_resolver.get_poll_votes_by(&note.id, &u.id).await?),
            None => None,
        };

        Ok(Some(PackPollBase::pack_from((
            Required(Id::from(poll.get_id())),
            Required(PollBase::extract(ctx, (&poll, votes.as_deref()))),
        ))))
    }

    async fn pack_single_attachments(
        &self,
        ctx: &PackingContext,
        drive_model: &DriveModel,
        note_data: &dyn NoteShapedData<'_>,
    ) -> PackResult<PackNoteMaybeAttachments> {
        let (PackNoteBase { id, note }, attachments_pack, poll_pack) = try_join!(
            self.extract_base(ctx, note_data),
            self.extract_attachments(ctx, drive_model, note_data.note()),
            self.extract_poll(ctx, note_data.note())
        )?;

        Ok(PackNoteMaybeAttachments::pack_from((
            id,
            note,
            Optional(self.extract_interaction(ctx, note_data)?),
            Optional(attachments_pack.map(|attachments| {
                NoteAttachmentExt::extract(
                    ctx,
                    NoteAttachmentSource {
                        attachments: &attachments,
                        poll: poll_pack.as_ref(),
                    },
                )
            })),
        )))
    }

    fn pack_full_single<'a>(
        &'a self,
        ctx: &'a PackingContext,
        note: &'a NoteData,
    ) -> BoxFuture<'a, PackResult<PackNoteMaybeFull>> {
        async move {
            let drive_model = DriveModel;

            let reply_target = async {
                match note.reply.as_ref() {
                    Some(r) if self.with_context => self.pack_full_single(ctx, r).await.map(Some),
                    _ => Ok(None),
                }
            };

            let renote_target = async {
                match note.renote.as_ref() {
                    Some(r) if self.with_context => self.pack_full_single(ctx, r).await.map(Some),
                    _ => Ok(None),
                }
            };

            let (
                PackNoteMaybeAttachments {
                    id,
                    note,
                    user_context,
                    attachment,
                },
                reply_target_pack,
                renote_target_pack,
            ) = try_join!(
                self.pack_single_attachments(ctx, &drive_model, &note),
                reply_target,
                renote_target
            )?;

            let detail = self.with_context.then(|| {
                NoteDetailExt::extract(
                    ctx,
                    NoteDetailSource {
                        parent_note: reply_target_pack.as_ref(),
                        renoted_note: renote_target_pack.as_ref(),
                    },
                )
            });

            Ok(PackNoteMaybeFull::pack_from((
                id,
                note,
                user_context,
                attachment,
                Optional(detail),
            )))
        }
        .boxed()
    }

    pub async fn fetch_single(
        &self,
        ctx: &PackingContext,
        id: &str,
    ) -> PackResult<Option<PackNoteMaybeFull>> {
        let note_resolver = ctx.service.db.get_note_resolver();
        let Some(note) = note_resolver
            .get_one(&NoteResolveOptions {
                ids: Some(vec![id.to_owned()]),
                visibility_filter: Box::new(NoteVisibilityFilterModel.new_note_visibility_filter(
                    ctx.self_user.as_deref().map(ck::user::Model::get_id),
                )),
                time_range: None,
                limit: None,
                with_user: self.with_context,
                with_reply_target: self.with_context,
                with_renote_target: self.with_context,
                with_interactions_from: self
                    .with_context
                    .then(|| {
                        ctx.self_user
                            .as_deref()
                            .map(ck::user::Model::get_id)
                            .map(str::to_string)
                    })
                    .flatten(),
                only_pins_from: None,
            })
            .await?
        else {
            return Ok(None);
        };

        Ok(Some(self.pack_full_single(ctx, &note).await?))
    }

    pub async fn fetch_pins(
        &self,
        ctx: &PackingContext,
        pin_user: &ck::user::Model,
    ) -> PackResult<Vec<PackNoteMaybeFull>> {
        let note_resolver = ctx.service.db.get_note_resolver();
        let notes = note_resolver
            .get_many(&NoteResolveOptions {
                ids: None,
                visibility_filter: Box::new(NoteVisibilityFilterModel.new_note_visibility_filter(
                    ctx.self_user.as_deref().map(ck::user::Model::get_id),
                )),
                time_range: None,
                limit: None,
                with_user: self.with_context,
                with_reply_target: self.with_context,
                with_renote_target: self.with_context,
                with_interactions_from: self
                    .with_context
                    .then(|| {
                        ctx.self_user
                            .as_deref()
                            .map(ck::user::Model::get_id)
                            .map(str::to_string)
                    })
                    .flatten(),
                only_pins_from: Some(pin_user.id.clone()),
            })
            .await?;

        let fut_iter = notes
            .iter()
            .map(|note| self.pack_full_single(ctx, note))
            .collect::<Vec<_>>();

        let processed = futures::stream::iter(fut_iter)
            .buffered(10)
            .err_into::<PackError>()
            .try_collect::<Vec<_>>()
            .await?;

        Ok(processed)
    }
}