//! Lease creation and keep-alive management.

use std::sync::atomic::{AtomicU8, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use etcd_client::{Client, PutOptions};
use tokio::sync::watch;
use tokio::task::JoinHandle;
use tokio::time;
use tracing::{debug, error, warn};

use crate::connection::node_key;
use crate::error::ClusterError;
use crate::types::Node;

use super::ClusterNode;

/// Creates an etcd lease and spawns a keep-alive task.
///
/// Returns the lease ID and a handle to the keep-alive task.
pub(super) async fn create_lease(
    node: &ClusterNode,
    client: &mut Client,
    ttl: i64,
    lease_healthy_tx: watch::Sender<bool>,
) -> Result<(i64, JoinHandle<()>), ClusterError> {
    let lease = client.lease_grant(ttl, None).await?;
    let (mut keeper, mut stream) = client.lease_keep_alive(lease.id()).await?;

    let node = node.clone(); // Cheap Arc clone
    let missed_ka = AtomicU8::new(0);

    let handle = tokio::spawn(async move {
        loop {
            // Track if send failed to avoid double-counting errors
            // When both send and recv fail (broken connection),
            // we should only count it as one failure, not two.
            let mut send_failed = false;

            // Send keep-alive request
            // Note: We don't `continue` on send failure - we still call stream.message()
            // to maintain the request/response protocol and detect connection state.
            if let Err(e) = keeper.keep_alive().await {
                send_failed = true;
                let missed = missed_ka.fetch_add(1, Ordering::AcqRel) + 1;
                if missed > 10 {
                    error!("Stopping keep-alive after {} send failures: {}", missed, e);
                    lease_healthy_tx.send_replace(false);
                    break;
                } else if missed >= 3 {
                    warn!(
                        "Keep-alive send failed {} times - marking unhealthy: {}",
                        missed, e
                    );
                    node.set_unhealthy();
                } else {
                    warn!("Keep-alive send failed (attempt {}): {}", missed, e);
                }
                // Fall through to message() to drain any pending responses
                // and properly detect stream state
            }

            // Wait for keep-alive response (always called, even if send failed)
            // This ensures we:
            // 1. Consume any pending responses from previous successful sends
            // 2. Detect connection errors promptly
            // 3. Maintain consistent stream state
            match stream.message().await {
                Ok(Some(resp)) => {
                    debug!("Lease {} TTL {}", resp.id(), resp.ttl());
                    missed_ka.store(0, Ordering::Release);
                    node.set_healthy();
                }
                Ok(None) => {
                    debug!("Keep-alive stream closed by server");
                    lease_healthy_tx.send_replace(false);
                    break;
                }
                Err(e) => {
                    // Only count recv error if send succeeded (avoid double-counting)
                    if !send_failed {
                        let missed = missed_ka.fetch_add(1, Ordering::AcqRel) + 1;
                        if missed > 10 {
                            error!("Stopping keep-alive after {} recv failures: {}", missed, e);
                            lease_healthy_tx.send_replace(false);
                            break;
                        } else if missed >= 3 {
                            warn!(
                                "Keep-alive recv failed {} times - marking unhealthy: {}",
                                missed, e
                            );
                            node.set_unhealthy();
                        } else {
                            warn!("Keep-alive recv failed (attempt {}): {}", missed, e);
                        }
                    } else {
                        // Send already failed and was counted, just log recv failure
                        debug!("Keep-alive recv also failed (send already counted): {}", e);
                    }
                }
            }

            time::sleep(Duration::from_secs((ttl / 3).max(1) as u64)).await;
        }
    });

    Ok((lease.id(), handle))
}

/// Registers this node in the etcd registry.
///
/// The node is registered with its current metadata from the metrics collector.
/// If the collector returns null, no metadata is included.
pub(super) async fn register_local_node(
    node: &ClusterNode,
    client: &mut Client,
    lease_id: i64,
) -> Result<(), ClusterError> {
    let key = node_key(&node.inner().group_name, &node.inner().node_id);
    let last_seen = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0);

    // Collect metrics from the configured collector
    let metadata = node.inner().metrics_collector.collect();

    let meta = Node {
        id: node.inner().node_id.clone(),
        ip: node.inner().node_ip,
        last_seen,
        metadata,
    };
    let value = serde_json::to_vec(&meta)
        .map_err(|e| ClusterError::Internal(format!("Serialization error: {}", e)))?;

    client
        .put(key, value, Some(PutOptions::new().with_lease(lease_id)))
        .await?;
    Ok(())
}
