//! Tenant migration state machine: Idle -> `PreparingTarget` -> `DualWrite` -> Finalizing -> Complete.

use std::sync::{Arc, Mutex};

use super::{
    control::ControlPlane,
    shard::ShardId,
    tenant::{TenantId, TenantRouter},
};
use crate::schema::TidalError;

/// Current state of a tenant migration.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum MigrationState {
    Idle,
    PreparingTarget { last_shipped_seqno: u64 },
    DualWrite { cutover_seqno: u64 },
    Finalizing { switched_at_ns: u64 },
    Complete,
}

/// Coordinates a single tenant's migration from source to target shard.
pub struct TenantMigration {
    pub tenant_id: TenantId,
    pub source_shard: ShardId,
    pub target_shard: ShardId,
    state: Mutex<MigrationState>,
    _control_plane: Arc<ControlPlane>,
    tenant_router: Arc<TenantRouter>,
    /// Durable sink for the tenant router's migration routing state. When set,
    /// every routing transition (enter dual-write, finalize) is persisted so a
    /// crash mid-migration cannot silently revert a finalized tenant to
    /// jump-hash routing (Accuracy-W). `None` (the default / ephemeral) keeps
    /// routing in memory only, identical to the previous behaviour.
    persistence: Option<Arc<dyn crate::storage::StorageEngine>>,
}

impl TenantMigration {
    /// Create a new migration in the `Idle` state.
    #[must_use]
    pub const fn new(
        tenant_id: TenantId,
        source_shard: ShardId,
        target_shard: ShardId,
        control_plane: Arc<ControlPlane>,
        tenant_router: Arc<TenantRouter>,
    ) -> Self {
        Self {
            tenant_id,
            source_shard,
            target_shard,
            state: Mutex::new(MigrationState::Idle),
            _control_plane: control_plane,
            tenant_router,
            persistence: None,
        }
    }

    /// Attach a durable storage sink so routing transitions are checkpointed.
    ///
    /// Call this in persistent (non-ephemeral) deployments with the items
    /// storage engine. Each subsequent `enter_dual_write` / `finalize` persists
    /// the tenant router's routing state via
    /// [`checkpoint_migration_routing`](crate::replication::tenant::checkpoint_migration_routing),
    /// so the migration survives a crash. Returns `self` for chaining.
    #[must_use]
    pub fn with_persistence(mut self, storage: Arc<dyn crate::storage::StorageEngine>) -> Self {
        self.persistence = Some(storage);
        self
    }

    /// Persist the router's routing state, propagating any checkpoint failure.
    ///
    /// Returns `Ok(())` when there is no durable sink (ephemeral mode) or the
    /// checkpoint succeeds. Returns the storage error if the durable write fails,
    /// so the caller can roll back the in-memory mutation and keep memory and
    /// disk in lock-step (durability before observable mutation, W38).
    fn persist_routing_strict(&self) -> crate::Result<()> {
        match &self.persistence {
            Some(storage) => crate::replication::tenant::checkpoint_migration_routing(
                storage.as_ref(),
                &self.tenant_router,
            ),
            None => Ok(()),
        }
    }

    /// Snapshot ONLY this migration's tenant routing for a targeted rollback.
    ///
    /// All `TenantMigration`s share one [`TenantRouter`] but each has only its
    /// own per-tenant state lock. A rollback that restored a *full-router*
    /// snapshot would clobber a concurrent migration's just-applied (and possibly
    /// already-persisted) routing for a different tenant. Snapshotting and
    /// restoring only `self.tenant_id` keeps every other tenant's routing intact
    /// (C10).
    fn snapshot_tenant_routing(&self) -> crate::replication::tenant::TenantRoutingSnapshot {
        self.tenant_router.snapshot_tenant_routing(self.tenant_id)
    }

    /// Restore ONLY this migration's tenant routing from a prior snapshot,
    /// leaving every other tenant untouched (C10).
    fn restore_tenant_routing(&self, snapshot: &crate::replication::tenant::TenantRoutingSnapshot) {
        self.tenant_router.restore_tenant_routing(snapshot);
    }

    /// Transition from `Idle` to `PreparingTarget`.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current state is not `Idle`.
    pub fn prepare_target(&self, last_shipped_seqno: u64) -> crate::Result<u64> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| TidalError::internal("migration", "state lock poisoned"))?;
        if *state != MigrationState::Idle {
            return Err(TidalError::InvalidState(
                "prepare_target called outside Idle".into(),
            ));
        }
        *state = MigrationState::PreparingTarget { last_shipped_seqno };
        drop(state);
        Ok(last_shipped_seqno)
    }

    /// Transition from `PreparingTarget` to `DualWrite`.
    ///
    /// The new dual-write routing is made durable before the transition is
    /// accepted, so a crash can never leave a tenant routing dual-writes that
    /// were never persisted (W38).
    ///
    /// # Errors
    ///
    /// - `InvalidState` if the current state is not `PreparingTarget`.
    /// - `Storage` if the routing checkpoint cannot be persisted; the in-memory
    ///   routing is rolled back and the migration stays in `PreparingTarget` so
    ///   the caller can retry once storage recovers.
    pub fn enter_dual_write(&self, cutover_seqno: u64) -> crate::Result<u64> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| TidalError::internal("migration", "state lock poisoned"))?;
        if !matches!(*state, MigrationState::PreparingTarget { .. }) {
            return Err(TidalError::InvalidState(
                "enter_dual_write called outside PreparingTarget".into(),
            ));
        }
        // Durability before observable mutation (W38): make the dual-write routing
        // durable BEFORE accepting the state-machine transition. Snapshot ONLY
        // this tenant's prior routing (C10) so a failed persist rolls back exactly
        // this tenant's entry without clobbering a concurrent migration's routing
        // for a different tenant on the shared router. The net guarantee is that
        // the in-memory dual-write entry exists iff it has been persisted, so a
        // crash can never strand an in-flight migration whose dual-write state
        // never reached disk.
        let prior = self.snapshot_tenant_routing();
        self.tenant_router
            .set_dual_write(self.tenant_id, self.source_shard, self.target_shard);
        if let Err(e) = self.persist_routing_strict() {
            // Roll back ONLY this tenant's routing, so the router never advertises
            // a dual-write that is not durable AND a concurrent migration for a
            // different tenant keeps its just-applied entry. The migration stays
            // in PreparingTarget; the caller can retry once storage recovers.
            self.restore_tenant_routing(&prior);
            return Err(e);
        }
        *state = MigrationState::DualWrite { cutover_seqno };
        drop(state);
        Ok(cutover_seqno)
    }

    /// Transition from `DualWrite` to `Finalizing`.
    ///
    /// The finalized routing (the cleared dual-write + the target pin) is made
    /// **durable before the transition is observable**, mirroring
    /// [`Self::enter_dual_write`]. This is the cutover that re-points every read
    /// for the tenant from source to target; a silent checkpoint failure here
    /// would let a crash revert the tenant to `DualWrite` (reads to the SOURCE
    /// shard) with no error surfaced — and the only later transition,
    /// [`Self::gc_source`], does not re-persist, so nothing would re-close the
    /// gap. So on a persist failure we restore the pre-finalize routing snapshot
    /// and leave the state in `DualWrite` so the caller retries once storage
    /// recovers (C9).
    ///
    /// # Errors
    ///
    /// - `InvalidState` if the current state is not `DualWrite`.
    /// - `NotReady` if `target_seqno` is behind the cutover sequence number —
    ///   the caller should wait and retry once the target has caught up.
    /// - `Storage` if the routing checkpoint cannot be persisted; the in-memory
    ///   routing is rolled back and the migration stays in `DualWrite` so the
    ///   caller can retry once storage recovers.
    pub fn finalize(&self, target_seqno: u64) -> crate::Result<()> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| TidalError::internal("migration", "state lock poisoned"))?;
        let MigrationState::DualWrite { cutover_seqno } = *state else {
            return Err(TidalError::InvalidState(
                "finalize called outside DualWrite".into(),
            ));
        };
        if target_seqno < cutover_seqno {
            return Err(TidalError::NotReady(format!(
                "target seqno {target_seqno} has not reached cutover seqno {cutover_seqno}"
            )));
        }
        // Durability before observable mutation (C9): snapshot only THIS tenant's
        // prior routing, apply the finalize in memory, then persist strictly. A
        // crash-after-finalize cannot revert the tenant to DualWrite, because the
        // pin is durable before `finalize` returns Ok. On a persist failure, roll
        // back ONLY this tenant's routing (never the whole router — a concurrent
        // migration for another tenant must not be clobbered, C10) and stay in
        // `DualWrite` so the caller retries.
        let prior = self.snapshot_tenant_routing();
        self.tenant_router
            .finalize_migration(self.tenant_id, self.target_shard);
        if let Err(e) = self.persist_routing_strict() {
            self.restore_tenant_routing(&prior);
            return Err(e);
        }
        let switched_at_ns = crate::replication::now_ns();
        *state = MigrationState::Finalizing { switched_at_ns };
        drop(state);
        Ok(())
    }

    /// Advance from `Finalizing` to `Complete` once the post-cutover safety
    /// window (`gc_window_ns`) has elapsed since the routing switch.
    ///
    /// # This is a state-only transition — it does NOT reclaim source data
    ///
    /// This flips the migration state to `Complete`; it performs **no** storage
    /// delete, no shipper teardown, and no source-shard reclamation. After it
    /// returns the tenant's data still lives on the old source shard. Physically
    /// reclaiming that keyspace (a delete-range over the source shard + flush,
    /// and tearing down the source-shard shipper) is a separate, **not-yet-wired**
    /// step tracked as the source-reclamation task in
    /// `docs/specs/14-scale-architecture.md` (per `CODING_GUIDELINES.md` §12).
    /// The window enforced here is the safety interval that step would wait out
    /// before deleting; the state flip marks "safe to reclaim", not "reclaimed".
    ///
    /// Named [`Self::mark_complete`] for that reason; [`Self::gc_source`] is kept
    /// as a backwards-compatible alias and forwards here.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current state is not `Finalizing`, or if
    /// `gc_window_ns` has not elapsed since the routing switch.
    pub fn mark_complete(&self, gc_window_ns: u64) -> crate::Result<()> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| TidalError::internal("migration", "state lock poisoned"))?;
        let MigrationState::Finalizing { switched_at_ns } = *state else {
            return Err(TidalError::InvalidState(
                "mark_complete called outside Finalizing".into(),
            ));
        };
        let now = crate::replication::now_ns();
        if now.saturating_sub(switched_at_ns) < gc_window_ns {
            return Err(TidalError::InvalidState(
                "gc_window_ns has not elapsed since routing switch".into(),
            ));
        }
        *state = MigrationState::Complete;
        drop(state);
        Ok(())
    }

    /// Backwards-compatible alias for [`Self::mark_complete`].
    ///
    /// The name historically implied source-data garbage collection, but the
    /// transition is state-only (no source-shard reclamation happens here — see
    /// [`Self::mark_complete`]). Retained so existing callers keep compiling;
    /// new code should call `mark_complete`.
    ///
    /// # Errors
    ///
    /// See [`Self::mark_complete`].
    pub fn gc_source(&self, gc_window_ns: u64) -> crate::Result<()> {
        self.mark_complete(gc_window_ns)
    }

    /// Return the current migration state.
    ///
    /// # Panics
    ///
    /// Panics if the internal state lock is poisoned.
    #[must_use]
    pub fn current_state(&self) -> MigrationState {
        self.state.lock().expect("state lock poisoned").clone()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::replication::{
        control::ControlPlane,
        lag::ReplicationLagGauge,
        state::ReplicationState,
        tenant::{ClusterTopology, TenantRouter},
    };

    fn make_migration() -> TenantMigration {
        let topo = Arc::new(std::sync::RwLock::new(ClusterTopology::single()));
        let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        let cp = Arc::new(ControlPlane::new(topo, Arc::clone(&router), lag));
        TenantMigration::new(TenantId(1), ShardId::SINGLE, ShardId(1), cp, router)
    }

    #[test]
    fn initial_state_is_idle() {
        let m = make_migration();
        assert_eq!(m.current_state(), MigrationState::Idle);
    }

    #[test]
    fn prepare_target_transitions_to_preparing() {
        let m = make_migration();
        let seqno = m.prepare_target(42).unwrap();
        assert_eq!(seqno, 42);
        assert!(matches!(
            m.current_state(),
            MigrationState::PreparingTarget {
                last_shipped_seqno: 42
            }
        ));
    }

    #[test]
    fn prepare_target_wrong_state_returns_invalid_state() {
        let m = make_migration();
        m.prepare_target(1).unwrap();
        let err = m.prepare_target(2).unwrap_err();
        assert!(matches!(err, TidalError::InvalidState(_)));
    }

    #[test]
    fn enter_dual_write_transitions() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        let seqno = m.enter_dual_write(20).unwrap();
        assert_eq!(seqno, 20);
        assert!(matches!(
            m.current_state(),
            MigrationState::DualWrite { cutover_seqno: 20 }
        ));
    }

    #[test]
    fn finalize_transitions_to_finalizing() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        m.enter_dual_write(20).unwrap();
        m.finalize(25).unwrap();
        assert!(matches!(
            m.current_state(),
            MigrationState::Finalizing { .. }
        ));
    }

    #[test]
    fn finalize_fails_if_target_behind_cutover() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        m.enter_dual_write(20).unwrap();
        let err = m.finalize(15).unwrap_err();
        // NotReady (not InvalidState): the target has not caught up yet — caller should retry.
        assert!(matches!(err, TidalError::NotReady(_)));
    }

    #[test]
    fn gc_source_completes_migration_after_window() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        m.enter_dual_write(20).unwrap();
        m.finalize(25).unwrap();
        m.gc_source(0).unwrap();
        assert_eq!(m.current_state(), MigrationState::Complete);
    }

    #[test]
    fn gc_source_fails_before_window_elapses() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        m.enter_dual_write(20).unwrap();
        m.finalize(25).unwrap();
        let err = m.gc_source(600_000_000_000).unwrap_err();
        assert!(matches!(err, TidalError::InvalidState(_)));
    }

    /// W38: durability before observable mutation. If the dual-write checkpoint
    /// fails to persist, `enter_dual_write` must roll the router back (no dangling,
    /// non-durable dual-write entry) and leave the migration in `PreparingTarget`
    /// so the caller can retry once storage recovers.
    #[test]
    fn enter_dual_write_rolls_back_on_persist_failure() {
        use std::sync::RwLock;

        use crate::{
            replication::tenant::ShardAssignment,
            storage::{
                InMemoryBackend, StorageEngine, StorageError, batch::WriteBatch,
                iterator::PrefixIterator,
            },
        };

        /// Backend that fails every durable write (`write_batch` / `flush`) but
        /// reads through to an inner backend, simulating a disk that cannot accept
        /// the migration routing checkpoint.
        struct WriteFailsBackend(Arc<InMemoryBackend>);

        impl StorageEngine for WriteFailsBackend {
            fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
                self.0.get(key)
            }
            fn put(&self, _key: &[u8], _value: &[u8]) -> Result<(), StorageError> {
                Err(StorageError::Closed)
            }
            fn delete(&self, key: &[u8]) -> Result<(), StorageError> {
                self.0.delete(key)
            }
            fn scan_prefix(&self, prefix: &[u8]) -> PrefixIterator<'_> {
                self.0.scan_prefix(prefix)
            }
            fn write_batch(&self, _batch: WriteBatch) -> Result<(), StorageError> {
                Err(StorageError::Closed)
            }
            fn flush(&self) -> Result<(), StorageError> {
                Err(StorageError::Closed)
            }
        }

        let topo = Arc::new(RwLock::new(ClusterTopology {
            shards: vec![
                ShardAssignment {
                    shard_id: ShardId(0),
                    region_id: crate::replication::RegionId(0),
                },
                ShardAssignment {
                    shard_id: ShardId(1),
                    region_id: crate::replication::RegionId(1),
                },
            ],
        }));
        let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        let cp = Arc::new(ControlPlane::new(
            Arc::clone(&topo),
            Arc::clone(&router),
            lag,
        ));

        let failing = Arc::new(WriteFailsBackend(Arc::new(InMemoryBackend::new())));
        let migration =
            TenantMigration::new(TenantId(7), ShardId(0), ShardId(1), cp, Arc::clone(&router))
                .with_persistence(Arc::clone(&failing) as Arc<dyn crate::storage::StorageEngine>);

        migration.prepare_target(10).unwrap();
        let err = migration.enter_dual_write(20).unwrap_err();
        assert!(
            matches!(err, TidalError::Storage(_)),
            "a failed durable checkpoint must propagate"
        );

        // The router must NOT advertise a dual-write that never reached disk.
        assert!(
            !router.is_dual_write(TenantId(7)),
            "router must roll back the dual-write when persistence fails"
        );
        // The state machine stays in PreparingTarget so the caller can retry.
        assert!(matches!(
            migration.current_state(),
            MigrationState::PreparingTarget { .. }
        ));
    }

    /// Accuracy-W: a finalized migration's routing must survive a crash. The
    /// coordinator persists routing on every transition; a fresh router restored
    /// from the durable checkpoint keeps the pin instead of reverting to
    /// jump-hash routing.
    #[test]
    fn migration_routing_survives_simulated_crash() {
        use std::sync::RwLock;

        use crate::{
            replication::tenant::{ShardAssignment, restore_migration_routing},
            storage::InMemoryBackend,
        };

        let topo = Arc::new(RwLock::new(ClusterTopology {
            shards: vec![
                ShardAssignment {
                    shard_id: ShardId(0),
                    region_id: crate::replication::RegionId(0),
                },
                ShardAssignment {
                    shard_id: ShardId(1),
                    region_id: crate::replication::RegionId(1),
                },
            ],
        }));
        let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        let cp = Arc::new(ControlPlane::new(
            Arc::clone(&topo),
            Arc::clone(&router),
            lag,
        ));

        // Durable sink shared with the migration coordinator.
        let storage = Arc::new(InMemoryBackend::new());

        let migration = TenantMigration::new(
            TenantId(42),
            ShardId(0),
            ShardId(1),
            cp,
            Arc::clone(&router),
        )
        .with_persistence(Arc::clone(&storage) as Arc<dyn crate::storage::StorageEngine>);

        // Drive the migration to finalization. Each transition persists routing.
        migration.prepare_target(10).unwrap();
        migration.enter_dual_write(20).unwrap();
        migration.finalize(25).unwrap();

        // "Crash": discard the in-memory router and rebuild routing from disk.
        let recovered = TenantRouter::new(Arc::clone(&topo));
        let n = restore_migration_routing(storage.as_ref(), &recovered).unwrap();
        assert_eq!(
            n, 1,
            "the finalized pin must be restored from the checkpoint"
        );
        assert_eq!(
            recovered.pinned_shard(TenantId(42)),
            Some(ShardId(1)),
            "finalized tenant must stay pinned to its target after a crash, not \
             revert to jump-hash routing"
        );
        // It must NOT still be in dual-write (finalize cleared it).
        assert!(!recovered.is_dual_write(TenantId(42)));
    }

    /// A storage backend whose `write_batch`/`flush` start succeeding and can be
    /// flipped to fail at a chosen point. Reads pass through. Used to fail the
    /// finalize checkpoint while the dual-write checkpoint succeeded (C9).
    struct ToggleFailBackend {
        inner: Arc<crate::storage::InMemoryBackend>,
        fail: std::sync::atomic::AtomicBool,
    }

    impl ToggleFailBackend {
        fn new() -> Self {
            Self {
                inner: Arc::new(crate::storage::InMemoryBackend::new()),
                fail: std::sync::atomic::AtomicBool::new(false),
            }
        }
        fn start_failing(&self) {
            self.fail.store(true, std::sync::atomic::Ordering::SeqCst);
        }
        fn failing(&self) -> bool {
            self.fail.load(std::sync::atomic::Ordering::SeqCst)
        }
    }

    impl crate::storage::StorageEngine for ToggleFailBackend {
        fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, crate::storage::StorageError> {
            self.inner.get(key)
        }
        fn put(&self, key: &[u8], value: &[u8]) -> Result<(), crate::storage::StorageError> {
            if self.failing() {
                return Err(crate::storage::StorageError::Closed);
            }
            self.inner.put(key, value)
        }
        fn delete(&self, key: &[u8]) -> Result<(), crate::storage::StorageError> {
            self.inner.delete(key)
        }
        fn scan_prefix(&self, prefix: &[u8]) -> crate::storage::iterator::PrefixIterator<'_> {
            self.inner.scan_prefix(prefix)
        }
        fn write_batch(
            &self,
            batch: crate::storage::batch::WriteBatch,
        ) -> Result<(), crate::storage::StorageError> {
            if self.failing() {
                return Err(crate::storage::StorageError::Closed);
            }
            self.inner.write_batch(batch)
        }
        fn flush(&self) -> Result<(), crate::storage::StorageError> {
            if self.failing() {
                return Err(crate::storage::StorageError::Closed);
            }
            self.inner.flush()
        }
    }

    fn two_shard_parts() -> (
        Arc<std::sync::RwLock<ClusterTopology>>,
        Arc<TenantRouter>,
        Arc<ControlPlane>,
    ) {
        use std::sync::RwLock;

        use crate::replication::tenant::ShardAssignment;
        let topo = Arc::new(RwLock::new(ClusterTopology {
            shards: vec![
                ShardAssignment {
                    shard_id: ShardId(0),
                    region_id: crate::replication::RegionId(0),
                },
                ShardAssignment {
                    shard_id: ShardId(1),
                    region_id: crate::replication::RegionId(1),
                },
            ],
        }));
        let router = Arc::new(TenantRouter::new(Arc::clone(&topo)));
        let state = Arc::new(ReplicationState::single());
        let lag = Arc::new(ReplicationLagGauge::new(ShardId::SINGLE, state));
        let cp = Arc::new(ControlPlane::new(
            Arc::clone(&topo),
            Arc::clone(&router),
            lag,
        ));
        (topo, router, cp)
    }

    /// C9: a finalize whose checkpoint cannot be persisted must NOT observably
    /// finalize. It rolls the router's routing for this tenant back to the prior
    /// dual-write, stays in `DualWrite`, and returns `Storage` — so a crash after
    /// a failed finalize replays as dual-write (correct), never as a half-finalized
    /// pin the durable store does not reflect. The previous code swallowed the
    /// error and returned Ok while disk still held the pre-finalize routing.
    #[test]
    fn finalize_rolls_back_on_persist_failure() {
        let (_topo, router, cp) = two_shard_parts();
        let backend = Arc::new(ToggleFailBackend::new());
        let migration =
            TenantMigration::new(TenantId(7), ShardId(0), ShardId(1), cp, Arc::clone(&router))
                .with_persistence(Arc::clone(&backend) as Arc<dyn crate::storage::StorageEngine>);

        migration.prepare_target(10).unwrap();
        migration.enter_dual_write(20).unwrap(); // dual-write persists OK
        assert!(router.is_dual_write(TenantId(7)));

        // Now make the durable store reject the finalize checkpoint.
        backend.start_failing();
        let err = migration.finalize(25).unwrap_err();
        assert!(
            matches!(err, TidalError::Storage(_)),
            "a failed finalize checkpoint must propagate, not be swallowed; got {err:?}"
        );

        // The router must be rolled back to dual-write (NOT pinned), and the state
        // machine must stay in DualWrite so the caller can retry once storage
        // recovers.
        assert!(
            router.is_dual_write(TenantId(7)),
            "finalize rollback must restore the prior dual-write routing"
        );
        assert_eq!(
            router.pinned_shard(TenantId(7)),
            None,
            "a finalize that failed to persist must NOT leave the tenant pinned"
        );
        assert!(matches!(
            migration.current_state(),
            MigrationState::DualWrite { cutover_seqno: 20 }
        ));
    }

    /// C10: a failed `enter_dual_write` persist for tenant A must roll back ONLY
    /// tenant A's routing — it must NOT clobber a concurrent migration's routing
    /// for tenant B that was applied on the shared router between A's snapshot and
    /// A's rollback. The old full-router-snapshot rollback wiped B's entry.
    #[test]
    fn rollback_preserves_other_tenant_routing() {
        let (_topo, router, cp_a) = two_shard_parts();

        // Tenant 99 has a finalized pin already on the shared router (a different,
        // already-completed migration). This is what A's rollback must not touch.
        router.finalize_migration(TenantId(99), ShardId(1));
        assert_eq!(router.pinned_shard(TenantId(99)), Some(ShardId(1)));

        // Tenant 50 is mid-dual-write on the same shared router, modeling a
        // CONCURRENT migration's just-applied entry that A's rollback must leave
        // intact.
        router.set_dual_write(TenantId(50), ShardId(0), ShardId(1));

        // Tenant A's (id 7) migration uses a backend that fails the dual-write
        // persist, forcing the rollback path.
        let failing = Arc::new(ToggleFailBackend::new());
        failing.start_failing();
        let migration_a = TenantMigration::new(
            TenantId(7),
            ShardId(0),
            ShardId(1),
            cp_a,
            Arc::clone(&router),
        )
        .with_persistence(Arc::clone(&failing) as Arc<dyn crate::storage::StorageEngine>);

        migration_a.prepare_target(10).unwrap();
        let err = migration_a.enter_dual_write(20).unwrap_err();
        assert!(matches!(err, TidalError::Storage(_)));

        // A's own dual-write entry must be gone (rolled back)...
        assert!(
            !router.is_dual_write(TenantId(7)),
            "A's non-durable dual-write must be rolled back"
        );
        // ...but EVERY other tenant's routing on the shared router survives.
        assert_eq!(
            router.pinned_shard(TenantId(99)),
            Some(ShardId(1)),
            "C10: A's rollback must not clobber B's finalized pin"
        );
        assert!(
            router.is_dual_write(TenantId(50)),
            "C10: A's rollback must not clobber a concurrent tenant's dual-write"
        );
    }

    /// W16: `mark_complete` (and its `gc_source` alias) is a pure state flip — it
    /// transitions Finalizing -> Complete and does NOT reclaim source data. This
    /// pins the documented honest behavior (no source-shard delete happens here).
    #[test]
    fn mark_complete_is_state_only_transition() {
        let m = make_migration();
        m.prepare_target(10).unwrap();
        m.enter_dual_write(20).unwrap();
        m.finalize(25).unwrap();
        // The alias and the canonical name are the same transition.
        m.mark_complete(0).unwrap();
        assert_eq!(m.current_state(), MigrationState::Complete);
    }
}
