macltok

use matrix_sdk::{

    Room,

    event_cache::{EventCacheError, HistoryGapFillOutcome, HistoryGapStatus},

};

use ruma::{OwnedRoomId, OwnedTransactionId};

use tokio::sync::watch;

use super::{

    executor::{JoinError, JoinHandle, spawn},

    timeline::SessionContext,

};

use crate::{

    CommandOutput, Error, ErrorKind, Event, HistoryContinuity, HistoryRepairCommand,

    HistoryRepairPhase, HistoryRepairReport, HistoryRepairStop, RecentHistoryState,

    transport::Responder,

};

const MAX_GAP_PAGES: u32 = 4;

struct Running {

    task: JoinHandle<Finished>,

    cancel: watch::Sender<bool>,

    responders: Vec<PendingResponder>,

    gap_id: Option<OwnedTransactionId>,

    closing: bool,

}

struct PendingResponder {

    responder: Responder,

    cancel_ack: bool,

}

pub(super) enum Outcome {

    Recent(RecentHistoryState),

    Fill(HistoryRepairReport),

}

pub(super) struct Completion {

    pub(super) outcome: Outcome,

    /// history went into the cache, so the room timeline has to show it before anyone hears

    pub(super) committed: bool,

    pub(super) closing: bool,

    pub(super) result: Result<CommandOutput, Error>,

    responders: Vec<PendingResponder>,

}

impl Completion {

    /// the room timeline couldn't show what was committed

    pub(super) fn unpublished(&mut self, error: Error) {

        match &mut self.outcome {

            Outcome::Recent(state) => *state = RecentHistoryState::Failed(error.message.clone()),

            Outcome::Fill(report) => {

                report.phase = HistoryRepairPhase::Paused;

                report.failure = Some(error.clone());

            }

        }

        self.result = Err(error);

    }

    pub(super) fn resolve(

        self,

        result: Result<CommandOutput, Error>,

        cancellation: Result<(), Error>,

    ) {

        for responder in self.responders {

            if responder.cancel_ack {

                responder.responder.done(cancellation.clone());

            } else {

                responder.responder.send(result.clone());

            }

        }

    }

}

enum Stopped {

    Done,

    Cancelled,

    Failed(Error),

}

struct Finished {

    stopped: Stopped,

    committed: bool,

    report: Option<HistoryRepairReport>,

}

pub(super) struct RepairController {

    room: Room,

    context: SessionContext,

    running: Option<Running>,

}

impl RepairController {

    pub(super) fn new(room: Room, context: SessionContext) -> Self {

        Self {

            room,

            context,

            running: None,

        }

    }

    pub(super) fn is_running(&self) -> bool {

        self.running.is_some()

    }

    /// a room with retained events already shows the live end, and a limited sync has put its

    /// own gap there, so only an empty room walks back from the live end. otherwise the newest

    /// gap gets one page, so a limited sync that skipped nothing leaves no gap behind

    pub(super) async fn start_recent(&mut self) {

        let retained = async {

            let (cache, _handles) = self.room.event_cache().await?;

            if !cache.has_retained_events().await? {

                return Ok(None);

            }

            let gaps = cache.history_gaps().await?;

            Ok::<_, EventCacheError>(Some(gaps.into_iter().rev().find(|gap| {

                matches!(

                    gap.status,

                    HistoryGapStatus::Idle | HistoryGapStatus::Failed { .. }

                )

            })))

        }

        .await;

        match retained {

            Ok(Some(Some(gap))) => self.start_gap_fill(gap.id, None, 1, false).await,

            Ok(Some(None)) => {}

            Ok(None) | Err(_) => self.start_recent_with(None).await,

        }

    }

    pub(super) async fn execute(&mut self, command: HistoryRepairCommand, reply: Responder) {

        match command {

            HistoryRepairCommand::Cancel { .. } => {

                if let Some(running) = self.running.as_mut() {

                    running.responders.push(PendingResponder {

                        responder: reply,

                        cancel_ack: true,

                    });

                    running.cancel.send_replace(true);

                } else {

                    reply.done(Ok(()));

                }

            }

            HistoryRepairCommand::RetryRecent { .. } => self.start_recent_with(Some(reply)).await,

            HistoryRepairCommand::FillGap { gap_id, .. } => {

                self.start_gap_fill(gap_id, Some(reply), MAX_GAP_PAGES, true)

                    .await;

            }

        }

    }

    async fn start_recent_with(&mut self, responder: Option<Responder>) {

        if let Some(running) = &mut self.running {

            if let Some(responder) = responder {

                if running.gap_id.is_none() {

                    running.responders.push(PendingResponder {

                        responder,

                        cancel_ack: false,

                    });

                } else {

                    responder.done(Err(room_busy()));

                }

            }

            return;

        }

        let event = Event::RecentHistoryChanged {

            room_id: self.room.room_id().to_owned(),

            state: RecentHistoryState::Loading,

        };

        if self.context.events.send(event).await.is_err() {

            if let Some(responder) = responder {

                responder.done(Err(Error::interrupted()));

            }

            return;

        }

        let (cancel, receiver) = watch::channel(false);

        let room = self.room.clone();

        let task = spawn(async move { load_recent(room, receiver).await });

        self.running = Some(Running {

            task,

            cancel,

            responders: responder

                .map(|responder| PendingResponder {

                    responder,

                    cancel_ack: false,

                })

                .into_iter()

                .collect(),

            gap_id: None,

            closing: false,

        });

    }

    async fn start_gap_fill(

        &mut self,

        gap_id: OwnedTransactionId,

        responder: Option<Responder>,

        pages: u32,

        spend_budget: bool,

    ) {

        if self.is_running() {

            if let Some(responder) = responder {

                responder.done(Err(room_busy()));

            }

            return;

        }

        let room = self.room.clone();

        let progress = report(self.room.room_id().to_owned(), gap_id.clone());

        if self

            .context

            .events

            .send(Event::HistoryRepairChanged(progress.clone()))

            .await

            .is_err()

        {

            if let Some(responder) = responder {

                responder.done(Err(Error::interrupted()));

            }

            return;

        }

        let (cancel, receiver) = watch::channel(false);

        let task =

            spawn(async move { fill_gap(room, progress, receiver, pages, spend_budget).await });

        self.running = Some(Running {

            task,

            cancel,

            responders: responder

                .map(|responder| PendingResponder {

                    responder,

                    cancel_ack: false,

                })

                .into_iter()

                .collect(),

            gap_id: Some(gap_id),

            closing: false,

        });

    }

    pub(super) async fn next_completion(&mut self) -> Option<Completion> {

        let running = self.running.as_mut()?;

        let result = (&mut running.task).await;

        Some(self.finish_task(result))

    }

    pub(super) async fn shutdown(&mut self) -> Option<Completion> {

        let running = self.running.as_mut()?;

        running.closing = true;

        running.cancel.send_replace(true);

        let result = (&mut running.task).await;

        Some(self.finish_task(result))

    }

    fn finish_task(&mut self, result: Result<Finished, JoinError>) -> Completion {

        let running = self.running.take().expect("repair task exists");

        let finished = result.unwrap_or(Finished {

            stopped: Stopped::Failed(Error::interrupted()),

            committed: false,

            report: None,

        });

        let result = match &finished.stopped {

            Stopped::Done => Ok(CommandOutput::Done),

            Stopped::Cancelled => Err(Error::interrupted()),

            Stopped::Failed(error) => Err(error.clone()),

        };

        let outcome = match running.gap_id {

            None => Outcome::Recent(match finished.stopped {

                _ if running.closing => RecentHistoryState::Idle,

                Stopped::Done => RecentHistoryState::Loaded,

                Stopped::Cancelled => RecentHistoryState::Cancelled,

                Stopped::Failed(error) => RecentHistoryState::Failed(error.message),

            }),

            Some(gap_id) => {

                let mut report = finished

                    .report

                    .unwrap_or_else(|| report(self.room.room_id().to_owned(), gap_id));

                match finished.stopped {

                    Stopped::Done => {}

                    Stopped::Cancelled => report.stop = Some(HistoryRepairStop::Cancelled),

                    Stopped::Failed(error) => report.failure = Some(error),

                }

                report.phase = if report.stop.is_none() && report.failure.is_none() {

                    HistoryRepairPhase::Ready

                } else {

                    HistoryRepairPhase::Paused

                };

                Outcome::Fill(report)

            }

        };

        Completion {

            outcome,

            committed: finished.committed,

            closing: running.closing,

            result,

            responders: running.responders,

        }

    }

}

async fn load_recent(room: Room, mut cancel: watch::Receiver<bool>) -> Finished {

    let outcome = tokio::select! {

        biased;

        _ = cancel.changed() => None,

        result = async {

            let (cache, _handles) = room.event_cache().await?;

            cache.load_recent_history().await

        } => Some(result),

    };

    let stopped = match outcome {

        None => Stopped::Cancelled,

        Some(Ok(_)) => Stopped::Done,

        Some(Err(error)) => Stopped::Failed(cache_error(error)),

    };

    Finished {

        committed: matches!(stopped, Stopped::Done),

        stopped,

        report: None,

    }

}

async fn fill_gap(

    room: Room,

    mut report: HistoryRepairReport,

    mut cancel: watch::Receiver<bool>,

    pages: u32,

    spend_budget: bool,

) -> Finished {

    let mut committed = false;

    let stopped = async {

        let (cache, _handles) = room.event_cache().await.map_err(cache_error)?;

        for _ in 0..pages {

            let next = tokio::select! {

                biased;

                _ = cancel.changed() => return Ok(Stopped::Cancelled),

                result = cache.fill_history_gap_once(report.gap_id.clone()) => match result {

                    Ok(next) => next,

                    Err(error) => {

                        committed |= matches!(error, EventCacheError::PaginationError(_));

                        return Err(cache_error(error));

                    }

                }

            };

            let (gap, events) = match next {

                HistoryGapFillOutcome::NotFound => {

                    return Err(Error::new(

                        ErrorKind::EventNotFound,

                        "history gap was not found",

                    ));

                }

                HistoryGapFillOutcome::Busy => return Err(room_busy()),

                HistoryGapFillOutcome::Updated { gap, events } => (gap, events),

            };

            // the sdk answers an already unavailable gap without asking the server, and a joined

            // one with NotFound. an empty first page that hits the end looks the same, uncounted

            let settled_before =

                report.pages == 0 && events == 0 && gap.status == HistoryGapStatus::Unavailable;

            if !settled_before {

                report.pages += 1;

                committed = true;

            }

            report.events = report

                .events

                .saturating_add(u32::try_from(events).unwrap_or(u32::MAX));

            match gap.status {

                HistoryGapStatus::Joined => {

                    report.continuity = HistoryContinuity::Verified;

                    return Ok(Stopped::Done);

                }

                HistoryGapStatus::Unavailable => {

                    report.continuity = HistoryContinuity::AccessibleStart;

                    report.stop = Some(HistoryRepairStop::HistoryUnavailable);

                    return Ok(Stopped::Done);

                }

                HistoryGapStatus::Paused => {

                    report.stop = Some(HistoryRepairStop::CursorStalled);

                    return Ok(Stopped::Done);

                }

                HistoryGapStatus::Failed { message } => {

                    return Err(Error::new(ErrorKind::Network, message));

                }

                HistoryGapStatus::Idle | HistoryGapStatus::BudgetSpent => {}

            }

        }

        report.stop = Some(HistoryRepairStop::Budget);

        if spend_budget {

            if let Err(error) = cache.spend_history_budget(report.gap_id.clone()).await {

                tracing::debug!(%error, "marking a history gap's spent budget failed");

            }

            committed = true;

        }

        Ok(Stopped::Done)

    }

    .await

    .unwrap_or_else(Stopped::Failed);

    Finished {

        stopped,

        committed,

        report: Some(report),

    }

}

fn report(room_id: OwnedRoomId, gap_id: OwnedTransactionId) -> HistoryRepairReport {

    HistoryRepairReport {

        room_id,

        gap_id,

        pages: 0,

        events: 0,

        phase: HistoryRepairPhase::Bridging,

        continuity: HistoryContinuity::Unresolved,

        stop: None,

        failure: None,

    }

}

fn cache_error(source: EventCacheError) -> Error {

    let kind = match &source {

        EventCacheError::PaginationError(error) => crate::error::classify_sdk_error(error),

        EventCacheError::InitialPaginationError(_) => ErrorKind::Network,

        EventCacheError::RoomNotFound { .. } => ErrorKind::RoomUnknown,

        EventCacheError::ClientDropped => ErrorKind::Interrupted,

        _ => ErrorKind::Store,

    };

    Error::with_source(kind, "room history operation failed", source)

}

fn room_busy() -> Error {

    Error::new(

        ErrorKind::RoomBusy,

        "room history operation is already running",

    )

}