//! Main watch loop implementation for ServiceDiscovery.

use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use etcd_client::{Client, GetOptions, SortOrder, SortTarget, WatchOptions};
use tokio::sync::broadcast;
use tracing::{debug, info, warn};

use crate::connection::{connect_with_retry, election_prefix, nodes_prefix};
use crate::error::ClusterError;
use crate::event::ClusterEvent;
use crate::types::Node;

use super::handlers::{emit, emit_diff_events, handle_election_events, handle_node_events};
use super::ServiceDiscovery;

/// Default timeout for etcd operations.
const ETCD_OPERATION_TIMEOUT: Duration = Duration::from_secs(5);

/// Runs the main discovery loop.
///
/// This function connects to etcd, loads initial state, sets up watches,
/// and processes events. It automatically reconnects on connection failures.
pub(super) async fn run(
    discovery: &ServiceDiscovery,
    mut shutdown: Option<&mut broadcast::Receiver<()>>,
) -> Result<(), ClusterError> {
    let nodes_prefix = nodes_prefix(&discovery.inner.group_name);
    let election_key = election_prefix(&discovery.inner.group_name);

    loop {
        debug!(
            "Starting discovery loop for group '{}'",
            discovery.inner.group_name
        );

        // Check shutdown before each iteration
        if let Some(ref mut rx) = shutdown {
            if rx.try_recv().is_ok() {
                info!("ServiceDiscovery shutdown requested");
                emit(discovery, ClusterEvent::Disconnected);
                return Ok(());
            }
        }

        // Connect to etcd (immediate connection attempt in connect_with_retry)
        let shutdown_rx = shutdown.as_mut().map(|rx| rx.resubscribe());
        let client = connect_with_retry(&discovery.inner.etcd_endpoints, shutdown_rx).await;
        let Some(mut client) = client else {
            return Err(ClusterError::Internal("Failed to connect to etcd".into()));
        };

        // Emit Reconnected event (except on first connect)
        if !discovery.inner.first_connect.swap(false, Ordering::AcqRel) {
            emit(discovery, ClusterEvent::Reconnected);
        }

        // Don't reset state - instead, load new state and diff
        // Store previous state for comparison
        let prev_nodes = discovery.nodes();
        let prev_leader = discovery.leader();

        // Load current nodes from etcd with timeout
        let registry_rev = match tokio::time::timeout(
            ETCD_OPERATION_TIMEOUT,
            load_initial_nodes(discovery, &mut client, &nodes_prefix),
        )
        .await
        {
            Ok(Ok(rev)) => rev,
            Ok(Err(e)) => {
                warn!("Failed to load nodes: {}", e);
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
            Err(_) => {
                warn!("Timeout loading nodes");
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        // Emit events only for actual differences after reconnection
        emit_diff_events(discovery, &prev_nodes, &discovery.nodes());

        // Load current leader from etcd with timeout
        match tokio::time::timeout(
            ETCD_OPERATION_TIMEOUT,
            load_initial_leader(discovery, &mut client, &election_key),
        )
        .await
        {
            Ok(Ok(())) => {}
            Ok(Err(e)) => {
                warn!("Failed to load leader: {}", e);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
            Err(_) => {
                warn!("Timeout loading leader");
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        }

        // Emit leader change event only if actually changed
        let new_leader = discovery.leader();
        if prev_leader.as_ref().map(|l| &l.id) != new_leader.as_ref().map(|l| &l.id) {
            match (&prev_leader, &new_leader) {
                (Some(_), None) => emit(discovery, ClusterEvent::LeaderLost),
                (_, Some(leader)) => emit(discovery, ClusterEvent::LeaderElected(leader.clone())),
                _ => {}
            }
        }

        // Set ready state via watch channel
        discovery.inner.ready_tx.send_replace(true);
        emit(discovery, ClusterEvent::Ready);

        // Set up watches with timeout
        let nodes_stream = match tokio::time::timeout(ETCD_OPERATION_TIMEOUT, async {
            client
                .watch(
                    nodes_prefix.clone(),
                    Some(
                        WatchOptions::new()
                            .with_prefix()
                            .with_start_revision(registry_rev + 1),
                    ),
                )
                .await
        })
        .await
        {
            Ok(Ok((_, stream))) => stream,
            Ok(Err(e)) => {
                warn!("Failed to watch nodes: {}", e);
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
            Err(_) => {
                warn!("Timeout setting up nodes watch");
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        let election_stream = match tokio::time::timeout(ETCD_OPERATION_TIMEOUT, async {
            client
                .watch(
                    election_key.clone(),
                    Some(
                        WatchOptions::new()
                            .with_prefix()
                            .with_start_revision(registry_rev + 1),
                    ),
                )
                .await
        })
        .await
        {
            Ok(Ok((_, stream))) => stream,
            Ok(Err(e)) => {
                warn!("Failed to watch election: {}", e);
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
            Err(_) => {
                warn!("Timeout setting up election watch");
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        // Run watch loop with shutdown support
        let watch_result = if let Some(ref mut rx) = shutdown {
            watch_loop_with_shutdown(
                discovery,
                &mut client,
                nodes_stream,
                election_stream,
                &nodes_prefix,
                &election_key,
                rx,
            )
            .await
        } else {
            watch_loop(
                discovery,
                &mut client,
                nodes_stream,
                election_stream,
                &nodes_prefix,
                &election_key,
            )
            .await
        };

        match watch_result {
            Ok(()) => {
                // Graceful shutdown
                info!("ServiceDiscovery shutting down gracefully");
                return Ok(());
            }
            Err(e) => {
                warn!("Watch loop error: {}", e);
                emit(discovery, ClusterEvent::Disconnected);
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        }
    }
}

/// Loads initial nodes from etcd.
async fn load_initial_nodes(
    discovery: &ServiceDiscovery,
    client: &mut Client,
    prefix: &str,
) -> Result<i64, ClusterError> {
    let resp = client
        .get(prefix, Some(GetOptions::new().with_prefix()))
        .await?;

    let mut nodes = Vec::new();
    for kv in resp.kvs() {
        match serde_json::from_slice::<Node>(kv.value()) {
            Ok(node) => {
                info!("Loaded node: {}", node.id);
                nodes.push(node);
            }
            Err(e) => {
                let key = String::from_utf8_lossy(kv.key());
                warn!("Malformed node at {}: {}", key, e);
            }
        }
    }

    discovery.inner.nodes_tx.send_replace(Arc::new(nodes));

    let revision = resp
        .header()
        .ok_or_else(|| ClusterError::Internal("Missing etcd response header".into()))?
        .revision();

    Ok(revision)
}

/// Loads initial leader from etcd.
async fn load_initial_leader(
    discovery: &ServiceDiscovery,
    client: &mut Client,
    election_key: &str,
) -> Result<(), ClusterError> {
    let resp = client
        .get(
            election_key,
            Some(
                GetOptions::new()
                    .with_prefix()
                    .with_sort(SortTarget::Create, SortOrder::Ascend)
                    .with_limit(1),
            ),
        )
        .await?;

    if let Some(kv) = resp.kvs().first() {
        if let Ok(leader_id) = String::from_utf8(kv.value().to_vec()) {
            let nodes = discovery.nodes();
            if let Some(leader) = nodes.iter().find(|n| n.id == leader_id).cloned() {
                info!("Current leader: {}", leader.id);
                discovery.inner.leader_tx.send_replace(Some(leader));
            }
        }
    }

    Ok(())
}

/// Main watch loop without shutdown support.
async fn watch_loop(
    discovery: &ServiceDiscovery,
    client: &mut Client,
    mut nodes_stream: etcd_client::WatchStream,
    mut election_stream: etcd_client::WatchStream,
    nodes_prefix: &str,
    election_key: &str,
) -> Result<(), ClusterError> {
    loop {
        tokio::select! {
            msg = nodes_stream.message() => {
                match msg {
                    Ok(Some(resp)) => {
                        handle_node_events(discovery, resp.events(), nodes_prefix);
                    }
                    Ok(None) => {
                        warn!("Nodes watch stream closed");
                        return Err(ClusterError::Internal("Nodes stream EOF".into()));
                    }
                    Err(e) => {
                        warn!("Nodes watch error: {}", e);
                        return Err(ClusterError::Etcd(e));
                    }
                }
            }
            msg = election_stream.message() => {
                match msg {
                    Ok(Some(resp)) => {
                        handle_election_events(discovery, client, resp.events(), election_key).await;
                    }
                    Ok(None) => {
                        warn!("Election watch stream closed");
                        return Err(ClusterError::Internal("Election stream EOF".into()));
                    }
                    Err(e) => {
                        warn!("Election watch error: {}", e);
                        return Err(ClusterError::Etcd(e));
                    }
                }
            }
        }
    }
}

/// Watch loop with shutdown support.
async fn watch_loop_with_shutdown(
    discovery: &ServiceDiscovery,
    client: &mut Client,
    mut nodes_stream: etcd_client::WatchStream,
    mut election_stream: etcd_client::WatchStream,
    nodes_prefix: &str,
    election_key: &str,
    shutdown: &mut broadcast::Receiver<()>,
) -> Result<(), ClusterError> {
    loop {
        tokio::select! {
            _ = shutdown.recv() => {
                info!("Shutdown signal received in watch loop");
                return Ok(()); // Graceful shutdown
            }
            msg = nodes_stream.message() => {
                match msg {
                    Ok(Some(resp)) => {
                        handle_node_events(discovery, resp.events(), nodes_prefix);
                    }
                    Ok(None) => {
                        warn!("Nodes watch stream closed");
                        return Err(ClusterError::Internal("Nodes stream EOF".into()));
                    }
                    Err(e) => {
                        warn!("Nodes watch error: {}", e);
                        return Err(ClusterError::Etcd(e));
                    }
                }
            }
            msg = election_stream.message() => {
                match msg {
                    Ok(Some(resp)) => {
                        handle_election_events(discovery, client, resp.events(), election_key).await;
                    }
                    Ok(None) => {
                        warn!("Election watch stream closed");
                        return Err(ClusterError::Internal("Election stream EOF".into()));
                    }
                    Err(e) => {
                        warn!("Election watch error: {}", e);
                        return Err(ClusterError::Etcd(e));
                    }
                }
            }
        }
    }
}
