macltok

use matrix_sdk::{Room, event_cache::HistoryGapFillOutcome};

use tokio::sync::watch;

use super::{

    executor::{JoinError, JoinHandle, spawn},

    timeline::SessionContext,

};

use crate::{

    CommandOutput, Error, ErrorKind, HistoryContinuity, HistoryRepairBudget, HistoryRepairCommand,

    HistoryRepairPhase, HistoryRepairPolicy, HistoryRepairReport, HistoryRepairStop,

    HistoryRepairStrategy, RecentHistoryState, transport::Responder,

};

const MAX_GAP_PAGES: u32 = 4;

struct Running {

    task: JoinHandle<CompletionResult>,

    cancel: watch::Sender<bool>,

    responders: Vec<PendingResponder>,

    recent: bool,

}

struct PendingResponder {

    responder: Responder,

    cancel_ack: bool,

}

pub(super) struct Completion {

    pub(super) report: HistoryRepairReport,

    pub(super) recent_state: Option<RecentHistoryState>,

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

    responders: Vec<PendingResponder>,

}

impl Completion {

    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());

            }

        }

    }

}

struct CompletionResult {

    report: HistoryRepairReport,

    recent_state: Option<RecentHistoryState>,

    result: Result<CommandOutput, Error>,

}

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()

    }

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

        self.start_recent_with(HistoryRepairPolicy::default(), None, true, Vec::new())

            .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::Start { policy, budget, .. } => {

                if let Err(error) = validate_budget(&budget) {

                    reply.done(Err(error));

                    return;

                }

                let policy = policy.unwrap_or_default();

                self.start_recent_with(policy, Some(reply), true, Vec::new())

                    .await;

            }

            HistoryRepairCommand::StartRecent { .. } => {

                self.start_recent_with(

                    HistoryRepairPolicy::default(),

                    Some(reply),

                    true,

                    Vec::new(),

                )

                .await;

            }

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

                self.start_gap_fill(gap_id, reply).await;

            }

            HistoryRepairCommand::Continue { .. } => reply.done(Err(Error::new(

                ErrorKind::Invalid,

                "continue requires a history gap ID, use FillGap",

            ))),

            HistoryRepairCommand::GetPolicy => {

                reply.send(

                    self.context

                        .mentions

                        .repair_policy()

                        .await

                        .map(CommandOutput::HistoryRepairPolicy),

                );

            }

            HistoryRepairCommand::SetPolicy(policy) => {

                reply.done(self.context.mentions.set_repair_policy(policy).await);

            }

        }

    }

    async fn start_recent_with(

        &mut self,

        policy: HistoryRepairPolicy,

        responder: Option<Responder>,

        publish_recent: bool,

        mut responders: Vec<PendingResponder>,

    ) {

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

            if let Some(responder) = responder {

                if running.recent {

                    running.responders.push(PendingResponder {

                        responder,

                        cancel_ack: false,

                    });

                } else {

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

                }

            }

            return;

        }

        if let Some(responder) = responder {

            responders.push(PendingResponder {

                responder,

                cancel_ack: false,

            });

        }

        if publish_recent

            && self

                .context

                .events

                .send(crate::Event::RecentHistoryChanged {

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

                    state: RecentHistoryState::Loading,

                })

                .await

                .is_err()

        {

            for responder in responders {

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

            }

            return;

        }

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

        let room = self.room.clone();

        let room_id = self.room.room_id().to_owned();

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

        self.running = Some(Running {

            task,

            cancel,

            responders,

            recent: true,

        });

    }

    async fn start_gap_fill(&mut self, gap_id: ruma::OwnedTransactionId, responder: Responder) {

        if self.is_running() {

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

            return;

        }

        let room = self.room.clone();

        let room_id = self.room.room_id().to_owned();

        let mut progress = report(room_id.clone(), HistoryRepairPolicy::default());

        progress.gap_id = Some(gap_id.clone());

        progress.strategy = HistoryRepairStrategy::Bridge;

        progress.phase = HistoryRepairPhase::Bridging;

        if self

            .context

            .events

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

            .await

            .is_err()

        {

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

            return;

        }

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

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

        self.running = Some(Running {

            task,

            cancel,

            responders: vec![PendingResponder {

                responder,

                cancel_ack: false,

            }],

            recent: 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.cancel.send_replace(true);

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

        Some(self.finish_task(result))

    }

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

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

        let result = match result {

            Ok(result) => result,

            Err(_) => CompletionResult {

                report: failed_report(

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

                    HistoryRepairPolicy::default(),

                    Error::interrupted(),

                ),

                recent_state: running.recent.then(|| {

                    RecentHistoryState::Failed(

                        "recent history operation was interrupted".to_owned(),

                    )

                }),

                result: Err(Error::interrupted()),

            },

        };

        Completion {

            report: result.report,

            recent_state: result.recent_state,

            result: result.result,

            responders: running.responders,

        }

    }

}

fn validate_budget(budget: &HistoryRepairBudget) -> Result<(), Error> {

    if !(1..=100).contains(&budget.max_pages)

        || !(50..=10_000).contains(&budget.max_events)

        || !(1..=60_000).contains(&budget.max_duration_ms)

    {

        return Err(Error::new(

            ErrorKind::Invalid,

            "repair budget requires 1 to 100 pages, 50 to 10000 events, and 1 to 60000 milliseconds",

        ));

    }

    Ok(())

}

async fn load_recent(

    room: Room,

    room_id: ruma::OwnedRoomId,

    policy: HistoryRepairPolicy,

    mut cancel: watch::Receiver<bool>,

) -> CompletionResult {

    let base = report(room_id.clone(), policy);

    let outcome = tokio::select! {

        biased;

        _ = cancel.changed() => Err(Error::interrupted()),

        result = async {

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

            cache.load_recent_history().await.map_err(cache_error)

        } => result,

    };

    match outcome {

        Ok(outcome) => {

            let mut report = base;

            report.pages = outcome.pages;

            report.events = outcome.events.min(u32::MAX as usize) as u32;

            report.recent_available = true;

            report.phase = HistoryRepairPhase::Ready;

            match outcome.gap.map(|gap| gap.status) {

                None => report.continuity = HistoryContinuity::AccessibleStart,

                Some(matrix_sdk::event_cache::HistoryGapStatus::Joined) => {

                    report.continuity = HistoryContinuity::Verified

                }

                Some(matrix_sdk::event_cache::HistoryGapStatus::Unavailable) => {

                    report.continuity = HistoryContinuity::AccessibleStart;

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

                    report.phase = HistoryRepairPhase::Paused;

                }

                Some(matrix_sdk::event_cache::HistoryGapStatus::Paused) => {

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

                    report.phase = HistoryRepairPhase::Paused;

                }

                Some(matrix_sdk::event_cache::HistoryGapStatus::Failed { message }) => {

                    let error = Error::new(ErrorKind::Network, message.clone());

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

                    report.phase = HistoryRepairPhase::Paused;

                }

                Some(matrix_sdk::event_cache::HistoryGapStatus::Idle) => {}

            }

            CompletionResult {

                report: report.clone(),

                recent_state: Some(RecentHistoryState::Loaded),

                result: Ok(CommandOutput::HistoryRepair(report)),

            }

        }

        Err(error) => CompletionResult {

            report: failed_report(room_id, base.requested_policy, error.clone()),

            recent_state: Some(RecentHistoryState::Failed(error.message.clone())),

            result: Err(error),

        },

    }

}

async fn fill_gap(

    room: Room,

    room_id: ruma::OwnedRoomId,

    gap_id: ruma::OwnedTransactionId,

    mut cancel: watch::Receiver<bool>,

) -> CompletionResult {

    let mut report = report(room_id, HistoryRepairPolicy::default());

    report.gap_id = Some(gap_id.clone());

    report.strategy = HistoryRepairStrategy::Bridge;

    let outcome = async {

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

        let mut last_gap = None;

        for _ in 0..MAX_GAP_PAGES {

            let next = tokio::select! {

                biased;

                _ = cancel.changed() => return Err(Error::interrupted()),

                result = cache.fill_history_gap_once(gap_id.clone()) => result.map_err(cache_error)?,

            };

            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 } => {

                    report.pages += 1;

                    report.events = report.events.saturating_add(events.min(u32::MAX as usize) as u32);

                    match gap.status.clone() {

                        matrix_sdk::event_cache::HistoryGapStatus::Joined => {

                            report.continuity = HistoryContinuity::Verified;

                            last_gap = Some(gap);

                            break;

                        }

                        matrix_sdk::event_cache::HistoryGapStatus::Unavailable => {

                            report.continuity = HistoryContinuity::AccessibleStart;

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

                            last_gap = Some(gap);

                            break;

                        }

                        matrix_sdk::event_cache::HistoryGapStatus::Paused => {

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

                            last_gap = Some(gap);

                            break;

                        }

                        matrix_sdk::event_cache::HistoryGapStatus::Failed { message } => {

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

                        }

                        matrix_sdk::event_cache::HistoryGapStatus::Idle => last_gap = Some(gap),

                    }

                }

            }

        }

        if report.pages == MAX_GAP_PAGES

            && report.continuity == HistoryContinuity::Unresolved

            && report.stop.is_none()

        {

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

        }

        Ok::<_, Error>(last_gap)

    }

    .await;

    let command_result = match outcome {

        Ok(_) => {

            report.recent_available = true;

            report.phase = if report.stop.is_some() {

                HistoryRepairPhase::Paused

            } else {

                HistoryRepairPhase::Ready

            };

            Ok(CommandOutput::Done)

        }

        Err(error) => {

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

            report.phase = HistoryRepairPhase::Paused;

            Err(error)

        }

    };

    CompletionResult {

        report,

        recent_state: None,

        result: command_result,

    }

}

fn report(room_id: ruma::OwnedRoomId, policy: HistoryRepairPolicy) -> HistoryRepairReport {

    HistoryRepairReport {

        room_id,

        gap_id: None,

        requested_policy: policy,

        strategy: HistoryRepairStrategy::Recent,

        coverage_duration_ms: None,

        threshold_evaluated: false,

        pages: 0,

        events: 0,

        phase: HistoryRepairPhase::FetchingRecent,

        continuity: HistoryContinuity::Unresolved,

        recent_available: false,

        stop: None,

        failure: None,

    }

}

fn failed_report(

    room_id: ruma::OwnedRoomId,

    policy: HistoryRepairPolicy,

    error: Error,

) -> HistoryRepairReport {

    let mut report = report(room_id, policy);

    report.phase = HistoryRepairPhase::Paused;

    report.failure = Some(error);

    report

}

fn cache_error(source: matrix_sdk::event_cache::EventCacheError) -> Error {

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

}

fn room_busy() -> Error {

    Error::new(

        ErrorKind::RoomBusy,

        "room history operation is already running",

    )

}