//! VIP Manager for health-based BGP route advertisement
//!
//! Monitors load balancer health status and automatically advertises/withdraws
//! VIP routes based on backend availability.

use std::collections::{HashMap, HashSet};
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::{watch, RwLock};
use tokio::time::sleep;
use tracing::{debug, error, info, warn};

use crate::bgp_client::BgpClient;
use crate::metadata::LbMetadataStore;
use crate::vip_owner::VipAddressOwner;
use fiberlb_types::LoadBalancerId;

const METRIC_VIP_DRAIN_ACTIVE: &str = "fiberlb_vip_drain_active";

/// Current local control-plane state for a VIP.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct VipState {
    /// The VIP is configured on the local node.
    owned: bool,
    /// The VIP is advertised to BGP peers.
    advertised: bool,
    /// When the node entered drain while this VIP was active.
    drain_started_at: Option<Instant>,
}

impl VipState {
    fn is_idle(self) -> bool {
        !self.owned && !self.advertised
    }
}

/// VIP Manager
///
/// Coordinates BGP route advertisements based on load balancer health.
/// Runs a background task that periodically checks health status and
/// updates BGP advertisements accordingly.
pub struct VipManager {
    /// BGP client for route operations
    bgp: Arc<dyn BgpClient>,
    /// Metadata store for load balancer state
    metadata: Arc<LbMetadataStore>,
    /// Current VIP advertisement state
    vip_state: Arc<RwLock<HashMap<IpAddr, VipState>>>,
    /// Optional local VIP owner.
    vip_owner: Option<Arc<dyn VipAddressOwner>>,
    /// Router's own IP address (used as BGP next hop)
    next_hop: IpAddr,
    /// Presence of this file activates drain mode.
    drain_file: String,
    /// How long to keep a locally-owned VIP after withdrawing it for drain.
    drain_hold_time: Duration,
    /// Shutdown signal for the background reconciliation task.
    shutdown: watch::Sender<bool>,
}

impl VipManager {
    /// Create a new VIP manager
    pub fn new(
        bgp: Arc<dyn BgpClient>,
        metadata: Arc<LbMetadataStore>,
        next_hop: IpAddr,
        vip_owner: Option<Arc<dyn VipAddressOwner>>,
        drain_file: String,
        drain_hold_time: Duration,
    ) -> Self {
        let (shutdown, _shutdown_rx) = watch::channel(false);
        Self {
            bgp,
            metadata,
            vip_state: Arc::new(RwLock::new(HashMap::new())),
            vip_owner,
            next_hop,
            drain_file,
            drain_hold_time,
            shutdown,
        }
    }

    /// Start the VIP management background task
    ///
    /// This spawns a tokio task that periodically:
    /// 1. Scans all load balancers with VIP addresses
    /// 2. Checks backend health status
    /// 3. Advertises VIPs when healthy backends exist
    /// 4. Withdraws VIPs when all backends are unhealthy
    pub fn spawn(self: Arc<Self>, check_interval: Duration) -> tokio::task::JoinHandle<()> {
        tokio::spawn(async move {
            info!(
                "VIP manager started (check interval: {}s)",
                check_interval.as_secs()
            );
            let mut shutdown_rx = self.shutdown.subscribe();

            loop {
                if *shutdown_rx.borrow() {
                    info!("VIP manager reconciliation loop stopping");
                    break;
                }

                if let Err(e) = self.check_and_update_vips().await {
                    error!("VIP manager check failed: {}", e);
                }

                tokio::select! {
                    _ = sleep(check_interval) => {}
                    changed = shutdown_rx.changed() => {
                        if changed.is_err() || *shutdown_rx.borrow_and_update() {
                            info!("VIP manager reconciliation loop stopping");
                            break;
                        }
                    }
                }
            }
        })
    }

    /// Perform one iteration of VIP health check and update
    async fn check_and_update_vips(&self) -> Result<(), Box<dyn std::error::Error>> {
        debug!("Checking VIP health status...");

        // Get all load balancers
        let load_balancers = self.metadata.list_all_lbs().await?;

        // Track which VIPs should be advertised
        let mut active_vips = HashSet::new();

        for lb in &load_balancers {
            // Skip LBs without VIP address
            let Some(ref vip_str) = lb.vip_address else {
                continue;
            };

            let vip: IpAddr = match vip_str.parse() {
                Ok(addr) => addr,
                Err(e) => {
                    warn!("Invalid VIP address '{}' for LB {}: {}", vip_str, lb.id, e);
                    continue;
                }
            };

            // Check if this LB has healthy backends
            let has_healthy_backends = self.has_healthy_backends(&lb.id).await?;

            if has_healthy_backends {
                active_vips.insert(vip);
            }
        }

        let drain_active = self.is_drain_active().await;
        metrics::gauge!(METRIC_VIP_DRAIN_ACTIVE).set(if drain_active { 1.0 } else { 0.0 });

        // Update BGP advertisements
        self.reconcile_advertisements(&active_vips, drain_active).await?;

        Ok(())
    }

    async fn is_drain_active(&self) -> bool {
        let path = self.drain_file.trim();
        if path.is_empty() {
            return false;
        }

        tokio::fs::metadata(path).await.is_ok()
    }

    /// Check if a load balancer has any healthy backends
    async fn has_healthy_backends(
        &self,
        lb_id: &LoadBalancerId,
    ) -> Result<bool, Box<dyn std::error::Error>> {
        // Get all pools for this load balancer
        let pools = self.metadata.list_pools(lb_id).await?;

        for pool in pools {
            // Get backends for this pool
            let backends = self.metadata.list_backends(&pool.id).await?;

            // Check if any backend is healthy
            for backend in backends {
                if backend.is_available() {
                    return Ok(true);
                }
            }
        }

        Ok(false)
    }

    /// Reconcile BGP advertisements with desired state
    ///
    /// Compares current advertisements with desired active VIPs and:
    /// - Announces new VIPs that should be active
    /// - Withdraws VIPs that should no longer be active
    async fn reconcile_advertisements(
        &self,
        active_vips: &HashSet<IpAddr>,
        drain_active: bool,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let mut state = self.vip_state.write().await;
        let now = Instant::now();

        for vip in active_vips {
            let mut vip_state = state.get(vip).copied().unwrap_or_default();
            let mut changed = false;

            if drain_active {
                if vip_state.advertised {
                    info!("Withdrawing VIP {} for node drain", vip);
                    if let Err(error) = self.bgp.withdraw_route(*vip).await {
                        error!("Failed to withdraw VIP {} for drain: {}", vip, error);
                    } else {
                        vip_state.advertised = false;
                        vip_state.drain_started_at = Some(now);
                        changed = true;
                    }
                } else if vip_state.owned && vip_state.drain_started_at.is_none() {
                    vip_state.drain_started_at = Some(now);
                    changed = true;
                }

                if vip_state.owned {
                    let drain_started_at = vip_state.drain_started_at.unwrap_or(now);
                    if now.duration_since(drain_started_at) >= self.drain_hold_time {
                        if let Some(vip_owner) = &self.vip_owner {
                            info!("Releasing local VIP {} after drain hold", vip);
                            if let Err(error) = vip_owner.ensure_absent(*vip).await {
                                error!("Failed to release local VIP {} after drain: {}", vip, error);
                            } else {
                                vip_state.owned = false;
                                changed = true;
                            }
                        } else {
                            vip_state.owned = false;
                            changed = true;
                        }
                    }
                }

                if vip_state.is_idle() {
                    state.remove(vip);
                } else if changed {
                    state.insert(*vip, vip_state);
                }
                continue;
            }

            vip_state.drain_started_at = None;

            if !vip_state.owned {
                if let Some(vip_owner) = &self.vip_owner {
                    info!("Claiming local VIP {} on this node", vip);
                    if let Err(error) = vip_owner.ensure_present(*vip).await {
                        error!("Failed to claim local VIP {}: {}", vip, error);
                        continue;
                    }
                }
                vip_state.owned = true;
                changed = true;
            } else {
                debug!("VIP {} already locally owned", vip);
            }

            if !vip_state.advertised {
                info!("Announcing VIP {} (healthy backends available)", vip);
                if let Err(error) = self.bgp.announce_route(*vip, self.next_hop).await {
                    error!("Failed to announce VIP {}: {}", vip, error);
                } else {
                    vip_state.advertised = true;
                    changed = true;
                }
            } else {
                debug!("VIP {} already advertised", vip);
            }

            if changed {
                state.insert(*vip, vip_state);
            }
        }

        let managed_vips: Vec<IpAddr> = state.keys().copied().collect();
        for vip in managed_vips {
            if !active_vips.contains(&vip) {
                let mut vip_state = state.get(&vip).copied().unwrap_or_default();
                vip_state.drain_started_at = None;

                if vip_state.owned {
                    if let Some(vip_owner) = &self.vip_owner {
                        info!("Releasing local VIP {} (no healthy backends)", vip);
                        if let Err(error) = vip_owner.ensure_absent(vip).await {
                            error!("Failed to release local VIP {}: {}", vip, error);
                        } else {
                            vip_state.owned = false;
                        }
                    } else {
                        vip_state.owned = false;
                    }
                }

                if vip_state.advertised {
                    info!("Withdrawing VIP {} (no healthy backends)", vip);
                    if let Err(error) = self.bgp.withdraw_route(vip).await {
                        error!("Failed to withdraw VIP {}: {}", vip, error);
                    } else {
                        vip_state.advertised = false;
                    }
                }

                if vip_state.is_idle() {
                    state.remove(&vip);
                } else {
                    state.insert(vip, vip_state);
                }
            }
        }

        Ok(())
    }

    /// Manually advertise a VIP (for testing or manual control)
    pub async fn advertise_vip(&self, vip: IpAddr) -> Result<(), Box<dyn std::error::Error>> {
        let mut state = self.vip_state.write().await;
        let mut vip_state = state.get(&vip).copied().unwrap_or_default();

        if !vip_state.owned {
            if let Some(vip_owner) = &self.vip_owner {
                info!("Claiming local VIP {} on this node", vip);
                vip_owner.ensure_present(vip).await?;
            }
            vip_state.owned = true;
        }
        vip_state.drain_started_at = None;

        if !vip_state.advertised {
            info!("Manually advertising VIP {}", vip);
            self.bgp.announce_route(vip, self.next_hop).await?;
            vip_state.advertised = true;
        }

        state.insert(vip, vip_state);

        Ok(())
    }

    /// Manually withdraw a VIP (for testing or manual control)
    pub async fn withdraw_vip(&self, vip: IpAddr) -> Result<(), Box<dyn std::error::Error>> {
        let mut state = self.vip_state.write().await;
        let mut vip_state = state.get(&vip).copied().unwrap_or_default();

        if vip_state.owned {
            if let Some(vip_owner) = &self.vip_owner {
                info!("Releasing local VIP {} on this node", vip);
                vip_owner.ensure_absent(vip).await?;
            }
            vip_state.owned = false;
        }
        vip_state.drain_started_at = None;

        if vip_state.advertised {
            info!("Manually withdrawing VIP {}", vip);
            self.bgp.withdraw_route(vip).await?;
            vip_state.advertised = false;
        }

        if vip_state.is_idle() {
            state.remove(&vip);
        } else {
            state.insert(vip, vip_state);
        }

        Ok(())
    }

    /// Gracefully shutdown - withdraw all advertised VIPs
    ///
    /// Should be called during server shutdown to ensure clean BGP state
    pub async fn shutdown(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("VIP manager shutting down, withdrawing all VIPs...");
        let _ = self.shutdown.send(true);

        let mut state = self.vip_state.write().await;
        let managed_vips: Vec<IpAddr> = state.keys().copied().collect();

        for vip in managed_vips {
            let mut vip_state = state.get(&vip).copied().unwrap_or_default();
            vip_state.drain_started_at = None;

            if vip_state.owned {
                info!("Releasing local VIP {} for shutdown", vip);
                if let Some(vip_owner) = &self.vip_owner {
                    if let Err(error) = vip_owner.ensure_absent(vip).await {
                        error!(
                            "Failed to release local VIP {} during shutdown: {}",
                            vip, error
                        );
                    } else {
                        vip_state.owned = false;
                    }
                } else {
                    vip_state.owned = false;
                }
            }

            if vip_state.advertised {
                info!("Withdrawing VIP {} for shutdown", vip);
                if let Err(error) = self.bgp.withdraw_route(vip).await {
                    error!("Failed to withdraw VIP {} during shutdown: {}", vip, error);
                } else {
                    vip_state.advertised = false;
                }
            }

            if vip_state.is_idle() {
                state.remove(&vip);
            } else {
                state.insert(vip, vip_state);
            }
        }

        info!("VIP manager shutdown complete");
        Ok(())
    }

    /// Get current advertisement state (for monitoring/debugging)
    pub async fn get_advertised_vips(&self) -> Vec<IpAddr> {
        let state = self.vip_state.read().await;
        state
            .iter()
            .filter(|(_, state)| state.advertised)
            .map(|(vip, _)| *vip)
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bgp_client::{BgpClient, Result};
    use crate::vip_owner::VipOwnershipError;
    use std::sync::Mutex;
    use std::time::{Duration, SystemTime, UNIX_EPOCH};

    /// Mock BGP client for testing
    struct MockBgpClient {
        announced: Arc<Mutex<HashSet<IpAddr>>>,
        withdrawn: Arc<Mutex<HashSet<IpAddr>>>,
        events: Arc<Mutex<Vec<String>>>,
    }

    impl MockBgpClient {
        fn new(events: Arc<Mutex<Vec<String>>>) -> Self {
            Self {
                announced: Arc::new(Mutex::new(HashSet::new())),
                withdrawn: Arc::new(Mutex::new(HashSet::new())),
                events,
            }
        }
    }

    #[tonic::async_trait]
    impl BgpClient for MockBgpClient {
        async fn announce_route(&self, prefix: IpAddr, _next_hop: IpAddr) -> Result<()> {
            self.events
                .lock()
                .unwrap()
                .push(format!("announce:{prefix}"));
            self.announced.lock().unwrap().insert(prefix);
            Ok(())
        }

        async fn withdraw_route(&self, prefix: IpAddr) -> Result<()> {
            self.events
                .lock()
                .unwrap()
                .push(format!("withdraw:{prefix}"));
            self.withdrawn.lock().unwrap().insert(prefix);
            Ok(())
        }

        async fn is_connected(&self) -> bool {
            true
        }
    }

    struct MockVipOwner {
        owned: Arc<Mutex<HashSet<IpAddr>>>,
        released: Arc<Mutex<HashSet<IpAddr>>>,
        events: Arc<Mutex<Vec<String>>>,
    }

    impl MockVipOwner {
        fn new(events: Arc<Mutex<Vec<String>>>) -> Self {
            Self {
                owned: Arc::new(Mutex::new(HashSet::new())),
                released: Arc::new(Mutex::new(HashSet::new())),
                events,
            }
        }
    }

    #[tonic::async_trait]
    impl VipAddressOwner for MockVipOwner {
        async fn ensure_present(&self, vip: IpAddr) -> std::result::Result<(), VipOwnershipError> {
            self.events.lock().unwrap().push(format!("own:{vip}"));
            self.owned.lock().unwrap().insert(vip);
            Ok(())
        }

        async fn ensure_absent(&self, vip: IpAddr) -> std::result::Result<(), VipOwnershipError> {
            self.events.lock().unwrap().push(format!("unown:{vip}"));
            self.released.lock().unwrap().insert(vip);
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_vip_advertisement_tracking() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let mock_bgp = Arc::new(MockBgpClient::new(events.clone()));
        let mock_owner = Arc::new(MockVipOwner::new(events.clone()));
        let metadata = Arc::new(LbMetadataStore::new_in_memory());
        let next_hop = "10.0.0.1".parse().unwrap();

        let manager = VipManager::new(
            mock_bgp.clone(),
            metadata,
            next_hop,
            Some(mock_owner.clone()),
            String::new(),
            Duration::from_secs(0),
        );

        let vip: IpAddr = "10.0.1.100".parse().unwrap();

        // Advertise VIP
        manager.advertise_vip(vip).await.unwrap();
        assert_eq!(manager.get_advertised_vips().await, vec![vip]);
        assert!(mock_bgp.announced.lock().unwrap().contains(&vip));

        // Withdraw VIP
        manager.withdraw_vip(vip).await.unwrap();
        assert!(manager.get_advertised_vips().await.is_empty());
        assert!(mock_bgp.withdrawn.lock().unwrap().contains(&vip));
        assert!(mock_owner.released.lock().unwrap().contains(&vip));

        assert_eq!(
            events.lock().unwrap().clone(),
            vec![
                format!("own:{vip}"),
                format!("announce:{vip}"),
                format!("unown:{vip}"),
                format!("withdraw:{vip}"),
            ]
        );
    }

    #[tokio::test]
    async fn test_vip_drain_withdraws_before_releasing_vip() {
        let events = Arc::new(Mutex::new(Vec::new()));
        let mock_bgp = Arc::new(MockBgpClient::new(events.clone()));
        let mock_owner = Arc::new(MockVipOwner::new(events.clone()));
        let metadata = Arc::new(LbMetadataStore::new_in_memory());
        let next_hop = "10.0.0.1".parse().unwrap();
        let drain_path = std::env::temp_dir().join(format!(
            "fiberlb-drain-{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));

        let manager = VipManager::new(
            mock_bgp,
            metadata,
            next_hop,
            Some(mock_owner.clone()),
            drain_path.display().to_string(),
            Duration::from_millis(200),
        );

        let vip: IpAddr = "10.0.1.100".parse().unwrap();
        manager.advertise_vip(vip).await.unwrap();

        tokio::fs::write(&drain_path, b"1").await.unwrap();
        let mut active_vips = HashSet::new();
        active_vips.insert(vip);

        manager
            .reconcile_advertisements(&active_vips, true)
            .await
            .unwrap();
        assert!(mock_owner.released.lock().unwrap().is_empty());
        assert_eq!(manager.get_advertised_vips().await, Vec::<IpAddr>::new());

        sleep(Duration::from_millis(250)).await;
        manager
            .reconcile_advertisements(&active_vips, true)
            .await
            .unwrap();
        assert!(mock_owner.released.lock().unwrap().contains(&vip));

        tokio::fs::remove_file(&drain_path).await.unwrap();
        assert_eq!(
            events.lock().unwrap().clone(),
            vec![
                format!("own:{vip}"),
                format!("announce:{vip}"),
                format!("withdraw:{vip}"),
                format!("unown:{vip}"),
            ]
        );
    }
}
