//! Leader election loop implementation.

use std::time::Duration;

use etcd_client::LeaderKey;
use etcd_client::ProclaimOptions;
use tokio::select;
use tokio::sync::{broadcast, watch};
use tokio::time;
use tracing::{debug, error, info, warn};

use crate::connection::{connect_with_retry, election_prefix};
use crate::error::ClusterError;

use super::ClusterNode;

/// Runs the leader election loop.
///
/// This function campaigns for leadership, and once elected, periodically
/// proclaims leadership via heartbeats. If leadership is lost or proclaim fails,
/// it will restart the election process.
pub(super) async fn run_election_loop(
    node: &ClusterNode,
    lease_id: i64,
    leader_tx: watch::Sender<Option<LeaderKey>>,
    mut shutdown_watcher: broadcast::Receiver<()>,
) -> Result<(), ClusterError> {
    loop {
        info!("🔵 Starting election loop ({})", node.inner().node_id);
        let client_watcher = shutdown_watcher.resubscribe();
        let client = connect_with_retry(&node.inner().etcd_endpoints, Some(client_watcher)).await;
        let Some(mut client) = client else {
            return Err(ClusterError::Internal("Failed to connect to etcd".into()));
        };

        // Wrap campaign in select! to allow graceful shutdown.
        // campaign() can block indefinitely waiting to become leader.
        // Without this, shutdown signals are ignored until leadership is acquired.
        let election_key = election_prefix(&node.inner().group_name);
        let campaign_resp = select! {
            _ = shutdown_watcher.recv() => {
                info!("Shutdown received while waiting for leadership");
                return Ok(()); // Graceful shutdown
            }
            result = client.campaign(
                election_key.as_bytes(),
                node.inner().node_id.as_bytes(),
                lease_id,
            ) => {
                match result {
                    Ok(resp) => resp,
                    Err(e) => {
                        break Err(ClusterError::Internal(format!(
                            "Failed to campaign for election: {}",
                            e
                        )));
                    }
                }
            }
        };

        let leader = campaign_resp
            .leader()
            .ok_or_else(|| ClusterError::Internal("Campaign response missing leader".into()))?;
        let _ = leader_tx.send(Some(leader.clone()));

        node.set_leader(true);
        info!("🟢 Became leader ({})", node.inner().node_id);

        // Proclaim loop
        let mut retries = 0;
        loop {
            select! {
                _ = shutdown_watcher.recv() => break,
                _ = time::sleep(Duration::from_secs((node.inner().ttl / 3).max(1) as u64)) => {
                    match client.proclaim(
                        node.inner().node_id.as_bytes(),
                        Some(ProclaimOptions::new().with_leader(leader.clone())),
                    ).await {
                        Ok(_) => {
                            retries = 0; // Reset retries on success
                            debug!("Proclaim succeeded");
                        }
                        Err(e) => {
                            retries += 1;
                            if retries > 3 {
                                error!("Proclaim failed after retries: {} - stepping down", e);
                                break;
                            }
                            warn!("Proclaim failed (retry {}): {}", retries, e);
                        }
                    }
                }
            }
        }

        // Cleanup this round
        node.resign(&mut client, leader.clone()).await.ok();
        let _ = leader_tx.send(None);
        node.set_leader(false);
        info!("🔴 Stepped down as leader, restarting election loop.");
    }
}
