// src/ip_tracker.rs
// Модуль для отслеживания и ограничения уникальных IP-адресов пользователей

use std::collections::{HashMap, HashSet};
use std::net::IpAddr;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Трекер уникальных IP-адресов для каждого пользователя MTProxy
/// 
/// Предоставляет thread-safe механизм для:
/// - Отслеживания активных IP-адресов каждого пользователя
/// - Ограничения количества уникальных IP на пользователя
/// - Автоматической очистки при отключении клиентов
#[derive(Debug, Clone)]
pub struct UserIpTracker {
    /// Маппинг: Имя пользователя -> Множество активных IP-адресов
    active_ips: Arc<RwLock<HashMap<String, HashSet<IpAddr>>>>,
    
    /// Маппинг: Имя пользователя -> Максимально разрешенное количество уникальных IP
    max_ips: Arc<RwLock<HashMap<String, usize>>>,
}

impl UserIpTracker {
    /// Создать новый пустой трекер
    pub fn new() -> Self {
        Self {
            active_ips: Arc::new(RwLock::new(HashMap::new())),
            max_ips: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Установить лимит уникальных IP для конкретного пользователя
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// * `max_ips` - Максимальное количество одновременно активных IP-адресов
    pub async fn set_user_limit(&self, username: &str, max_ips: usize) {
        let mut limits = self.max_ips.write().await;
        limits.insert(username.to_string(), max_ips);
    }

    /// Загрузить лимиты из конфигурации
    /// 
    /// # Arguments
    /// * `limits` - HashMap с лимитами из config.toml
    pub async fn load_limits(&self, limits: &HashMap<String, usize>) {
        let mut max_ips = self.max_ips.write().await;
        for (user, limit) in limits {
            max_ips.insert(user.clone(), *limit);
        }
    }

    /// Проверить, может ли пользователь подключиться с данного IP-адреса
    /// и добавить IP в список активных, если проверка успешна
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// * `ip` - IP-адрес клиента
    /// 
    /// # Returns
    /// * `Ok(())` - Подключение разрешено, IP добавлен в активные
    /// * `Err(String)` - Подключение отклонено с описанием причины
    pub async fn check_and_add(&self, username: &str, ip: IpAddr) -> Result<(), String> {
        // Получаем лимит для пользователя
        let max_ips = self.max_ips.read().await;
        let limit = match max_ips.get(username) {
            Some(limit) => *limit,
            None => {
                // Если лимит не задан - разрешаем безлимитный доступ
                drop(max_ips);
                let mut active_ips = self.active_ips.write().await;
                let user_ips = active_ips
                    .entry(username.to_string())
                    .or_insert_with(HashSet::new);
                user_ips.insert(ip);
                return Ok(());
            }
        };
        drop(max_ips);

        // Проверяем и обновляем активные IP
        let mut active_ips = self.active_ips.write().await;
        let user_ips = active_ips
            .entry(username.to_string())
            .or_insert_with(HashSet::new);

        // Если IP уже есть в списке - это повторное подключение, разрешаем
        if user_ips.contains(&ip) {
            return Ok(());
        }

        // Проверяем, не превышен ли лимит
        if user_ips.len() >= limit {
            return Err(format!(
                "IP limit reached for user '{}': {}/{} unique IPs already connected",
                username,
                user_ips.len(),
                limit
            ));
        }

        // Лимит не превышен - добавляем новый IP
        user_ips.insert(ip);
        Ok(())
    }

    /// Удалить IP-адрес из списка активных при отключении клиента
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// * `ip` - IP-адрес отключившегося клиента
    pub async fn remove_ip(&self, username: &str, ip: IpAddr) {
        let mut active_ips = self.active_ips.write().await;
        
        if let Some(user_ips) = active_ips.get_mut(username) {
            user_ips.remove(&ip);
            
            // Если у пользователя не осталось активных IP - удаляем запись
            // для экономии памяти
            if user_ips.is_empty() {
                active_ips.remove(username);
            }
        }
    }

    /// Получить текущее количество активных IP-адресов для пользователя
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// 
    /// # Returns
    /// Количество уникальных активных IP-адресов
    pub async fn get_active_ip_count(&self, username: &str) -> usize {
        let active_ips = self.active_ips.read().await;
        active_ips
            .get(username)
            .map(|ips| ips.len())
            .unwrap_or(0)
    }

    /// Получить список всех активных IP-адресов для пользователя
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// 
    /// # Returns
    /// Вектор с активными IP-адресами
    pub async fn get_active_ips(&self, username: &str) -> Vec<IpAddr> {
        let active_ips = self.active_ips.read().await;
        active_ips
            .get(username)
            .map(|ips| ips.iter().copied().collect())
            .unwrap_or_else(Vec::new)
    }

    /// Получить статистику по всем пользователям
    /// 
    /// # Returns
    /// Вектор кортежей: (имя_пользователя, количество_активных_IP, лимит)
    pub async fn get_stats(&self) -> Vec<(String, usize, usize)> {
        let active_ips = self.active_ips.read().await;
        let max_ips = self.max_ips.read().await;

        let mut stats = Vec::new();
        
        // Собираем статистику по пользователям с активными подключениями
        for (username, user_ips) in active_ips.iter() {
            let limit = max_ips.get(username).copied().unwrap_or(0);
            stats.push((username.clone(), user_ips.len(), limit));
        }
        
        stats.sort_by(|a, b| a.0.cmp(&b.0)); // Сортируем по имени пользователя
        stats
    }

    /// Очистить все активные IP для пользователя (при необходимости)
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    pub async fn clear_user_ips(&self, username: &str) {
        let mut active_ips = self.active_ips.write().await;
        active_ips.remove(username);
    }

    /// Очистить всю статистику (использовать с осторожностью!)
    pub async fn clear_all(&self) {
        let mut active_ips = self.active_ips.write().await;
        active_ips.clear();
    }

    /// Проверить, подключен ли пользователь с данного IP
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// * `ip` - IP-адрес для проверки
    /// 
    /// # Returns
    /// `true` если IP активен, `false` если нет
    pub async fn is_ip_active(&self, username: &str, ip: IpAddr) -> bool {
        let active_ips = self.active_ips.read().await;
        active_ips
            .get(username)
            .map(|ips| ips.contains(&ip))
            .unwrap_or(false)
    }

    /// Получить лимит для пользователя
    /// 
    /// # Arguments
    /// * `username` - Имя пользователя
    /// 
    /// # Returns
    /// Лимит IP-адресов или None, если лимит не установлен
    pub async fn get_user_limit(&self, username: &str) -> Option<usize> {
        let max_ips = self.max_ips.read().await;
        max_ips.get(username).copied()
    }

    /// Форматировать статистику в читаемый текст
    /// 
    /// # Returns
    /// Строка со статистикой для логов или мониторинга
    pub async fn format_stats(&self) -> String {
        let stats = self.get_stats().await;
        
        if stats.is_empty() {
            return String::from("No active users");
        }
        
        let mut output = String::from("User IP Statistics:\n");
        output.push_str("==================\n");
        
        for (username, active_count, limit) in stats {
            output.push_str(&format!(
                "User: {:<20} Active IPs: {}/{}\n",
                username,
                active_count,
                if limit > 0 { limit.to_string() } else { "unlimited".to_string() }
            ));
            
            let ips = self.get_active_ips(&username).await;
            for ip in ips {
                output.push_str(&format!("  └─ {}\n", ip));
            }
        }
        
        output
    }
}

impl Default for UserIpTracker {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// ТЕСТЫ
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

    fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr {
        IpAddr::V4(Ipv4Addr::new(oct1, oct2, oct3, oct4))
    }

    fn test_ipv6() -> IpAddr {
        IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1))
    }

    #[tokio::test]
    async fn test_basic_ip_limit() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("test_user", 2).await;

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);
        let ip3 = test_ipv4(192, 168, 1, 3);

        // Первые два IP должны быть приняты
        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
        assert!(tracker.check_and_add("test_user", ip2).await.is_ok());

        // Третий IP должен быть отклонен
        assert!(tracker.check_and_add("test_user", ip3).await.is_err());

        // Проверяем счетчик
        assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
    }

    #[tokio::test]
    async fn test_reconnection_from_same_ip() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("test_user", 2).await;

        let ip1 = test_ipv4(192, 168, 1, 1);

        // Первое подключение
        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
        
        // Повторное подключение с того же IP должно пройти
        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
        
        // Счетчик не должен увеличиться
        assert_eq!(tracker.get_active_ip_count("test_user").await, 1);
    }

    #[tokio::test]
    async fn test_ip_removal() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("test_user", 2).await;

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);
        let ip3 = test_ipv4(192, 168, 1, 3);

        // Добавляем два IP
        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
        assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
        
        // Третий не должен пройти
        assert!(tracker.check_and_add("test_user", ip3).await.is_err());

        // Удаляем первый IP
        tracker.remove_ip("test_user", ip1).await;
        
        // Теперь третий должен пройти
        assert!(tracker.check_and_add("test_user", ip3).await.is_ok());
        
        assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
    }

    #[tokio::test]
    async fn test_no_limit() {
        let tracker = UserIpTracker::new();
        // Не устанавливаем лимит для test_user

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);
        let ip3 = test_ipv4(192, 168, 1, 3);

        // Без лимита все IP должны проходить
        assert!(tracker.check_and_add("test_user", ip1).await.is_ok());
        assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
        assert!(tracker.check_and_add("test_user", ip3).await.is_ok());
        
        assert_eq!(tracker.get_active_ip_count("test_user").await, 3);
    }

    #[tokio::test]
    async fn test_multiple_users() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("user1", 2).await;
        tracker.set_user_limit("user2", 1).await;

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);

        // user1 может использовать 2 IP
        assert!(tracker.check_and_add("user1", ip1).await.is_ok());
        assert!(tracker.check_and_add("user1", ip2).await.is_ok());

        // user2 может использовать только 1 IP
        assert!(tracker.check_and_add("user2", ip1).await.is_ok());
        assert!(tracker.check_and_add("user2", ip2).await.is_err());
    }

    #[tokio::test]
    async fn test_ipv6_support() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("test_user", 2).await;

        let ipv4 = test_ipv4(192, 168, 1, 1);
        let ipv6 = test_ipv6();

        // Должны работать оба типа адресов
        assert!(tracker.check_and_add("test_user", ipv4).await.is_ok());
        assert!(tracker.check_and_add("test_user", ipv6).await.is_ok());
        
        assert_eq!(tracker.get_active_ip_count("test_user").await, 2);
    }

    #[tokio::test]
    async fn test_get_active_ips() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("test_user", 3).await;

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);

        tracker.check_and_add("test_user", ip1).await.unwrap();
        tracker.check_and_add("test_user", ip2).await.unwrap();

        let active_ips = tracker.get_active_ips("test_user").await;
        assert_eq!(active_ips.len(), 2);
        assert!(active_ips.contains(&ip1));
        assert!(active_ips.contains(&ip2));
    }

    #[tokio::test]
    async fn test_stats() {
        let tracker = UserIpTracker::new();
        tracker.set_user_limit("user1", 3).await;
        tracker.set_user_limit("user2", 2).await;

        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);

        tracker.check_and_add("user1", ip1).await.unwrap();
        tracker.check_and_add("user2", ip2).await.unwrap();

        let stats = tracker.get_stats().await;
        assert_eq!(stats.len(), 2);
        
        // Проверяем наличие обоих пользователей в статистике
        assert!(stats.iter().any(|(name, _, _)| name == "user1"));
        assert!(stats.iter().any(|(name, _, _)| name == "user2"));
    }

    #[tokio::test]
    async fn test_clear_user_ips() {
        let tracker = UserIpTracker::new();
        let ip1 = test_ipv4(192, 168, 1, 1);
        
        tracker.check_and_add("test_user", ip1).await.unwrap();
        assert_eq!(tracker.get_active_ip_count("test_user").await, 1);
        
        tracker.clear_user_ips("test_user").await;
        assert_eq!(tracker.get_active_ip_count("test_user").await, 0);
    }

    #[tokio::test]
    async fn test_is_ip_active() {
        let tracker = UserIpTracker::new();
        let ip1 = test_ipv4(192, 168, 1, 1);
        let ip2 = test_ipv4(192, 168, 1, 2);
        
        tracker.check_and_add("test_user", ip1).await.unwrap();
        
        assert!(tracker.is_ip_active("test_user", ip1).await);
        assert!(!tracker.is_ip_active("test_user", ip2).await);
    }

    #[tokio::test]
    async fn test_load_limits_from_config() {
        let tracker = UserIpTracker::new();
        
        let mut config_limits = HashMap::new();
        config_limits.insert("user1".to_string(), 5);
        config_limits.insert("user2".to_string(), 3);
        
        tracker.load_limits(&config_limits).await;
        
        assert_eq!(tracker.get_user_limit("user1").await, Some(5));
        assert_eq!(tracker.get_user_limit("user2").await, Some(3));
        assert_eq!(tracker.get_user_limit("user3").await, None);
    }
}