//! Load balancing algorithms for upstream selection.
//!
//! Uses Pingora's built-in load balancing primitives where possible,
//! with custom extensions for session-affinity consistent hashing.

use crate::config::Endpoint;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Load balancer that selects an upstream endpoint from a pool.
#[derive(Debug)]
pub struct LoadBalancer {
    endpoints: Vec<Endpoint>,
    counter: AtomicUsize,
}

impl LoadBalancer {
    /// Create a new load balancer with the given endpoints.
    pub fn new(endpoints: Vec<Endpoint>) -> Self {
        Self {
            endpoints,
            counter: AtomicUsize::new(0),
        }
    }

    /// Update the endpoint pool (e.g., on hot-reload).
    pub fn update_endpoints(&mut self, endpoints: Vec<Endpoint>) {
        self.endpoints = endpoints;
    }

    /// Select an endpoint using round-robin.
    pub fn round_robin(&self) -> Option<&Endpoint> {
        let healthy: Vec<&Endpoint> = self.endpoints.iter().filter(|e| e.ready).collect();
        if healthy.is_empty() {
            return None;
        }
        let idx = self.counter.fetch_add(1, Ordering::Relaxed) % healthy.len();
        Some(healthy[idx])
    }

    /// Select an endpoint using the least-connections strategy (placeholder).
    ///
    /// Full implementation would track active connections per endpoint.
    pub fn least_connections(&self) -> Option<&Endpoint> {
        self.endpoints.iter().filter(|e| e.ready).min_by_key(|_| {
            // Placeholder: return 0 for now. Real impl would track connection counts.
            0usize
        })
    }

    /// Consistent hash selection based on a key (for session affinity).
    pub fn consistent_hash(&self, key: &str) -> Option<&Endpoint> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let healthy: Vec<&Endpoint> = self.endpoints.iter().filter(|e| e.ready).collect();
        if healthy.is_empty() {
            return None;
        }

        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        let hash = hasher.finish();
        let idx = hash as usize % healthy.len();
        Some(healthy[idx])
    }
}