macltok

use std::collections::{HashMap, VecDeque};

use matrix_sdk::{Room, RoomState};

use matrix_sdk_ui::timeline::{RoomExt, Timeline, TimelineFocus, TimelineReadReceiptTracking};

use ruma::{OwnedEventId, OwnedRoomId};

use tokio::{

    sync::mpsc,

    task::{AbortHandle, JoinSet},

};

use super::timeline::{SessionContext, TimelineSession};

use crate::{

    Error, ErrorKind, Event, TimelineCommand, TimelineId, TypingState, transport::Responder,

};

pub(super) enum ActorRequest {

    /// opening an open room answers with a fresh reset

    Open(Responder),

    Execute(TimelineCommand, Responder),

    Close(Responder),

}

impl ActorRequest {

    pub(super) fn into_reply(self) -> Responder {

        match self {

            Self::Open(reply) | Self::Execute(_, reply) | Self::Close(reply) => reply,

        }

    }

}

pub(super) struct Exited {

    pub(super) room_id: OwnedRoomId,

    pub(super) generation: u64,

}

pub(super) struct RoomActor {

    pub(super) room_id: OwnedRoomId,

    pub(super) generation: u64,

    pub(super) room: Room,

    pub(super) requests: mpsc::Receiver<ActorRequest>,

    pub(super) exits: mpsc::UnboundedSender<Exited>,

    pub(super) context: SessionContext,

    pub(super) panics: mpsc::UnboundedReceiver<&'static str>,

    pub(super) opening: Vec<Responder>,

}

impl RoomActor {

    pub(super) async fn run(mut self) {

        let exits = self.exits.clone();

        let exited = Exited {

            room_id: self.room_id.clone(),

            generation: self.generation,

        };

        self.serve().await;

        let _ = exits.send(exited);

    }

    async fn serve(&mut self) {

        let timeline = match build_timeline(

            &self.room,

            TimelineFocus::Live {

                hide_threaded_events: true,

            },

        )

        .await

        {

            Ok(timeline) => timeline,

            Err(error) => return self.refuse_everything(error).await,

        };

        let Ok(room_timeline) = TimelineSession::start(

            TimelineId::Room(self.room_id.clone()),

            timeline,

            self.context.clone(),

        )

        .await

        else {

            return;

        };

        for reply in self.opening.drain(..) {

            reply.done(Ok(()));

        }

        let (typing_guard, typing) = self.room.subscribe_to_typing_notifications();

        let typing_watcher = self.context.supervisor.spawn(

            "typing watcher",

            watch_typing(self.room_id.clone(), typing, self.context.events.clone()),

        );

        let mut timelines = Timelines {

            room: room_timeline,

            threads: Threads::new(self.requests.max_capacity()),

            work: JoinSet::new(),

        };

        let close = self.serve_open(&mut timelines).await;

        // let redactions and receipts that already left finish rather than lie about them

        while timelines.work.join_next().await.is_some() {}

        drop(typing_guard);

        typing_watcher.abort();

        timelines.threads.close().await;

        timelines.room.close().await;

        if let Some(reply) = close {

            reply.done(Ok(()));

        }

    }

    /// returns the close request's responder, if a close is why we stopped

    async fn serve_open(&mut self, timelines: &mut Timelines) -> Option<Responder> {

        let mut room_info = self.room.subscribe_info();

        loop {

            tokio::select! {

                Some(info) = room_info.next() => {

                    if matches!(info.state(), RoomState::Left | RoomState::Banned) {

                        tracing::debug!(room_id = %self.room_id, "room actor stopped because we left the room");

                        return None;

                    }

                }

                request = self.requests.recv() => match request {

                    None => return None,

                    Some(ActorRequest::Close(reply)) => return Some(reply),

                    Some(ActorRequest::Open(reply)) => reply.done(timelines.room.resync().await),

                    Some(ActorRequest::Execute(command, reply)) => {

                        timelines.execute(command, reply, &self.room).await;

                    }

                },

                Some(built) = timelines.threads.builds.join_next(), if !timelines.threads.builds.is_empty() => {

                    timelines.thread_built(built, &self.room_id, &self.context).await;

                }

                Some(label) = self.panics.recv() => panic!("{label} panicked in room {}", self.room_id),

                Some(done) = timelines.work.join_next(), if !timelines.work.is_empty() => {

                    if let Err(error) = done && error.is_panic() {

                        tracing::warn!(room_id = %self.room_id, %error, "timeline operation panicked");

                    }

                }

            }

        }

    }

    async fn refuse_everything(&mut self, error: matrix_sdk_ui::timeline::Error) {

        tracing::warn!(%error, room_id = %self.room_id, "room actor could not build its timeline");

        let error = Error::with_source(ErrorKind::Timeline, "room timeline build failed", error);

        for reply in self.opening.drain(..) {

            reply.done(Err(error.clone()));

        }

        self.requests.close();

        while let Some(request) = self.requests.recv().await {

            request.into_reply().done(Err(error.clone()));

        }

    }

}

async fn build_timeline(

    room: &Room,

    focus: TimelineFocus,

) -> Result<Timeline, matrix_sdk_ui::timeline::Error> {

    room.timeline_builder()

        .with_focus(focus)

        .track_read_marker_and_receipts(TimelineReadReceiptTracking::AllEvents)

        .build()

        .await

}

async fn watch_typing(

    room_id: OwnedRoomId,

    mut typing: tokio::sync::broadcast::Receiver<Vec<ruma::OwnedUserId>>,

    events: crate::transport::EventSender,

) {

    loop {

        match typing.recv().await {

            Ok(users) => {

                let state = TypingState {

                    room_id: room_id.clone(),

                    users,

                };

                if events.send(Event::TypingChanged(state)).await.is_err() {

                    return;

                }

            }

            Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {

                tracing::warn!(%room_id, skipped, "typing watcher lagged");

            }

            Err(tokio::sync::broadcast::error::RecvError::Closed) => return,

        }

    }

}

struct Timelines {

    room: TimelineSession,

    threads: Threads,

    work: JoinSet<()>,

}

impl Timelines {

    async fn execute(&mut self, command: TimelineCommand, reply: Responder, room: &Room) {

        match command {

            TimelineCommand::OpenThread { root_event_id, .. } => {

                self.threads.open(root_event_id, reply, room);

            }

            TimelineCommand::CloseThread { root_event_id, .. } => {

                reply.done(self.threads.close_one(&root_event_id).await);

            }

            TimelineCommand::Resync {

                timeline_id: TimelineId::Room(_),

            } => reply.done(self.room.resync().await),

            command => {

                let root = command

                    .timeline_id()

                    .and_then(TimelineId::thread_root)

                    .cloned();

                let Some(root) = root else {

                    return self.room.execute(command, reply, &mut self.work).await;

                };

                match self.threads.route(root, command, reply) {

                    Routed::Open(session, TimelineCommand::Resync { .. }, reply) => {

                        reply.done(session.resync().await);

                    }

                    Routed::Open(session, command, reply) => {

                        session.execute(command, reply, &mut self.work).await;

                    }

                    Routed::Queued => {}

                }

            }

        }

    }

    async fn thread_built(

        &mut self,

        built: Result<Built, tokio::task::JoinError>,

        room_id: &OwnedRoomId,

        context: &SessionContext,

    ) {

        let Some((root, queued, timeline)) = self.threads.built(built) else {

            return;

        };

        let timeline = match timeline {

            Ok(timeline) => timeline,

            Err(error) => {

                tracing::warn!(%error, %room_id, %root, "thread timeline build failed");

                let error =

                    Error::with_source(ErrorKind::Timeline, "thread timeline build failed", error);

                for (_, reply) in queued {

                    reply.done(Err(error.clone()));

                }

                return;

            }

        };

        let id = TimelineId::thread(room_id.clone(), root.clone());

        let Ok(session) = TimelineSession::start(id, timeline, context.clone()).await else {

            return;

        };

        let session = self

            .threads

            .open

            .entry(root)

            .insert_entry(session)

            .into_mut();

        for (command, reply) in queued {

            match command {

                None => reply.done(Ok(())),

                Some(TimelineCommand::Resync { .. }) => reply.done(session.resync().await),

                Some(command) => session.execute(command, reply, &mut self.work).await,

            }

        }

    }

}

type Built = (

    OwnedEventId,

    u64,

    Result<Timeline, matrix_sdk_ui::timeline::Error>,

);

/// `None` is the open itself, everything else waits behind it

type Queued = VecDeque<(Option<TimelineCommand>, Responder)>;

enum Routed<'a> {

    Open(&'a mut TimelineSession, TimelineCommand, Responder),

    Queued,

}

struct OpeningThread {

    generation: u64,

    build: AbortHandle,

    queued: Queued,

}

struct Threads {

    open: HashMap<OwnedEventId, TimelineSession>,

    opening: HashMap<OwnedEventId, OpeningThread>,

    builds: JoinSet<Built>,

    next_generation: u64,

    queued: usize,

    capacity: usize,

}

impl Threads {

    fn new(capacity: usize) -> Self {

        Self {

            open: HashMap::new(),

            opening: HashMap::new(),

            builds: JoinSet::new(),

            next_generation: 0,

            queued: 0,

            capacity,

        }

    }

    fn enqueue(&mut self, root: &OwnedEventId, command: Option<TimelineCommand>, reply: Responder) {

        if self.queued >= self.capacity {

            return reply.done(Err(Error::new(

                ErrorKind::QueueFull,

                "thread command queue is full",

            )));

        }

        let Some(opening) = self.opening.get_mut(root) else {

            return reply.done(Err(Error::new(

                ErrorKind::TimelineNotOpen,

                "thread timeline is not open",

            )));

        };

        opening.queued.push_back((command, reply));

        self.queued += 1;

    }

    fn open(&mut self, root: OwnedEventId, reply: Responder, room: &Room) {

        if self.open.contains_key(&root) {

            return reply.done(Ok(()));

        }

        if !self.opening.contains_key(&root) {

            let generation = self.next_generation;

            self.next_generation = self.next_generation.wrapping_add(1);

            let room = room.clone();

            let build_root = root.clone();

            let build = self.builds.spawn(async move {

                let focus = TimelineFocus::Thread {

                    root_event_id: build_root.clone(),

                };

                (build_root, generation, build_timeline(&room, focus).await)

            });

            self.opening.insert(

                root.clone(),

                OpeningThread {

                    generation,

                    build,

                    queued: VecDeque::new(),

                },

            );

        }

        self.enqueue(&root, None, reply);

    }

    fn route(

        &mut self,

        root: OwnedEventId,

        command: TimelineCommand,

        reply: Responder,

    ) -> Routed<'_> {

        if !self.open.contains_key(&root) {

            self.enqueue(&root, Some(command), reply);

            return Routed::Queued;

        }

        let session = self.open.get_mut(&root).expect("thread was just checked");

        Routed::Open(session, command, reply)

    }

    /// `None` when the build belongs to a thread that was closed while it ran

    fn built(

        &mut self,

        built: Result<Built, tokio::task::JoinError>,

    ) -> Option<(

        OwnedEventId,

        Queued,

        Result<Timeline, matrix_sdk_ui::timeline::Error>,

    )> {

        let (root, generation, timeline) = match built {

            Ok(built) => built,

            Err(error) if error.is_cancelled() => return None,

            Err(error) => std::panic::resume_unwind(error.into_panic()),

        };

        if self

            .opening

            .get(&root)

            .is_none_or(|opening| opening.generation != generation)

        {

            return None;

        }

        let opening = self.opening.remove(&root)?;

        self.queued -= opening.queued.len();

        Some((root, opening.queued, timeline))

    }

    async fn close_one(&mut self, root: &OwnedEventId) -> Result<(), Error> {

        if let Some(session) = self.open.remove(root) {

            session.close().await;

            return Ok(());

        }

        let Some(opening) = self.opening.remove(root) else {

            return Err(Error::new(

                ErrorKind::TimelineNotOpen,

                "thread timeline is not open",

            ));

        };

        opening.build.abort();

        self.queued -= opening.queued.len();

        Ok(())

    }

    async fn close(&mut self) {

        self.builds.abort_all();

        while self.builds.join_next().await.is_some() {}

        self.opening.clear();

        for (_, session) in self.open.drain() {

            session.close().await;

        }

    }

}