"""Tests for the resource cleanup manager."""

import pytest
import tempfile
import time
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime, timezone, timedelta

from src.services.resource_cleanup_manager import ResourceCleanupManager, CleanupStats


class TestResourceCleanupManager:
    """Test cases for ResourceCleanupManager."""
    
    @pytest.fixture
    def cleanup_manager(self):
        """Create a ResourceCleanupManager instance for testing."""
        return ResourceCleanupManager(cleanup_interval_seconds=1, max_temp_files=10)
    
    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory for testing."""
        with tempfile.TemporaryDirectory() as temp_dir:
            yield Path(temp_dir)
    
    def test_initialization(self, cleanup_manager):
        """Test cleanup manager initialization."""
        assert cleanup_manager.cleanup_interval_seconds == 1
        assert cleanup_manager.max_temp_files == 10
        assert isinstance(cleanup_manager.temp_directories, set)
        assert isinstance(cleanup_manager.cache_directories, set)
        assert isinstance(cleanup_manager.model_references, set)
        assert cleanup_manager._cleanup_thread.is_alive()
    
    def test_register_temp_directory(self, cleanup_manager, temp_dir):
        """Test registering temporary directory."""
        cleanup_manager.register_temp_directory(temp_dir)
        
        assert temp_dir in cleanup_manager.temp_directories
    
    def test_register_cache_directory(self, cleanup_manager, temp_dir):
        """Test registering cache directory."""
        cleanup_manager.register_cache_directory(temp_dir)
        
        assert temp_dir in cleanup_manager.cache_directories
    
    def test_register_model_reference(self, cleanup_manager):
        """Test registering model reference."""
        mock_model = Mock()
        cleanup_manager.register_model_reference(mock_model)
        
        assert mock_model in cleanup_manager.model_references
    
    @patch('src.services.resource_cleanup_manager.psutil.virtual_memory')
    @patch('src.services.resource_cleanup_manager.torch.cuda.is_available')
    @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated')
    @patch('src.services.resource_cleanup_manager.torch.cuda.empty_cache')
    @patch('src.services.resource_cleanup_manager.gc.collect')
    def test_cleanup_memory(self, mock_gc_collect, mock_empty_cache, mock_memory_allocated, 
                           mock_cuda_available, mock_virtual_memory, cleanup_manager):
        """Test memory cleanup."""
        # Mock memory usage
        mock_memory = Mock()
        mock_memory.used = 8 * 1024 * 1024  # 8MB initial
        mock_virtual_memory.return_value = mock_memory
        
        mock_cuda_available.return_value = True
        mock_memory_allocated.return_value = 4 * 1024 * 1024  # 4MB GPU memory
        mock_gc_collect.return_value = 10  # 10 objects collected
        
        # Add some model references
        mock_model1 = Mock()
        mock_model2 = Mock()
        cleanup_manager.register_model_reference(mock_model1)
        cleanup_manager.register_model_reference(mock_model2)
        
        # Perform cleanup
        result = cleanup_manager.cleanup_memory()
        
        # Verify cleanup was performed
        assert "memory_freed_mb" in result
        assert "gpu_memory_freed_mb" in result
        assert "objects_collected" in result
        assert "cleanup_duration_seconds" in result
        
        # Verify model references were cleared
        assert len(cleanup_manager.model_references) == 0
        
        # Verify garbage collection was called
        mock_gc_collect.assert_called_once()
        
        # Verify GPU cache was cleared
        mock_empty_cache.assert_called_once()
    
    def test_cleanup_temp_files(self, cleanup_manager, temp_dir):
        """Test temporary file cleanup."""
        # Create some test files
        test_file1 = temp_dir / "test1.txt"
        test_file2 = temp_dir / "test2.txt"
        test_file3 = temp_dir / "test3.txt"
        
        test_file1.write_text("test content 1")
        test_file2.write_text("test content 2")
        test_file3.write_text("test content 3")
        
        # Set file modification times
        old_time = datetime.now(timezone.utc) - timedelta(hours=2)  # 2 hours ago
        new_time = datetime.now(timezone.utc) - timedelta(minutes=30)  # 30 minutes ago
        
        # Make two files old
        test_file1.touch()
        test_file2.touch()
        test_file3.touch()
        
        # Register temp directory
        cleanup_manager.register_temp_directory(temp_dir)
        
        # Perform cleanup
        result = cleanup_manager.cleanup_temp_files()
        
        # Verify cleanup statistics
        assert "files_removed" in result
        assert "directories_removed" in result
        assert "cleanup_duration_seconds" in result
        
        # Note: In real scenarios, files would be removed based on modification time
        # This test verifies the cleanup logic structure
    
    def test_cleanup_cache(self, cleanup_manager, temp_dir):
        """Test cache cleanup."""
        # Create some test cache files
        cache_file1 = temp_dir / "cache1.pkl"
        cache_file2 = temp_dir / "cache2.pkl"
        
        cache_file1.write_text("cache content 1")
        cache_file2.write_text("cache content 2")
        
        # Register cache directory
        cleanup_manager.register_cache_directory(temp_dir)
        
        # Perform cleanup
        result = cleanup_manager.cleanup_cache()
        
        # Verify cleanup statistics
        assert "cache_entries_cleared" in result
        assert "cleanup_duration_seconds" in result
    
    @patch('src.services.resource_cleanup_manager.psutil.virtual_memory')
    @patch('src.services.resource_cleanup_manager.torch.cuda.is_available')
    @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated')
    @patch('src.services.resource_cleanup_manager.torch.cuda.empty_cache')
    @patch('src.services.resource_cleanup_manager.gc.collect')
    def test_perform_full_cleanup(self, mock_gc_collect, mock_empty_cache, mock_memory_allocated,
                                 mock_cuda_available, mock_virtual_memory, cleanup_manager, temp_dir):
        """Test full cleanup operation."""
        # Mock memory usage
        mock_memory = Mock()
        mock_memory.used = 8 * 1024 * 1024  # 8MB
        mock_virtual_memory.return_value = mock_memory
        
        mock_cuda_available.return_value = True
        mock_memory_allocated.return_value = 4 * 1024 * 1024  # 4MB GPU memory
        mock_gc_collect.return_value = 5  # 5 objects collected
        
        # Register directories and models
        cleanup_manager.register_temp_directory(temp_dir)
        cleanup_manager.register_cache_directory(temp_dir)
        cleanup_manager.register_model_reference(Mock())
        
        # Perform full cleanup
        stats = cleanup_manager.perform_full_cleanup()
        
        # Verify cleanup stats
        assert isinstance(stats, CleanupStats)
        assert stats.memory_freed_mb >= 0
        assert stats.gpu_memory_freed_mb >= 0
        assert stats.temp_files_removed >= 0
        assert stats.cache_entries_cleared >= 0
        assert stats.cleanup_duration_seconds >= 0
        assert isinstance(stats.timestamp, datetime)
        
        # Verify stats were added to history
        assert len(cleanup_manager.cleanup_stats) > 0
        assert cleanup_manager.cleanup_stats[-1] == stats
    
    def test_get_cleanup_stats(self, cleanup_manager):
        """Test getting cleanup statistics."""
        # Add some mock stats
        mock_stats = [
            CleanupStats(
                memory_freed_mb=100.0,
                gpu_memory_freed_mb=50.0,
                temp_files_removed=5,
                cache_entries_cleared=10,
                cleanup_duration_seconds=1.5,
                timestamp=datetime.now(timezone.utc)
            ),
            CleanupStats(
                memory_freed_mb=200.0,
                gpu_memory_freed_mb=75.0,
                temp_files_removed=8,
                cache_entries_cleared=15,
                cleanup_duration_seconds=2.0,
                timestamp=datetime.now(timezone.utc)
            )
        ]
        
        cleanup_manager.cleanup_stats = mock_stats
        
        # Get stats with limit
        stats = cleanup_manager.get_cleanup_stats(limit=1)
        
        assert len(stats) == 1
        assert stats[0] == mock_stats[-1]
    
    @patch('src.services.resource_cleanup_manager.psutil.virtual_memory')
    @patch('src.services.resource_cleanup_manager.torch.cuda.is_available')
    @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated')
    @patch('src.services.resource_cleanup_manager.torch.cuda.get_device_properties')
    def test_get_memory_usage(self, mock_gpu_props, mock_memory_allocated, mock_cuda_available,
                             mock_virtual_memory, cleanup_manager):
        """Test memory usage monitoring."""
        # Mock system memory
        mock_memory = Mock()
        mock_memory.total = 16 * 1024**3  # 16GB
        mock_memory.available = 8 * 1024**3  # 8GB
        mock_memory.used = 8 * 1024**3  # 8GB
        mock_memory.percent = 50.0
        mock_virtual_memory.return_value = mock_memory
        
        # Mock GPU memory
        mock_cuda_available.return_value = True
        mock_memory_allocated.return_value = 4 * 1024**3  # 4GB allocated
        mock_gpu_props.return_value.total_memory = 8 * 1024**3  # 8GB total
        
        # Get memory usage
        usage = cleanup_manager.get_memory_usage()
        
        # Verify system memory
        assert usage["total_memory_gb"] == 16.0
        assert usage["available_memory_gb"] == 8.0
        assert usage["used_memory_gb"] == 8.0
        assert usage["memory_percent"] == 50.0
        
        # Verify GPU memory
        assert usage["gpu_memory_allocated_gb"] == 4.0
        assert usage["gpu_memory_total_gb"] == 8.0
        assert usage["gpu_memory_percent"] == 50.0
    
    def test_should_perform_cleanup_high_memory(self, cleanup_manager):
        """Test cleanup decision with high memory usage."""
        # Mock high memory usage
        with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage:
            mock_usage.return_value = {
                "memory_percent": 85.0,  # High memory usage
                "gpu_memory_percent": 30.0
            }
            
            should_cleanup = cleanup_manager.should_perform_cleanup()
            
            assert should_cleanup is True
    
    def test_should_perform_cleanup_high_gpu_memory(self, cleanup_manager):
        """Test cleanup decision with high GPU memory usage."""
        # Mock high GPU memory usage
        with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage:
            mock_usage.return_value = {
                "memory_percent": 50.0,
                "gpu_memory_percent": 85.0  # High GPU memory usage
            }
            
            should_cleanup = cleanup_manager.should_perform_cleanup()
            
            assert should_cleanup is True
    
    def test_should_perform_cleanup_time_based(self, cleanup_manager):
        """Test cleanup decision based on time since last cleanup."""
        # Set last cleanup to be a long time ago
        cleanup_manager.last_cleanup = datetime.now(timezone.utc) - timedelta(hours=2)
        
        # Mock normal memory usage
        with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage:
            mock_usage.return_value = {
                "memory_percent": 50.0,
                "gpu_memory_percent": 30.0
            }
            
            should_cleanup = cleanup_manager.should_perform_cleanup()
            
            assert should_cleanup is True
    
    def test_should_perform_cleanup_normal_usage(self, cleanup_manager):
        """Test cleanup decision with normal memory usage."""
        # Mock normal memory usage
        with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage:
            mock_usage.return_value = {
                "memory_percent": 50.0,
                "gpu_memory_percent": 30.0
            }
            
            should_cleanup = cleanup_manager.should_perform_cleanup()
            
            # Should not cleanup unless it's been a long time
            assert should_cleanup is False
    
    def test_shutdown(self, cleanup_manager):
        """Test cleanup manager shutdown."""
        # Verify thread is running
        assert cleanup_manager._cleanup_thread.is_alive()
        
        # Shutdown
        cleanup_manager.shutdown()
        
        # Verify thread is stopped
        assert not cleanup_manager._cleanup_thread.is_alive()
    
    def test_cleanup_stats_dataclass(self):
        """Test CleanupStats dataclass."""
        timestamp = datetime.now(timezone.utc)
        
        stats = CleanupStats(
            memory_freed_mb=100.0,
            gpu_memory_freed_mb=50.0,
            temp_files_removed=5,
            cache_entries_cleared=10,
            cleanup_duration_seconds=1.5,
            timestamp=timestamp
        )
        
        assert stats.memory_freed_mb == 100.0
        assert stats.gpu_memory_freed_mb == 50.0
        assert stats.temp_files_removed == 5
        assert stats.cache_entries_cleared == 10
        assert stats.cleanup_duration_seconds == 1.5
        assert stats.timestamp == timestamp
