macltok

use std::collections::{BTreeMap, HashSet};

use matrix_sdk::{

    Room, RoomState, cross_process_lock::MappedCrossProcessLockState,

    deserialized_responses::TimelineEvent,

};

use matrix_sdk_base::event_cache::store::EventCacheStoreLockGuard;

use ruma::{

    OwnedEventId,

    events::{

        receipt::{ReceiptThread, ReceiptType},

        relation::RelationType,

    },

};

use serde_json::{Value, json};

use crate::ThreadUnread;

fn key(room: &Room) -> Vec<u8> {

    format!("macltok.thread_unreads.{}", room.room_id()).into_bytes()

}

/// what the last recount left, so a room list never touches the event cache

pub(crate) async fn stored_thread_unreads(room: &Room) -> BTreeMap<OwnedEventId, ThreadUnread> {

    if room.state() != RoomState::Joined {

        return BTreeMap::new();

    }

    match room.client().state_store().get_custom_value(&key(room)).await {

        Ok(Some(bytes)) => decode(&bytes),

        Ok(None) => BTreeMap::new(),

        Err(error) => {

            tracing::debug!(room_id = %room.room_id(), %error, "reading thread unreads failed");

            BTreeMap::new()

        }

    }

}

/// the sdk leaves thread events out of the room's counts once threads are on and keeps none

/// of its own for them. this counts `touched` and every thread that still had something

/// unread from the event cache store, which holds every reply it ever took in whatever is

/// in memory, against a receipt for that thread (msc3771), an unthreaded one or the user's

/// own reply. `None` when nothing moved

pub(crate) async fn recount_thread_unreads(

    room: &Room,

    touched: HashSet<OwnedEventId>,

) -> Option<BTreeMap<OwnedEventId, ThreadUnread>> {

    let known = stored_thread_unreads(room).await;

    let mut roots = touched;

    roots.extend(known.keys().cloned());

    if roots.is_empty() || room.state() != RoomState::Joined {

        return None;

    }

    let store = match room.client().event_cache_store().lock().await {

        Ok(MappedCrossProcessLockState::Clean(store) | MappedCrossProcessLockState::Dirty(store)) => {

            store

        }

        Err(error) => {

            tracing::debug!(room_id = %room.room_id(), %error, "locking the event cache store failed");

            return None;

        }

    };

    let own_user_id = room.own_user_id();

    let unthreaded = read_through(room, &store, ReceiptThread::Unthreaded).await;

    let mut counted = BTreeMap::new();

    for root in roots {

        let replies = match store

            .find_event_relations(room.room_id(), &root, Some(&[RelationType::Thread][..]))

            .await

        {

            Ok(replies) => replies,

            Err(error) => {

                tracing::debug!(room_id = %room.room_id(), %root, %error, "reading a thread failed");

                if let Some(unread) = known.get(&root) {

                    counted.insert(root, *unread);

                }

                continue;

            }

        };

        let mut read = unthreaded.max(read_through(room, &store, ReceiptThread::Thread(root.clone())).await);

        for (reply, _) in &replies {

            if reply.sender().as_deref() == Some(own_user_id) {

                read = read.max(timestamp_ms(reply));

            }

        }

        let mut unread = ThreadUnread::default();

        for (reply, _) in &replies {

            let Some(at) = timestamp_ms(reply) else {

                continue;

            };

            if read.is_some_and(|read| at <= read)

                || reply.sender().as_deref() == Some(own_user_id)

                || !counts_as_unread(reply)

            {

                continue;

            }

            let actions = reply.push_actions().unwrap_or_default();

            unread.messages += 1;

            unread.notifications += u64::from(actions.iter().any(|action| action.should_notify()));

            unread.mentions += u64::from(actions.iter().any(|action| action.is_highlight()));

        }

        if unread != ThreadUnread::default() {

            counted.insert(root, unread);

        }

    }

    if counted == known {

        return None;

    }

    let stored = room

        .client()

        .state_store()

        .set_custom_value(&key(room), encode(&counted))

        .await;

    if let Err(error) = stored {

        tracing::debug!(room_id = %room.room_id(), %error, "saving thread unreads failed");

    }

    Some(counted)

}

/// receipts sit on events, and stream order between a thread and the rest of the room is

/// gone once it's in the store, so the event's timestamp stands for its place. a receipt

/// on an event the store never had falls back to when it was sent

async fn read_through(

    room: &Room,

    store: &EventCacheStoreLockGuard,

    thread: ReceiptThread,

) -> Option<u64> {

    let mut latest = None;

    for receipt_type in [ReceiptType::Read, ReceiptType::ReadPrivate] {

        let (event_id, receipt) = match room

            .load_user_receipt(receipt_type, thread.clone(), room.own_user_id())

            .await

        {

            Ok(Some(found)) => found,

            Ok(None) => continue,

            Err(error) => {

                tracing::debug!(room_id = %room.room_id(), %error, "reading a receipt failed");

                continue;

            }

        };

        let at = match store.find_event(room.room_id(), &event_id).await {

            Ok(Some(event)) => timestamp_ms(&event),

            _ => None,

        }

        .or_else(|| receipt.ts.map(|ts| u64::from(ts.get())));

        latest = latest.max(at);

    }

    latest

}

fn timestamp_ms(event: &TimelineEvent) -> Option<u64> {

    event.timestamp().map(|ts| u64::from(ts.get()))

}

/// the sdk's `marks_as_unread` for its room counts, which is private to it

fn counts_as_unread(event: &TimelineEvent) -> bool {

    let raw = event.raw();

    let redacted = raw

        .get_field::<Value>("unsigned")

        .ok()

        .flatten()

        .is_some_and(|unsigned| unsigned.get("redacted_because").is_some());

    let edit = raw

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

        .ok()

        .flatten()

        .and_then(|content| {

            content

                .pointer("/m.relates_to/rel_type")

                .and_then(Value::as_str)

                .map(|rel_type| rel_type == "m.replace")

        })

        .unwrap_or(false);

    let message_like = matches!(

        raw.get_field::<String>("type").ok().flatten().as_deref(),

        Some(

            "m.message"

                | "m.room.message"

                | "m.room.encrypted"

                | "m.sticker"

                | "m.poll.start"

                | "org.matrix.msc3381.poll.start"

                | "m.poll.end"

                | "org.matrix.msc3381.poll.end"

        )

    );

    message_like && !edit && !redacted

}

fn encode(unreads: &BTreeMap<OwnedEventId, ThreadUnread>) -> Vec<u8> {

    let value: serde_json::Map<String, Value> = unreads

        .iter()

        .map(|(root, unread)| {

            (

                root.to_string(),

                json!([unread.messages, unread.mentions, unread.notifications]),

            )

        })

        .collect();

    serde_json::to_vec(&value).unwrap_or_default()

}

fn decode(bytes: &[u8]) -> BTreeMap<OwnedEventId, ThreadUnread> {

    let Ok(value) = serde_json::from_slice::<BTreeMap<OwnedEventId, [u64; 3]>>(bytes) else {

        return BTreeMap::new();

    };

    value

        .into_iter()

        .map(|(root, [messages, mentions, notifications])| {

            (

                root,

                ThreadUnread {

                    messages,

                    mentions,

                    notifications,

                },

            )

        })

        .collect()

}