macltok

use std::{

    collections::HashSet,

    sync::{

        Arc,

        atomic::{AtomicBool, Ordering},

    },

    time::Duration,

};

use matrix_sdk::{Room, deserialized_responses::TimelineEvent, room::MessagesOptions};

use matrix_sdk_ui::sync_service::SyncService;

use ruma::{OwnedEventId, time::Instant};

use tokio::sync::{OwnedMutexGuard, mpsc, oneshot, watch};

use super::{

    executor::{JoinHandle, spawn},

    timeline::SessionContext,

};

use crate::{

    CommandOutput, Error, ErrorKind, Event, HistoryContinuity, HistoryRepairBudget,

    HistoryRepairCommand, HistoryRepairPhase, HistoryRepairPolicy, HistoryRepairReport,

    HistoryRepairStop, HistoryRepairStrategy, transport::Responder,

};

const PAGE_SIZE: u16 = 50;

const STAGING_LIMIT: usize = 10_000;

pub(super) struct InstallRequest {

    pub(super) events: Vec<TimelineEvent>,

    pub(super) prev_batch: Option<String>,

    pub(super) reply: oneshot::Sender<Result<(), Error>>,

}

struct Stage {

    events: Vec<TimelineEvent>,

    cursor: Option<String>,

    anchor: Option<OwnedEventId>,

    report: HistoryRepairReport,

}

struct Running {

    task: JoinHandle<Stage>,

    cancel: watch::Sender<bool>,

    reply: Responder,

}

pub(super) struct RepairController {

    room: Room,

    context: SessionContext,

    installs: mpsc::Sender<InstallRequest>,

    requests: mpsc::Receiver<InstallRequest>,

    running: Option<Running>,

    stage: Option<Stage>,

}

impl RepairController {

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

        let (installs, requests) = mpsc::channel(1);

        Self {

            room,

            context,

            installs,

            requests,

            running: None,

            stage: None,

        }

    }

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

        self.running.is_some()

    }

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

        if matches!(command, HistoryRepairCommand::Cancel { .. }) {

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

                running.cancel.send_replace(true);

            }

            reply.done(Ok(()));

            return;

        }

        if self.is_running() {

            reply.done(Err(Error::new(

                ErrorKind::RoomBusy,

                "room history repair is running",

            )));

            return;

        }

        let (stage, budget) = match command {

            HistoryRepairCommand::Start { policy, budget, .. } => {

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

                    reply.done(Err(error));

                    return;

                }

                let policy = match policy {

                    Some(policy) => policy,

                    None => match self.context.mentions.repair_policy().await {

                        Ok(policy) => policy,

                        Err(error) => {

                            reply.done(Err(error));

                            return;

                        }

                    },

                };

                let coverage = match self

                    .context

                    .mentions

                    .repair_coverage(self.room.room_id().to_owned())

                    .await

                {

                    Ok(coverage) => coverage,

                    Err(error) => {

                        reply.done(Err(error));

                        return;

                    }

                };

                let duration = coverage.duration_ms();

                let report = HistoryRepairReport {

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

                    requested_policy: policy.clone(),

                    strategy: HistoryRepairStrategy::Recent,

                    coverage_duration_ms: duration,

                    threshold_evaluated: duration.is_some()

                        && coverage.unresolved

                        && matches!(policy, HistoryRepairPolicy::BridgeShortGaps { .. }),

                    pages: 0,

                    events: 0,

                    phase: HistoryRepairPhase::FetchingRecent,

                    continuity: HistoryContinuity::Unresolved,

                    recent_available: false,

                    stop: None,

                    failure: None,

                };

                (

                    Stage {

                        events: Vec::new(),

                        cursor: None,

                        anchor: None,

                        report,

                    },

                    budget,

                )

            }

            HistoryRepairCommand::Continue { budget, .. } => {

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

                    reply.done(Err(error));

                    return;

                }

                let Some(mut stage) = self.stage.take() else {

                    reply.done(Err(Error::new(

                        ErrorKind::Invalid,

                        "no history bridge to continue",

                    )));

                    return;

                };

                if !stage.report.recent_available

                    || stage.report.continuity != HistoryContinuity::Unresolved

                {

                    reply.done(Err(Error::new(

                        ErrorKind::Invalid,

                        "history bridge is not resumable",

                    )));

                    self.stage = Some(stage);

                    return;

                }

                stage.report.strategy = HistoryRepairStrategy::Bridge;

                stage.report.phase = HistoryRepairPhase::Bridging;

                stage.report.stop = None;

                stage.report.failure = None;

                (stage, budget)

            }

            HistoryRepairCommand::StartRecent { .. } => {

                let Some(mut stage) = self.stage.take() else {

                    reply.done(Err(Error::new(

                        ErrorKind::Invalid,

                        "no recent history repair",

                    )));

                    return;

                };

                if !stage.report.recent_available {

                    self.stage = Some(stage);

                    reply.done(Err(Error::new(

                        ErrorKind::Invalid,

                        "recent history is not available, retry the repair",

                    )));

                    return;

                }

                stage.report.strategy = HistoryRepairStrategy::Recent;

                stage.report.phase = HistoryRepairPhase::Ready;

                stage.report.stop = None;

                stage.report.failure = None;

                let report = stage.report.clone();

                let published = self

                    .context

                    .events

                    .publish_repair(report.clone())

                    .await

                    .map_err(|_| Error::interrupted());

                reply.send(published.map(|_| CommandOutput::HistoryRepair(report)));

                return;

            }

            _ => {

                reply.done(Err(Error::new(

                    ErrorKind::Invalid,

                    "repair command was misrouted",

                )));

                return;

            }

        };

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

        let room = self.room.clone();

        let context = self.context.clone();

        let installs = self.installs.clone();

        let task =

            spawn(async move { run(room, context, installs, receiver, stage, budget).await });

        self.running = Some(Running {

            task,

            cancel,

            reply,

        });

    }

    pub(super) async fn next_install(&mut self) -> Option<InstallRequest> {

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

        tokio::select! {

            request = self.requests.recv() => request,

            result = &mut running.task => {

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

                match result {

                    Ok(stage) => {

                        let report = stage.report.clone();

                        let published = self.context.events.publish_repair(report.clone()).await

                            .map_err(|_| Error::interrupted());

                        running.reply.send(published.map(|_| CommandOutput::HistoryRepair(report)));

                        self.stage = Some(stage);

                    }

                    Err(_) => running.reply.done(Err(Error::interrupted())),

                }

                None

            }

        }

    }

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

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

            running.task.abort();

            match (&mut running.task).await {

                Ok(stage) => {

                    let report = stage.report;

                    let published = self

                        .context

                        .events

                        .publish_repair(report.clone())

                        .await

                        .map_err(|_| Error::interrupted());

                    running

                        .reply

                        .send(published.map(|_| CommandOutput::HistoryRepair(report)));

                }

                Err(_) => running.reply.done(Err(Error::interrupted())),

            }

        }

    }

}

impl Drop for RepairController {

    fn drop(&mut self) {

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

            running.task.abort();

        }

    }

}

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

}

fn should_bridge(policy: &HistoryRepairPolicy, duration: Option<u64>, boundary: bool) -> bool {

    match policy {

        HistoryRepairPolicy::Bridge => true,

        HistoryRepairPolicy::StartRecent => false,

        HistoryRepairPolicy::BridgeShortGaps { threshold_ms } => {

            boundary

                && *threshold_ms > 0

                && duration.is_some_and(|duration| duration <= *threshold_ms)

        }

    }

}

async fn run(

    room: Room,

    context: SessionContext,

    installs: mpsc::Sender<InstallRequest>,

    mut cancel: watch::Receiver<bool>,

    mut stage: Stage,

    budget: HistoryRepairBudget,

) -> Stage {

    let started = Instant::now();

    let initial_pages = stage.report.pages;

    let initial_events = stage.report.events;

    let result = async {

        if !stage.report.recent_available {

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

            stage.anchor = cache

                .events()

                .await

                .map_err(cache_error)?

                .iter()

                .find_map(TimelineEvent::event_id);

            let _ = context

                .events

                .send(Event::HistoryRepairChanged(stage.report.clone()))

                .await;

            let guard = pause_sync(&context, &mut cancel, &started, &budget).await?;

            let page = fetch(&room, None, &mut cancel, &started, &budget).await?;

            stage.report.pages += 1;

            stage.report.events += page.chunk.len() as u32;

            stage.cursor = page.end;

            stage.events = page.chunk.into_iter().rev().collect();

            deduplicate(&mut stage.events);

            install(&installs, &stage).await?;

            stage.report.recent_available = true;

            if contains_anchor(&stage) {

                stage.report.strategy = HistoryRepairStrategy::Overlap;

                stage.report.continuity = HistoryContinuity::Verified;

            } else if stage.cursor.is_none() {

                stage.report.continuity = HistoryContinuity::AccessibleStart;

            }

            if stage.report.continuity != HistoryContinuity::Unresolved {

                context

                    .mentions

                    .repair_resolved(room.room_id().to_owned())

                    .await?;

            }

            guard.finish().await;

            let _ = context

                .events

                .send(Event::HistoryRepairChanged(stage.report.clone()))

                .await;

            if stage.report.continuity != HistoryContinuity::Unresolved {

                return Ok(());

            }

            if !should_bridge(

                &stage.report.requested_policy,

                stage.report.coverage_duration_ms,

                stage.report.threshold_evaluated,

            ) {

                return Ok(());

            }

            stage.report.strategy = HistoryRepairStrategy::Bridge;

        }

        stage.report.phase = HistoryRepairPhase::Bridging;

        let _ = context

            .events

            .send(Event::HistoryRepairChanged(stage.report.clone()))

            .await;

        if contains_anchor(&stage) {

            stage.report.continuity = HistoryContinuity::Verified;

        } else if stage.cursor.is_none() {

            stage.report.continuity = HistoryContinuity::AccessibleStart;

        }

        while stage.report.continuity == HistoryContinuity::Unresolved {

            if *cancel.borrow() {

                stage.report.stop = Some(HistoryRepairStop::Cancelled);

                break;

            }

            if stage.report.pages - initial_pages + 1 >= u32::from(budget.max_pages)

                || stage.report.events - initial_events + 2 * u32::from(PAGE_SIZE)

                    > budget.max_events

                || stage.events.len() + 2 * usize::from(PAGE_SIZE) > STAGING_LIMIT

                || started.elapsed() >= Duration::from_millis(budget.max_duration_ms)

            {

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

                break;

            }

            if stage.anchor.is_none() {

                stage.report.stop = Some(HistoryRepairStop::NoAnchor);

                break;

            }

            let Some(cursor) = stage.cursor.clone() else {

                break;

            };

            let page = fetch(&room, Some(&cursor), &mut cancel, &started, &budget).await?;

            stage.report.pages += 1;

            stage.report.events += page.chunk.len() as u32;

            if page.end.as_deref() == Some(&cursor) {

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

                break;

            }

            let mut older: Vec<_> = page.chunk.into_iter().rev().collect();

            older.append(&mut stage.events);

            deduplicate(&mut older);

            stage.events = older;

            stage.cursor = page.end;

            if contains_anchor(&stage) {

                stage.report.continuity = HistoryContinuity::Verified;

            } else if stage.cursor.is_none() {

                stage.report.continuity = HistoryContinuity::AccessibleStart;

            }

            if stage.report.continuity == HistoryContinuity::Unresolved {

                let _ = context

                    .events

                    .send(Event::HistoryRepairChanged(stage.report.clone()))

                    .await;

            }

        }

        if stage.report.continuity != HistoryContinuity::Unresolved

            && stage.report.strategy == HistoryRepairStrategy::Bridge

        {

            let guard = pause_sync(&context, &mut cancel, &started, &budget).await?;

            let head = fetch(&room, None, &mut cancel, &started, &budget).await?;

            stage.report.pages += 1;

            stage.report.events += head.chunk.len() as u32;

            let head: Vec<_> = head.chunk.into_iter().rev().collect();

            if !extend_head(&mut stage.events, head) {

                stage.report.continuity = HistoryContinuity::Unresolved;

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

                return Ok(());

            }

            install(&installs, &stage).await?;

            context

                .mentions

                .repair_resolved(room.room_id().to_owned())

                .await?;

            guard.finish().await;

        }

        Ok::<_, Error>(())

    }

    .await;

    if let Err(error) = result {

        if error.kind == ErrorKind::Interrupted {

            stage.report.stop = Some(HistoryRepairStop::Cancelled);

        } else if error.kind == ErrorKind::RoomBusy {

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

        } else {

            stage.report.failure = Some(error);

        }

        stage.report.continuity = HistoryContinuity::Unresolved;

    }

    stage.report.phase = if stage.report.stop.is_some() || stage.report.failure.is_some() {

        HistoryRepairPhase::Paused

    } else {

        HistoryRepairPhase::Ready

    };

    stage

}

fn contains_anchor(stage: &Stage) -> bool {

    stage.anchor.as_ref().is_some_and(|anchor| {

        stage

            .events

            .iter()

            .any(|event| event.event_id().as_ref() == Some(anchor))

    })

}

fn deduplicate(events: &mut Vec<TimelineEvent>) {

    let mut seen = HashSet::new();

    events.retain(|event| event.event_id().is_some_and(|id| seen.insert(id)));

}

fn extend_head(events: &mut Vec<TimelineEvent>, head: Vec<TimelineEvent>) -> bool {

    let ids: HashSet<_> = events.iter().filter_map(TimelineEvent::event_id).collect();