//! Partition location service backed by etcd.
//!
//! Tracks which Kubernetes pod owns each graph partition and their health
//! status. Uses etcd as the distributed coordination store so all pods share
//! a consistent view of partition ownership.
//!
//! ## etcd key schema
//!
//! ```text
//! /ocg/partitions/{partition_id}/owner          → "ocg-partition-2"
//! /ocg/partitions/{partition_id}/status         → "healthy" | "recovering" | "failed"
//! /ocg/partitions/{partition_id}/last_heartbeat → unix timestamp (seconds)
//! /ocg/partitions/{partition_id}/storage_path   → "s3://bucket/graphs/g1/partitions/2"
//! ```

#[cfg(feature = "distributed")]
use std::collections::HashMap;
#[cfg(feature = "distributed")]
use std::sync::Arc;
#[cfg(feature = "distributed")]
use std::time::{SystemTime, UNIX_EPOCH};

#[cfg(feature = "distributed")]
use etcd_client::{Client, Compare, CompareOp, GetOptions, PutOptions, Txn, TxnOp};
#[cfg(feature = "distributed")]
use tokio::sync::RwLock;

#[cfg(feature = "distributed")]
use crate::error::{CypherError, CypherResult, ExecutionError};

// ── Types ─────────────────────────────────────────────────────────────────────

/// Health status of a partition.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PartitionStatus {
    Healthy,
    Recovering,
    Failed,
}

#[cfg(feature = "distributed")]
impl PartitionStatus {
    fn as_str(&self) -> &'static str {
        match self {
            PartitionStatus::Healthy    => "healthy",
            PartitionStatus::Recovering => "recovering",
            PartitionStatus::Failed     => "failed",
        }
    }

    fn from_str(s: &str) -> Self {
        match s {
            "recovering" => PartitionStatus::Recovering,
            "failed"     => PartitionStatus::Failed,
            _            => PartitionStatus::Healthy,
        }
    }
}

/// Metadata for a single partition as stored in etcd.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct PartitionLocation {
    pub partition_id:   u32,
    pub owner:          String,
    pub status:         PartitionStatus,
    pub last_heartbeat: u64,
    pub storage_path:   String,
}

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

/// Client for the etcd-backed partition location service.
///
/// Maintains an in-memory cache of partition → pod mappings that is kept
/// up-to-date by periodic refreshes and watch events.
#[cfg(feature = "distributed")]
pub struct PartitionLocationService {
    client:  Client,
    /// In-memory cache: partition_id → location
    cache:   Arc<RwLock<HashMap<u32, PartitionLocation>>>,
}

#[cfg(feature = "distributed")]
impl PartitionLocationService {
    // ── Construction ──────────────────────────────────────────────────────

    /// Connect to an etcd cluster.
    ///
    /// `endpoints` examples: `["http://etcd:2379"]`,
    ///                        `["http://127.0.0.1:2379"]`
    pub async fn connect(endpoints: Vec<String>) -> CypherResult<Self> {
        let client = Client::connect(endpoints, None)
            .await
            .map_err(|e| etcd_err(format!("connect: {e}")))?;
        Ok(Self {
            client,
            cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Construct a [`PartitionLocationService`] from a pre-connected etcd client.
    ///
    /// Use this when the client was already established (e.g. via
    /// [`crate::distributed::partition::etcd_retry::connect_with_retry`]) to
    /// avoid a redundant connection attempt.
    pub fn from_client(client: Client) -> Self {
        Self {
            client,
            cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    // ── Registration ──────────────────────────────────────────────────────

    /// Register this pod as the owner of `partition_id`.
    pub async fn register_partition(
        &mut self,
        partition_id:  u32,
        pod_name:      String,
        storage_path:  String,
    ) -> CypherResult<()> {
        let ts = unix_now();
        self.put(key(partition_id, "owner"),          pod_name.clone()).await?;
        self.put(key(partition_id, "status"),         "healthy".to_string()).await?;
        self.put(key(partition_id, "last_heartbeat"), ts.to_string()).await?;
        self.put(key(partition_id, "storage_path"),   storage_path.clone()).await?;

        // Update local cache
        let loc = PartitionLocation {
            partition_id,
            owner: pod_name,
            status: PartitionStatus::Healthy,
            last_heartbeat: ts,
            storage_path,
        };
        self.cache.write().await.insert(partition_id, loc);
        Ok(())
    }

    // ── Status updates ────────────────────────────────────────────────────

    /// Update the health status of a partition.
    pub async fn update_partition_status(
        &mut self,
        partition_id: u32,
        status:       PartitionStatus,
    ) -> CypherResult<()> {
        self.put(key(partition_id, "status"), status.as_str().to_string()).await?;
        if let Some(loc) = self.cache.write().await.get_mut(&partition_id) {
            loc.status = status;
        }
        Ok(())
    }

    /// Record a heartbeat for `partition_id`.
    pub async fn record_heartbeat(&mut self, partition_id: u32) -> CypherResult<()> {
        let ts = unix_now();
        self.put(key(partition_id, "last_heartbeat"), ts.to_string()).await?;
        if let Some(loc) = self.cache.write().await.get_mut(&partition_id) {
            loc.last_heartbeat = ts;
        }
        Ok(())
    }

    // ── Lookups ───────────────────────────────────────────────────────────

    /// Return the owner pod name for `partition_id`, from the cache if available.
    pub async fn get_partition_owner(&self, partition_id: u32) -> Option<String> {
        self.cache.read().await
            .get(&partition_id)
            .map(|loc| loc.owner.clone())
    }

    /// Return the full location info for `partition_id`.
    pub async fn get_partition_location(&self, partition_id: u32) -> Option<PartitionLocation> {
        self.cache.read().await
            .get(&partition_id)
            .cloned()
    }

    /// Fetch and cache the full partition directory from etcd.
    ///
    /// Call this on startup and after suspected changes to rebuild the cache.
    pub async fn refresh_cache(&mut self) -> CypherResult<()> {
        let prefix = "/ocg/partitions/";
        let opts   = GetOptions::new().with_prefix();
        let resp   = self.client.get(prefix, Some(opts)).await
            .map_err(|e| etcd_err(format!("refresh_cache: {e}")))?;

        let mut raw: HashMap<u32, HashMap<String, String>> = HashMap::new();
        for kv in resp.kvs() {
            let full_key = kv.key_str()
                .map_err(|e| etcd_err(format!("key utf-8: {e}")))?;
            // key format: /ocg/partitions/{id}/{field}
            let parts: Vec<&str> = full_key.splitn(5, '/').collect();
            // parts = ["", "ocg", "partitions", "{id}", "{field}"]
            if parts.len() != 5 { continue; }
            let Ok(pid) = parts[3].parse::<u32>() else { continue; };
            let field   = parts[4].to_string();
            let value   = kv.value_str()
                .map_err(|e| etcd_err(format!("value utf-8: {e}")))?
                .to_string();
            raw.entry(pid).or_default().insert(field, value);
        }

        let mut cache = self.cache.write().await;
        cache.clear();
        for (pid, fields) in raw {
            let loc = PartitionLocation {
                partition_id:   pid,
                owner:          fields.get("owner").cloned().unwrap_or_default(),
                status:         PartitionStatus::from_str(
                    fields.get("status").map(String::as_str).unwrap_or("healthy")
                ),
                last_heartbeat: fields.get("last_heartbeat")
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0),
                storage_path:   fields.get("storage_path").cloned().unwrap_or_default(),
            };
            cache.insert(pid, loc);
        }
        Ok(())
    }

    /// Return all partitions currently marked as failed.
    pub async fn failed_partitions(&self) -> Vec<u32> {
        self.cache.read().await
            .values()
            .filter(|loc| loc.status == PartitionStatus::Failed)
            .map(|loc| loc.partition_id)
            .collect()
    }

    /// Return partitions whose heartbeat is older than `max_age_secs`.
    pub async fn stale_partitions(&self, max_age_secs: u64) -> Vec<u32> {
        let now = unix_now();
        self.cache.read().await
            .values()
            .filter(|loc| {
                loc.status == PartitionStatus::Healthy
                    && now.saturating_sub(loc.last_heartbeat) > max_age_secs
            })
            .map(|loc| loc.partition_id)
            .collect()
    }

    /// Build a pod → partition_ids assignment map from the cache.
    pub async fn assignment_map(&self) -> HashMap<String, Vec<u32>> {
        let mut map: HashMap<String, Vec<u32>> = HashMap::new();
        for loc in self.cache.read().await.values() {
            map.entry(loc.owner.clone()).or_default().push(loc.partition_id);
        }
        map
    }

    // ── CAS ownership ─────────────────────────────────────────────────

    /// Claim ownership of a partition using etcd compare-and-swap.
    ///
    /// Succeeds only if the current owner in etcd matches `expected_owner`.
    /// This prevents two pods from simultaneously claiming the same partition.
    ///
    /// If `lease_id` is `Some`, the ownership key is attached to that etcd
    /// lease — when the lease expires (pod dies), the key is auto-deleted.
    ///
    /// Returns `Ok(true)` if ownership was claimed, `Ok(false)` if the CAS
    /// failed (another pod already claimed it).
    pub async fn claim_partition(
        &mut self,
        partition_id:   u32,
        expected_owner: &str,
        new_owner:      String,
        storage_path:   String,
        lease_id:       Option<i64>,
    ) -> CypherResult<bool> {
        let owner_key = key(partition_id, "owner");
        let ts = unix_now();

        // Build put options — attach lease if provided
        let put_opts = lease_id.map(|lid| PutOptions::new().with_lease(lid));

        // CAS: only update if current owner == expected_owner
        let cmp = Compare::value(owner_key.clone(), CompareOp::Equal, expected_owner);
        let put_owner  = TxnOp::put(owner_key.clone(), new_owner.as_bytes(), put_opts.clone());
        let put_status = TxnOp::put(key(partition_id, "status"), "healthy", put_opts.clone());
        let put_hb     = TxnOp::put(key(partition_id, "last_heartbeat"), ts.to_string(), put_opts.clone());
        let put_path   = TxnOp::put(key(partition_id, "storage_path"), storage_path.as_bytes(), put_opts);

        let txn = Txn::new()
            .when([cmp])
            .and_then([put_owner, put_status, put_hb, put_path])
            .or_else([]);

        let resp = self.client.txn(txn).await
            .map_err(|e| etcd_err(format!("claim_partition txn: {e}")))?;

        if resp.succeeded() {
            let loc = PartitionLocation {
                partition_id,
                owner: new_owner,
                status: PartitionStatus::Healthy,
                last_heartbeat: ts,
                storage_path,
            };
            self.cache.write().await.insert(partition_id, loc);
            Ok(true)
        } else {
            tracing::warn!(
                partition = partition_id,
                expected_owner,
                "CAS claim failed — another pod owns this partition"
            );
            Ok(false)
        }
    }

    // ── Lease management ──────────────────────────────────────────────

    /// Grant an etcd lease with the given TTL (seconds).
    ///
    /// The caller should keep the lease alive by calling `keep_alive_lease()`
    /// periodically (at intervals < TTL). If the lease expires, all keys
    /// attached to it are automatically deleted by etcd.
    pub async fn grant_lease(&mut self, ttl_secs: i64) -> CypherResult<i64> {
        let resp = self.client.lease_grant(ttl_secs, None).await
            .map_err(|e| etcd_err(format!("lease_grant: {e}")))?;
        Ok(resp.id())
    }

    /// Send a single keep-alive for the given lease.
    pub async fn keep_alive_lease(&mut self, lease_id: i64) -> CypherResult<()> {
        let (mut keeper, _stream) = self.client.lease_keep_alive(lease_id).await
            .map_err(|e| etcd_err(format!("lease_keep_alive: {e}")))?;
        keeper.keep_alive().await
            .map_err(|e| etcd_err(format!("keep_alive send: {e}")))?;
        Ok(())
    }

    /// Register a partition with ownership attached to an etcd lease.
    ///
    /// Like `register_partition()` but all keys are associated with the
    /// lease so they auto-delete when the lease expires.
    pub async fn register_partition_with_lease(
        &mut self,
        partition_id:  u32,
        pod_name:      String,
        storage_path:  String,
        lease_id:      i64,
    ) -> CypherResult<()> {
        let ts = unix_now();
        let opts = Some(PutOptions::new().with_lease(lease_id));
        self.put_with_opts(key(partition_id, "owner"),          pod_name.clone(), opts.clone()).await?;
        self.put_with_opts(key(partition_id, "status"),         "healthy".to_string(), opts.clone()).await?;
        self.put_with_opts(key(partition_id, "last_heartbeat"), ts.to_string(), opts.clone()).await?;
        self.put_with_opts(key(partition_id, "storage_path"),   storage_path.clone(), opts).await?;

        let loc = PartitionLocation {
            partition_id,
            owner: pod_name,
            status: PartitionStatus::Healthy,
            last_heartbeat: ts,
            storage_path,
        };
        self.cache.write().await.insert(partition_id, loc);
        Ok(())
    }

    /// Record a heartbeat attached to an etcd lease.
    pub async fn record_heartbeat_with_lease(
        &mut self,
        partition_id: u32,
        lease_id: i64,
    ) -> CypherResult<()> {
        let ts = unix_now();
        let opts = Some(PutOptions::new().with_lease(lease_id));
        self.put_with_opts(key(partition_id, "last_heartbeat"), ts.to_string(), opts).await?;
        if let Some(loc) = self.cache.write().await.get_mut(&partition_id) {
            loc.last_heartbeat = ts;
        }
        Ok(())
    }

    // ── Internal helpers ──────────────────────────────────────────────────

    async fn put(&mut self, k: String, v: String) -> CypherResult<()> {
        self.client.put(k, v, None::<PutOptions>).await
            .map_err(|e| etcd_err(format!("put: {e}")))?;
        Ok(())
    }

    async fn put_with_opts(
        &mut self,
        k: String,
        v: String,
        opts: Option<PutOptions>,
    ) -> CypherResult<()> {
        self.client.put(k, v, opts).await
            .map_err(|e| etcd_err(format!("put: {e}")))?;
        Ok(())
    }
}

// ── Free helpers ──────────────────────────────────────────────────────────────

#[cfg(feature = "distributed")]
fn key(partition_id: u32, field: &str) -> String {
    format!("/ocg/partitions/{partition_id}/{field}")
}

#[cfg(feature = "distributed")]
fn unix_now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

#[cfg(feature = "distributed")]
fn etcd_err(msg: impl Into<String>) -> CypherError {
    CypherError::Execution(ExecutionError::Internal(msg.into()))
}

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

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn test_partition_status_round_trip() {
        for (s, expected) in [
            ("healthy",    PartitionStatus::Healthy),
            ("recovering", PartitionStatus::Recovering),
            ("failed",     PartitionStatus::Failed),
            ("unknown",    PartitionStatus::Healthy), // default
        ] {
            assert_eq!(PartitionStatus::from_str(s), expected, "status: {s}");
        }
    }

    #[test]
    fn test_status_as_str() {
        assert_eq!(PartitionStatus::Healthy.as_str(),    "healthy");
        assert_eq!(PartitionStatus::Recovering.as_str(), "recovering");
        assert_eq!(PartitionStatus::Failed.as_str(),     "failed");
    }

    #[test]
    fn test_key_format() {
        assert_eq!(key(3, "owner"), "/ocg/partitions/3/owner");
        assert_eq!(key(0, "status"), "/ocg/partitions/0/status");
    }

    #[test]
    fn test_unix_now_nonzero() {
        assert!(unix_now() > 0);
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_register_and_lookup_live() {
        let mut svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();

        svc.register_partition(
            1,
            "ocg-partition-1".to_string(),
            "s3://test-bucket/graphs/g1/partitions/1".to_string(),
        ).await.unwrap();

        let owner = svc.get_partition_owner(1).await;
        assert_eq!(owner.as_deref(), Some("ocg-partition-1"));
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_claim_partition_cas_live() {
        let mut svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();

        // Register initial owner
        svc.register_partition(
            200,
            "old-pod".to_string(),
            "s3://test/p200".to_string(),
        ).await.unwrap();

        // CAS succeeds when expected_owner matches
        let ok = svc.claim_partition(200, "old-pod", "new-pod".to_string(), "s3://test/p200".to_string(), None)
            .await.unwrap();
        assert!(ok, "CAS should succeed with matching owner");

        let owner = svc.get_partition_owner(200).await;
        assert_eq!(owner.as_deref(), Some("new-pod"));

        // CAS fails when expected_owner is stale
        let fail = svc.claim_partition(200, "old-pod", "rogue-pod".to_string(), "s3://test/p200".to_string(), None)
            .await.unwrap();
        assert!(!fail, "CAS should fail with stale owner");

        // Owner should still be new-pod
        let owner = svc.get_partition_owner(200).await;
        assert_eq!(owner.as_deref(), Some("new-pod"));
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_lease_registration_and_expiry_live() {
        let mut svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();

        // Grant a very short lease (2 seconds)
        let lease_id = svc.grant_lease(2).await.unwrap();
        assert!(lease_id > 0);

        svc.register_partition_with_lease(
            201,
            "lease-pod".to_string(),
            "s3://test/p201".to_string(),
            lease_id,
        ).await.unwrap();

        let owner = svc.get_partition_owner(201).await;
        assert_eq!(owner.as_deref(), Some("lease-pod"));

        // After lease expires, etcd deletes the keys.
        // We wait 3s and refresh cache to observe deletion.
        tokio::time::sleep(Duration::from_secs(3)).await;
        svc.refresh_cache().await.unwrap();

        // The partition should be gone from cache after lease expiry
        let owner = svc.get_partition_owner(201).await;
        assert!(owner.is_none(), "partition should be absent after lease expiry");
    }

    #[tokio::test]
    #[ignore = "requires a running etcd server on localhost:2379"]
    async fn test_assignment_map_live() {
        let mut svc = PartitionLocationService::connect(
            vec!["http://127.0.0.1:2379".to_string()]
        ).await.unwrap();

        svc.register_partition(300, "pod-a".to_string(), "/p300".to_string()).await.unwrap();
        svc.register_partition(301, "pod-a".to_string(), "/p301".to_string()).await.unwrap();
        svc.register_partition(302, "pod-b".to_string(), "/p302".to_string()).await.unwrap();

        let map = svc.assignment_map().await;
        assert_eq!(map.get("pod-a").map(|v| v.len()), Some(2));
        assert_eq!(map.get("pod-b").map(|v| v.len()), Some(1));
    }
}
