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

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

use super::control::ControlPlane;
use super::shard::ShardId;
use super::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>,
}

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,
        }
    }

    /// 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`.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current state is not `PreparingTarget`.
    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(),
            ));
        }
        self.tenant_router
            .set_dual_write(self.tenant_id, self.source_shard, self.target_shard);
        *state = MigrationState::DualWrite { cutover_seqno };
        drop(state);
        Ok(cutover_seqno)
    }

    /// Transition from `DualWrite` to `Finalizing`.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current state is not `DualWrite`.
    /// Returns `NotReady` if `target_seqno` is behind the cutover sequence number —
    /// the caller should wait and retry once the target has caught up.
    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}"
            )));
        }
        self.tenant_router
            .finalize_migration(self.tenant_id, self.target_shard);
        let switched_at_ns = crate::replication::now_ns();
        *state = MigrationState::Finalizing { switched_at_ns };
        drop(state);
        Ok(())
    }

    /// Transition from `Finalizing` to `Complete` after the GC window has elapsed.
    ///
    /// # Errors
    ///
    /// Returns `InvalidState` if the current state is not `Finalizing`, or if
    /// `gc_window_ns` has not elapsed since the routing switch.
    pub fn gc_source(&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(
                "gc_source 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(())
    }

    /// 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;
    use crate::replication::lag::ReplicationLagGauge;
    use crate::replication::state::ReplicationState;
    use crate::replication::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(_)));
    }
}
