//! Heartbeat service — updates the partition's `last_heartbeat` timestamp
//! in etcd every `interval_secs` seconds.
//!
//! Each partition pod runs one `HeartbeatService` per partition it owns.
//! The coordinator's `FailureDetector` watches these timestamps and marks
//! a partition as failed if no heartbeat is received within the timeout.

#[cfg(feature = "distributed")]
use std::sync::Arc;
#[cfg(feature = "distributed")]
use std::time::Duration;

#[cfg(feature = "distributed")]
use tokio::sync::RwLock;
#[cfg(feature = "distributed")]
use tokio::task::JoinHandle;
#[cfg(feature = "distributed")]
use tokio_util::sync::CancellationToken;

#[cfg(feature = "distributed")]
use crate::distributed::partition::location_service::PartitionLocationService;
#[cfg(feature = "distributed")]
use crate::error::CypherResult;

// ── Service ───────────────────────────────────────────────────────────────────

/// Emits heartbeats for a set of partitions at a fixed interval.
///
/// Supports two modes:
/// - **Timestamp mode** (`start()`): updates the `last_heartbeat` key in etcd.
/// - **Lease mode** (`start_with_lease()`): additionally sends etcd lease
///   keep-alive messages so the lease (and all keys attached to it) auto-expire
///   if this pod dies.
#[cfg(feature = "distributed")]
pub struct HeartbeatService {
    location_svc: Arc<RwLock<PartitionLocationService>>,
    partition_ids: Vec<u32>,
    interval:      Duration,
    cancel:        CancellationToken,
    handle:        Option<JoinHandle<()>>,
}

#[cfg(feature = "distributed")]
impl HeartbeatService {
    /// Create a new heartbeat service.
    ///
    /// - `partition_ids`: partitions this pod is responsible for.
    /// - `interval_secs`: how often (in seconds) to update the heartbeat.
    pub fn new(
        location_svc:  Arc<RwLock<PartitionLocationService>>,
        partition_ids: Vec<u32>,
        interval_secs: u64,
    ) -> Self {
        Self {
            location_svc,
            partition_ids,
            interval: Duration::from_secs(interval_secs),
            cancel:   CancellationToken::new(),
            handle:   None,
        }
    }

    /// Start the heartbeat loop in a background tokio task.
    ///
    /// Returns `Err` if the service is already running.
    pub fn start(&mut self) -> CypherResult<()> {
        if self.handle.is_some() {
            return Ok(()); // already running
        }

        let svc       = Arc::clone(&self.location_svc);
        let ids       = self.partition_ids.clone();
        let interval  = self.interval;
        let cancel    = self.cancel.clone();

        let handle = tokio::spawn(async move {
            let mut ticker = tokio::time::interval(interval);
            loop {
                tokio::select! {
                    _ = ticker.tick() => {
                        let mut guard = svc.write().await;
                        for &pid in &ids {
                            if let Err(e) = guard.record_heartbeat(pid).await {
                                tracing::warn!(
                                    partition = pid,
                                    error = %e,
                                    "heartbeat update failed"
                                );
                            }
                        }
                    }
                    _ = cancel.cancelled() => {
                        tracing::info!("heartbeat service shutting down");
                        break;
                    }
                }
            }
        });

        self.handle = Some(handle);
        Ok(())
    }

    /// Start the heartbeat loop with etcd lease keep-alive.
    ///
    /// In addition to updating the `last_heartbeat` timestamp, this mode
    /// sends a keep-alive for the given etcd lease on each tick. If the pod
    /// dies without calling `stop()`, the lease expires after its TTL and
    /// etcd automatically deletes all keys attached to it — effectively
    /// marking the partition as unavailable.
    pub fn start_with_lease(&mut self, lease_id: i64) -> CypherResult<()> {
        if self.handle.is_some() {
            return Ok(());
        }

        let svc       = Arc::clone(&self.location_svc);
        let ids       = self.partition_ids.clone();
        let interval  = self.interval;
        let cancel    = self.cancel.clone();

        let handle = tokio::spawn(async move {
            let mut ticker = tokio::time::interval(interval);
            loop {
                tokio::select! {
                    _ = ticker.tick() => {
                        let mut guard = svc.write().await;
                        // Keep the lease alive
                        if let Err(e) = guard.keep_alive_lease(lease_id).await {
                            tracing::warn!(
                                lease_id,
                                error = %e,
                                "lease keep-alive failed"
                            );
                        }
                        // Update partition heartbeat timestamps
                        for &pid in &ids {
                            if let Err(e) = guard.record_heartbeat_with_lease(pid, lease_id).await {
                                tracing::warn!(
                                    partition = pid,
                                    error = %e,
                                    "heartbeat update failed"
                                );
                            }
                        }
                    }
                    _ = cancel.cancelled() => {
                        tracing::info!("heartbeat service shutting down");
                        break;
                    }
                }
            }
        });

        self.handle = Some(handle);
        Ok(())
    }

    /// Stop the heartbeat loop gracefully.
    pub fn stop(&mut self) {
        self.cancel.cancel();
        self.handle = None;
    }
}

#[cfg(feature = "distributed")]
impl Drop for HeartbeatService {
    fn drop(&mut self) {
        self.cancel.cancel();
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};

    #[test]
    fn test_heartbeat_config() {
        // Verify that HeartbeatService can be constructed with correct parameters
        // (without actually connecting to etcd)
        let interval = Duration::from_secs(5);
        assert_eq!(interval.as_secs(), 5);
    }

    #[test]
    fn test_cancellation_token() {
        let token = CancellationToken::new();
        assert!(!token.is_cancelled());
        token.cancel();
        assert!(token.is_cancelled());
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_heartbeat_updates_timestamp_live() {
        use crate::distributed::partition::location_service::PartitionLocationService;

        let svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();
        let svc = Arc::new(RwLock::new(svc));

        // Register partition first
        {
            let mut guard = svc.write().await;
            guard.register_partition(
                42,
                "ocg-test-pod".to_string(),
                "s3://test/p42".to_string(),
            ).await.unwrap();
        }

        let t0 = {
            let guard = svc.read().await;
            guard.get_partition_location(42).await.unwrap().last_heartbeat
        };

        let mut hb = HeartbeatService::new(Arc::clone(&svc), vec![42], 1);
        hb.start().unwrap();
        tokio::time::sleep(Duration::from_millis(1500)).await;
        hb.stop();

        let t1 = {
            let guard = svc.read().await;
            guard.get_partition_location(42).await.unwrap().last_heartbeat
        };

        assert!(t1 >= t0, "heartbeat should advance timestamp");
    }

    #[test]
    fn test_atomic_counter() {
        // Simple counter to verify atomic operations work as expected in tests
        let counter = Arc::new(AtomicU32::new(0));
        counter.fetch_add(1, Ordering::SeqCst);
        assert_eq!(counter.load(Ordering::SeqCst), 1);
    }
}
