macltok

use std::{

    collections::{HashMap, VecDeque},

    sync::Arc,

};

use futures_util::future::{join_all, try_join};

use matrix_sdk::{

    Client, Room, RoomState, deserialized_responses::TimelineEvent as SdkTimelineEvent,

    room::MessagesOptions,

};

use ruma::{OwnedEventId, OwnedRoomId, OwnedUserId, UInt, push::Action};

use serde_json::Value;

use tokio::sync::{Mutex, broadcast::error::RecvError};

use crate::{

    Error, ErrorKind, EventContext, EventContextCommand, EventContextPage, EventContextSession,

    MediaId, Sender, TimelineEvent, TimelineItem,

    engine::MediaService,

    normalise::{

        MediaRegistration, MediaScope, apply_profiles, normalise_context, normalise_observed,

        resolve_replies,

    },

};

const MAX_CONTEXT_LIMIT: u16 = 50;

/// a consumer that never closes its jump views still holds at most this many

const MAX_SESSIONS: usize = 32;

/// open event contexts, oldest first, each owning the media its pages registered

#[derive(Default)]

pub(crate) struct Sessions {

    open: Mutex<VecDeque<Session>>,

}

struct Session {

    id: String,

    room_id: OwnedRoomId,

    scope: MediaScope,

}

impl Sessions {

    async fn open(

        &self,

        media: &MediaService,

        room_id: OwnedRoomId,

        registrations: Vec<MediaRegistration>,

    ) -> String {

        let scope = media.context_scope();

        let id = format!("context-{}", scope.nonce());

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

        media.register(&scope, registrations).await;

        open.push_back(Session {

            id: id.clone(),

            room_id,

            scope,

        });

        while open.len() > MAX_SESSIONS {

            if let Some(oldest) = open.pop_front() {

                media.release(&oldest.scope).await;

            }

        }

        id

    }

    async fn room(&self, id: &str) -> Result<OwnedRoomId, Error> {

        let open = self.open.lock().await;

        let session = open.iter().find(|session| session.id == id);

        session

            .map(|session| session.room_id.clone())

            .ok_or_else(not_open)

    }

    /// under the lock, so a page that lands after its close registers nothing

    async fn register(

        &self,

        media: &MediaService,

        id: &str,

        registrations: Vec<MediaRegistration>,

    ) -> Result<(), Error> {

        let open = self.open.lock().await;

        let session = open

            .iter()

            .find(|session| session.id == id)

            .ok_or_else(not_open)?;

        media.register(&session.scope, registrations).await;

        Ok(())

    }

    async fn close(&self, media: &MediaService, id: &str) -> Result<(), Error> {

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

        let index = open

            .iter()

            .position(|session| session.id == id)

            .ok_or_else(not_open)?;

        if let Some(session) = open.remove(index) {

            media.release(&session.scope).await;

        }

        Ok(())

    }

    async fn close_left(&self, media: &MediaService, client: &Client) {

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

        let mut kept = VecDeque::with_capacity(open.len());

        for session in open.drain(..) {

            let left = client

                .get_room(&session.room_id)

                .is_none_or(|room| matches!(room.state(), RoomState::Left | RoomState::Banned));

            if left {

                media.release(&session.scope).await;

            } else {

                kept.push_back(session);

            }

        }

        *open = kept;

    }

    /// the media service is cleared beside this, so nothing is released one by one

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

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

    }

}

fn not_open() -> Error {

    Error::new(

        ErrorKind::EventContextNotOpen,

        "event context is closed or was never opened",

    )

}

/// leaving a room closes its event contexts

pub(super) async fn close_on_leave(

    client: Client,

    media: Arc<MediaService>,

    sessions: Arc<Sessions>,

) {

    let mut info = client.room_info_notable_update_receiver();

    loop {

        match info.recv().await {

            Ok(_) | Err(RecvError::Lagged(_)) => sessions.close_left(&media, &client).await,

            Err(RecvError::Closed) => return,

        }

    }

}

pub(crate) async fn execute(

    client: Client,

    media: Arc<MediaService>,

    sessions: Arc<Sessions>,

    command: EventContextCommand,

) -> Result<crate::CommandOutput, Error> {

    let current_user_id = client.user_id().map(ToOwned::to_owned);

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

    let (context_id, token, limit, backwards) = match command {

        EventContextCommand::Fetch {

            room_id,

            event_id,

            before_limit,

            after_limit,

        } => {

            let (context, registrations, _) = capture(

                room(&room_id)?,

                current_user_id,

                event_id,

                before_limit,

                after_limit,

            )

            .await?;

            let context_id = sessions.open(&media, room_id, registrations).await;

            return Ok(crate::CommandOutput::EventContext(EventContextSession {

                context_id,

                context,

            }));

        }

        EventContextCommand::Close { context_id } => {

            sessions.close(&media, &context_id).await?;

            return Ok(crate::CommandOutput::Done);

        }

        EventContextCommand::Before {

            context_id,

            token,

            limit,

        } => (context_id, token, limit, true),

        EventContextCommand::After {

            context_id,

            token,

            limit,

        } => (context_id, token, limit, false),

    };

    if limit == 0 {

        return Err(Error::new(

            ErrorKind::Invalid,

            "context page limit must be at least 1",

        ));

    }

    validate_limit(limit)?;

    if token.is_empty() {

        return Err(Error::new(

            ErrorKind::Invalid,

            "context pagination token is empty",

        ));

    }

    let room = room(&sessions.room(&context_id).await?)?;

    let side = fetch_side(&room, Some(&token), limit, backwards).await?;

    let mut page = Assembly::new(

        &room,

        current_user_id.as_deref(),

        &side.values,

        &side.highlights,

        None,

    )

    .await;

    let events = page.take(0..side.values.len());

    sessions.register(&media, &context_id, page.media).await?;

    Ok(crate::CommandOutput::EventContextPage(EventContextPage {

        events,

        next_token: side.token,

    }))

}

pub(crate) async fn capture(

    room: Room,

    current_user_id: Option<OwnedUserId>,

    event_id: OwnedEventId,

    before_limit: u16,

    after_limit: u16,

) -> Result<(EventContext, Vec<MediaRegistration>, Vec<Value>), Error> {

    validate_limit(before_limit)?;

    validate_limit(after_limit)?;

    let raw = room

        .event_with_context(&event_id, false, UInt::from(0_u16), None)

        .await

        .map_err(|source| {

            let error = Error::from_sdk("fetching event context failed", source);

            if error

                .server

                .as_ref()

                .is_some_and(|server| server.status == 404)

            {

                Error::new(ErrorKind::EventNotFound, "event was not found")

            } else {

                error

            }

        })?;

    let (target, target_highlight) = raw

        .event

        .as_ref()

        .and_then(raw_event)

        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event was not found"))?;

    let (before, after) = try_join(

        fetch_side(&room, raw.prev_batch_token.as_deref(), before_limit, true),

        fetch_side(&room, raw.next_batch_token.as_deref(), after_limit, false),

    )

    .await?;

    let target_index = before.values.len();

    let mut values = before.values;

    values.push(target);

    values.extend(after.values);

    let mut highlights = before.highlights;

    highlights.push(target_highlight);

    highlights.extend(after.highlights);

    let mut page = Assembly::new(

        &room,

        current_user_id.as_deref(),

        &values,

        &highlights,

        Some(target_index),

    )

    .await;

    let events_before = page.take(0..target_index);

    let event = page

        .take(target_index..target_index + 1)

        .pop()

        .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "event could not be read"))?;

    let events_after = page.take(target_index + 1..values.len());

    Ok((

        EventContext {

            events_before,

            event,

            events_after,

            start_token: before.token,

            end_token: after.token,

        },

        page.media,

        values,

    ))

}

struct Side {

    /// in timeline order

    values: Vec<Value>,

    /// beside `values`, from the push actions the sdk worked out as it fetched

    highlights: Vec<bool>,

    token: Option<String>,

}

fn raw_event(event: &SdkTimelineEvent) -> Option<(Value, bool)> {

    let value = serde_json::from_str(event.raw().json().get()).ok()?;

    let highlight = event

        .push_actions()

        .is_some_and(|actions| actions.iter().any(Action::is_highlight));

    Some((value, highlight))

}

async fn fetch_side(

    room: &Room,

    token: Option<&str>,

    limit: u16,

    backwards: bool,

) -> Result<Side, Error> {

    let Some(token) = token else {

        return Ok(Side {

            values: Vec::new(),

            highlights: Vec::new(),

            token: None,

        });

    };

    if limit == 0 {

        return Ok(Side {

            values: Vec::new(),

            highlights: Vec::new(),

            token: Some(token.to_owned()),

        });

    }

    let mut options = if backwards {

        MessagesOptions::backward()

    } else {

        MessagesOptions::forward()

    };

    options = options.from(token);

    options.limit = UInt::from(limit);

    let messages = room

        .messages(options)

        .await

        .map_err(|source| Error::from_sdk("fetching event context events failed", source))?;

    let mut events: Vec<_> = messages.chunk.iter().filter_map(raw_event).collect();

    if backwards {

        events.reverse();

    }

    let (values, highlights) = events.into_iter().unzip();

    Ok(Side {

        values,

        highlights,

        token: messages.end,

    })

}

/// a window of raw events normalised together, so relations inside it land on their targets

struct Assembly {

    items: Vec<(usize, TimelineItem)>,

    media: Vec<MediaRegistration>,

}

impl Assembly {

    async fn new(

        room: &Room,

        current_user_id: Option<&ruma::UserId>,

        values: &[Value],

        highlights: &[bool],

        keep: Option<usize>,

    ) -> Self {

        let page = normalise_context(values, keep, current_user_id, room.room_id());

        let mut items = page.items;

        for (index, item) in &mut items {

            if let TimelineItem::Event(event) = item {

                event.highlight = highlights[*index];

            }

        }

        let mut media = page.media;

        let mut targets: HashMap<OwnedEventId, TimelineEvent> = HashMap::new();

        let mut missing = Vec::new();

        for (_, item) in &items {

            if let TimelineItem::Event(event) = item

                && let Some(event_id) = &event.event_id

            {

                targets.insert(event_id.clone(), (**event).clone());

            }

        }

        for event in events(&mut items) {

            if let Some(reply) = &event.reply_to

                && !targets.contains_key(&reply.event_id)

                && !missing.contains(&reply.event_id)

            {

                missing.push(reply.event_id.clone());

            }

        }

        let fetched = join_all(

            missing

                .into_iter()

                .map(|event_id| fetch_reply(room, current_user_id, event_id)),

        )

        .await;

        for (event, registrations) in fetched.into_iter().flatten() {

            media.extend(registrations);

            if let Some(event_id) = event.event_id.clone() {

                targets.insert(event_id, event);

            }

        }

        let profiles = profiles(

            room,

            items

                .iter()

                .filter_map(|(_, item)| match item {

                    TimelineItem::Event(event) => Some(&event.sender.user_id),

                    _ => None,

                })

                .chain(targets.values().map(|event| &event.sender.user_id)),

        )

        .await;

        apply_profiles(targets.values_mut(), &profiles);

        apply_profiles(events(&mut items), &profiles);

        resolve_replies(events(&mut items), &targets);

        Self { items, media }

    }

    /// the items that came from raw events in `range`

    fn take(&mut self, range: std::ops::Range<usize>) -> Vec<TimelineItem> {

        let (taken, kept) = std::mem::take(&mut self.items)

            .into_iter()

            .partition(|(index, _)| range.contains(index));

        self.items = kept;

        taken.into_iter().map(|(_, item)| item).collect()

    }

}

fn events(items: &mut [(usize, TimelineItem)]) -> impl Iterator<Item = &mut TimelineEvent> {

    items.iter_mut().filter_map(|(_, item)| match item {

        TimelineItem::Event(event) => Some(&mut **event),

        _ => None,

    })

}

async fn fetch_reply(

    room: &Room,

    current_user_id: Option<&ruma::UserId>,

    event_id: OwnedEventId,

) -> Option<(TimelineEvent, Vec<MediaRegistration>)> {

    let event = room

        .load_or_fetch_event(&event_id, None)

        .await

        .inspect_err(|error| tracing::debug!(%event_id, %error, "fetching a reply target failed"))

        .ok()?;

    let value: Value = serde_json::from_str(event.raw().json().get()).ok()?;

    let identity = format!("{}\n{event_id}", room.room_id());

    match normalise_observed(&value, current_user_id, &identity)? {

        (TimelineItem::Event(event), media) => Some((*event, media)),

        _ => None,

    }

}

/// under lazy-loaded members a sender may be unknown locally, and `get_member` fetches the

/// member list once for the room when it is

async fn profiles<'a>(

    room: &Room,

    user_ids: impl Iterator<Item = &'a OwnedUserId>,

) -> HashMap<OwnedUserId, Sender> {

    let mut unique: Vec<&OwnedUserId> = user_ids.collect();

    unique.sort();

    unique.dedup();

    join_all(unique.into_iter().map(|user_id| async move {

        let member = room.get_member(user_id).await.ok().flatten()?;

        Some((

            user_id.clone(),

            Sender {

                user_id: user_id.clone(),

                display_name: member.display_name().map(str::to_owned),

                avatar: member.avatar_url().map(MediaId::plain),

                name_ambiguous: member.name_ambiguous(),

            },

        ))

    }))

    .await

    .into_iter()

    .flatten()

    .collect()

}

fn validate_limit(limit: u16) -> Result<(), Error> {

    if limit > MAX_CONTEXT_LIMIT {

        Err(Error::new(ErrorKind::Invalid, "context limit exceeds 50"))

    } else {

        Ok(())

    }

}