batata-client 0.0.2

use std::sync::atomic::{AtomicUsize, Ordering};

use crate::common::{parse_server_address, GRPC_PORT_OFFSET};
use crate::error::{BatataError, Result};

/// Server address information
#[derive(Clone, Debug)]
pub struct ServerAddress {
    pub host: String,
    pub port: u16,
    pub grpc_port: u16,
    pub tls_enabled: bool,
}

impl ServerAddress {
    pub fn new(host: &str, port: u16) -> Self {
        Self {
            host: host.to_string(),
            port,
            grpc_port: port + GRPC_PORT_OFFSET,
            tls_enabled: false,
        }
    }

    /// Enable TLS for this server address
    pub fn with_tls(mut self, enabled: bool) -> Self {
        self.tls_enabled = enabled;
        self
    }

    /// Get the host
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Get the HTTP port
    pub fn port(&self) -> u16 {
        self.port
    }

    /// Get the address string (host:port)
    pub fn address(&self) -> String {
        format!("{}:{}", self.host, self.port)
    }

    pub fn grpc_endpoint(&self) -> String {
        let scheme = if self.tls_enabled { "https" } else { "http" };
        format!("{}://{}:{}", scheme, self.host, self.grpc_port)
    }

    pub fn http_endpoint(&self) -> String {
        let scheme = if self.tls_enabled { "https" } else { "http" };
        format!("{}://{}:{}", scheme, self.host, self.port)
    }
}

/// Server list manager for load balancing
pub struct ServerListManager {
    servers: Vec<ServerAddress>,
    current_index: AtomicUsize,
}

impl ServerListManager {
    /// Create a new server list manager from address strings
    pub fn new(addresses: Vec<String>) -> Result<Self> {
        if addresses.is_empty() {
            return Err(BatataError::InvalidParameter(
                "Server addresses cannot be empty".to_string(),
            ));
        }

        let servers: Vec<ServerAddress> = addresses
            .iter()
            .map(|addr| {
                let (host, port) = parse_server_address(addr);
                ServerAddress::new(&host, port)
            })
            .collect();

        Ok(Self {
            servers,
            current_index: AtomicUsize::new(0),
        })
    }

    /// Get the next server using round-robin
    pub fn next_server(&self) -> &ServerAddress {
        let index = self.current_index.fetch_add(1, Ordering::Relaxed) % self.servers.len();
        &self.servers[index]
    }

    /// Get current server
    pub fn current_server(&self) -> &ServerAddress {
        let index = self.current_index.load(Ordering::Relaxed) % self.servers.len();
        &self.servers[index]
    }

    /// Get all servers
    pub fn all_servers(&self) -> &[ServerAddress] {
        &self.servers
    }

    /// Get server count
    pub fn server_count(&self) -> usize {
        self.servers.len()
    }

    /// Mark current server as failed and move to next
    pub fn fail_current(&self) -> &ServerAddress {
        self.next_server()
    }
}

impl Clone for ServerListManager {
    fn clone(&self) -> Self {
        Self {
            servers: self.servers.clone(),
            current_index: AtomicUsize::new(self.current_index.load(Ordering::Relaxed)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_server_address() {
        let addr = ServerAddress::new("localhost", 8848);
        assert_eq!(addr.grpc_port, 9848);
        assert_eq!(addr.grpc_endpoint(), "http://localhost:9848");
        assert_eq!(addr.http_endpoint(), "http://localhost:8848");
    }

    #[test]
    fn test_server_list_manager() {
        let manager = ServerListManager::new(vec![
            "localhost:8848".to_string(),
            "localhost:8849".to_string(),
        ])
        .unwrap();

        assert_eq!(manager.server_count(), 2);

        let first = manager.next_server().host.clone();
        let second = manager.next_server().host.clone();
        let third = manager.next_server().host.clone();

        assert_eq!(first, "localhost");
        assert_eq!(second, "localhost");
        assert_eq!(third, "localhost");
    }
}