lattice-embed 0.9.0

//! Coordinates an embedding backfill above the migration state machine.
//!
//! It selects model routes, accounts for batches, and maintains the post-cutover
//! rollback window; it does not perform embedding work or persistence itself.
//!
//! See [docs/backfill.md](../../docs/backfill.md) for the routing and batching design.

use std::time::{Duration, Instant};

use crate::migration::{
    MigrationController, MigrationError, MigrationPlan, MigrationProgress, MigrationState,
};
use crate::model::EmbeddingModel;

use super::types::{BackfillConfig, EmbeddingRoute, EmbeddingRoutingConfig, RoutingPhase};

/// Coordinates the backfill process during embedding migration.
///
/// It layers request/query routing, batch sizing, and rollback timing over the migration
/// controller; callers perform the actual embedding and index updates.
///
/// See [docs/backfill.md](../../docs/backfill.md) for the lifecycle and routing tables.
#[derive(Debug)]
pub struct BackfillCoordinator {
    pub(super) config: BackfillConfig,
    pub(super) controller: MigrationController,
    /// Count accumulated by `record_batch` calls.
    backfilled_count: usize,
    /// Timestamp when cutover occurred (for rollback window tracking).
    pub(super) cutover_at: Option<Instant>,
    /// Duration of rollback window (computed from config).
    rollback_window: Duration,
}

impl BackfillCoordinator {
    /// Create a new backfill coordinator with the given plan and config.
    pub fn new(plan: MigrationPlan, config: BackfillConfig) -> Self {
        let rollback_window = Duration::from_secs(config.rollback_window_secs);
        Self {
            config,
            controller: MigrationController::new(plan),
            backfilled_count: 0,
            cutover_at: None,
            rollback_window,
        }
    }

    /// Create a backfill coordinator with default configuration.
    pub fn with_defaults(plan: MigrationPlan) -> Self {
        Self::new(plan, BackfillConfig::default())
    }

    /// Start the backfill process.
    ///
    /// Transitions the underlying migration from `Planned` to `InProgress`.
    ///
    /// # Errors
    ///
    /// Returns [`MigrationError::InvalidTransition`] if not in `Planned` state.
    pub fn start(&mut self) -> Result<(), MigrationError> {
        self.controller.start()
    }

    /// Routes an embedding write request according to migration state.
    ///
    /// `is_new_document` distinguishes initial indexing from ordinary re-embedding.
    /// See [`docs/backfill.md`](../../docs/backfill.md#backfillcoordinatorroute_request) for dual-write semantics.
    pub fn route_request(&self, is_new_document: bool) -> EmbeddingRoute {
        match self.controller.state() {
            MigrationState::Planned => EmbeddingRoute::Legacy,
            MigrationState::InProgress { .. } => {
                if is_new_document && self.config.dual_write {
                    EmbeddingRoute::DualWrite
                } else {
                    EmbeddingRoute::Legacy
                }
            }
            MigrationState::Paused { .. } => EmbeddingRoute::Legacy,
            MigrationState::Completed { .. } => EmbeddingRoute::Target,
            MigrationState::Failed { .. } => EmbeddingRoute::Legacy,
            MigrationState::Cancelled { .. } => EmbeddingRoute::Legacy,
        }
    }

    /// Route a query request (which model's embeddings to search against).
    ///
    /// Returns [`EmbeddingRoute::Target`] once migration is complete, or
    /// once progress exceeds [`BackfillConfig::target_query_threshold`]
    /// during an active migration. Otherwise returns [`EmbeddingRoute::Legacy`].
    pub fn route_query(&self) -> EmbeddingRoute {
        match self.controller.state() {
            MigrationState::Completed { .. } => EmbeddingRoute::Target,
            MigrationState::InProgress {
                processed, total, ..
            } => {
                let progress = if *total == 0 {
                    1.0
                } else {
                    *processed as f64 / *total as f64
                };
                if progress >= self.config.target_query_threshold {
                    EmbeddingRoute::Target
                } else {
                    EmbeddingRoute::Legacy
                }
            }
            _ => EmbeddingRoute::Legacy,
        }
    }

    /// Records a completed backfill batch and advances the migration.
    ///
    /// Starts rollback timing when the batch completes the migration.
    ///
    /// # Errors
    ///
    /// Returns [`MigrationError::InvalidTransition`] outside `InProgress`.
    ///
    /// See [`docs/backfill.md`](../../docs/backfill.md#backfillcoordinatorrecord_batch) for accounting and cutover timing.
    pub fn record_batch(&mut self, count: usize) -> Result<(), MigrationError> {
        let was_in_progress = matches!(self.controller.state(), MigrationState::InProgress { .. });
        self.backfilled_count += count;
        self.controller.record_progress(count)?;

        // Track cutover time when transitioning to Completed
        if was_in_progress && matches!(self.controller.state(), MigrationState::Completed { .. }) {
            self.cutover_at = Some(Instant::now());
        }

        Ok(())
    }

    /// Record a non-fatal error encountered during backfill processing.
    ///
    /// Increments the error counter without changing state. Useful for
    /// tracking transient failures (e.g., individual embedding retries).
    pub fn record_error(&mut self) {
        self.controller.record_error();
    }

    /// Pause the backfill process.
    ///
    /// # Errors
    ///
    /// Returns [`MigrationError::InvalidTransition`] if not in `InProgress` state.
    pub fn pause(&mut self, reason: impl Into<String>) -> Result<(), MigrationError> {
        self.controller.pause(reason)
    }

    /// Resume the backfill process.
    ///
    /// # Errors
    ///
    /// Returns [`MigrationError::InvalidTransition`] if not in `Paused` or `Failed` state.
    pub fn resume(&mut self) -> Result<(), MigrationError> {
        self.controller.resume()
    }

    /// Cancel the backfill process.
    ///
    /// # Errors
    ///
    /// Returns [`MigrationError::InvalidTransition`] if already in a terminal state.
    pub fn cancel(&mut self) -> Result<(), MigrationError> {
        self.controller.cancel()
    }

    /// Get the current migration state.
    #[inline]
    pub fn state(&self) -> &MigrationState {
        self.controller.state()
    }

    /// Get a snapshot of current migration progress.
    #[inline]
    pub fn progress(&self) -> MigrationProgress {
        self.controller.progress()
    }

    /// Get the source model being migrated from.
    #[inline]
    pub fn source_model(&self) -> EmbeddingModel {
        self.controller.plan().source_model
    }

    /// Get the target model being migrated to.
    #[inline]
    pub fn target_model(&self) -> EmbeddingModel {
        self.controller.plan().target_model
    }

    /// Check if we are currently in the post-cutover rollback window.
    ///
    /// Returns `true` only after this coordinator observed completion and before its
    /// configured rollback duration has elapsed.
    #[inline]
    pub fn in_rollback_window(&self) -> bool {
        self.cutover_at
            .map(|t| t.elapsed() < self.rollback_window)
            .unwrap_or(false)
    }

    /// Get the routing configuration for the current state.
    ///
    /// The post-cutover rollback window selects target queries and both write models.
    /// See [docs/backfill.md](../../docs/backfill.md) for the complete routing contract.
    pub fn routing_config(&self) -> EmbeddingRoutingConfig {
        match self.controller.state() {
            MigrationState::Planned => EmbeddingRoutingConfig {
                query_model: self.source_model(),
                write_models: vec![self.source_model()],
                phase: RoutingPhase::Stable,
                migration_id: None,
            },
            MigrationState::InProgress { .. } => {
                let write_models = if self.config.dual_write {
                    vec![self.source_model(), self.target_model()]
                } else {
                    vec![self.source_model()]
                };
                EmbeddingRoutingConfig {
                    query_model: self.source_model(), // query legacy during migration
                    write_models,
                    phase: RoutingPhase::Migrating,
                    migration_id: Some(self.controller.plan().id.clone()),
                }
            }
            MigrationState::Completed { .. } => {
                // During rollback window, still dual-write
                if self.in_rollback_window() {
                    EmbeddingRoutingConfig {
                        query_model: self.target_model(), // query new
                        write_models: vec![self.source_model(), self.target_model()], // still dual
                        phase: RoutingPhase::RollbackWindow,
                        migration_id: Some(self.controller.plan().id.clone()),
                    }
                } else {
                    EmbeddingRoutingConfig {
                        query_model: self.target_model(),
                        write_models: vec![self.target_model()],
                        phase: RoutingPhase::Stable,
                        migration_id: None,
                    }
                }
            }
            // Paused, Failed, Cancelled: fall back to source-only
            _ => EmbeddingRoutingConfig {
                query_model: self.source_model(),
                write_models: vec![self.source_model()],
                phase: RoutingPhase::Stable,
                migration_id: None,
            },
        }
    }

    /// Get the backfill configuration.
    #[inline]
    pub fn config(&self) -> &BackfillConfig {
        &self.config
    }

    /// Get the total count accumulated by `record_batch` calls.
    #[inline]
    pub fn backfilled_count(&self) -> usize {
        self.backfilled_count
    }

    /// Returns the number of items that may be processed in the next batch.
    ///
    /// Returns zero outside `InProgress` and caps effective remaining work at the configured size.
    /// See [`docs/backfill.md`](../../docs/backfill.md#backfillcoordinatornext_batch_size) for its accounting rule.
    pub fn next_batch_size(&self) -> usize {
        match self.controller.state() {
            MigrationState::InProgress {
                processed,
                total,
                skipped,
            } => {
                let remaining = total.saturating_sub(*processed).saturating_sub(*skipped);
                remaining.min(self.config.batch_size)
            }
            _ => 0,
        }
    }
}