#!/usr/bin/env python3
"""
Test suite for AI content generation functionality
"""

import os
import sys
from unittest.mock import patch, MagicMock

# Add parent directory to path to import modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from social_agent import AIContentGenerator


class TestAIContentGeneration:
    """Test class for AI content generation functionality"""
    
    def __init__(self):
        self.test_results = []
        
    def test_ai_generator_initialization(self):
        """Test AI content generator initialization"""
        try:
            ai_generator = AIContentGenerator()
            
            # Check if required attributes exist
            required_attrs = ['model', 'tokenizer', 'model_loaded', 'model_name']
            
            missing_attrs = []
            for attr in required_attrs:
                if not hasattr(ai_generator, attr):
                    missing_attrs.append(attr)
            
            if not missing_attrs:
                # Check initial state
                if (ai_generator.model is None and 
                    ai_generator.tokenizer is None and 
                    not ai_generator.model_loaded):
                    self.test_results.append(("AI Generator Initialization", True, "AIContentGenerator initialized correctly"))
                    return True
                else:
                    self.test_results.append(("AI Generator Initialization", False, f"Unexpected initial state - model: {ai_generator.model}, tokenizer: {ai_generator.tokenizer}, loaded: {ai_generator.model_loaded}"))
                    return False
            else:
                self.test_results.append(("AI Generator Initialization", False, f"Missing attributes: {missing_attrs}"))
                return False
                
        except Exception as e:
            self.test_results.append(("AI Generator Initialization", False, f"Exception: {e}"))
            return False
    
    def test_model_loading_simulation(self):
        """Test model loading functionality (simulated)"""
        try:
            ai_generator = AIContentGenerator()
            
            # Mock the model loading to avoid actually loading the large model
            with patch.object(ai_generator, 'load_model') as mock_load:
                mock_load.return_value = True
                
                result = ai_generator.load_model()
                
                if result:
                    self.test_results.append(("Model Loading Simulation", True, "Model loading method called successfully"))
                    return True
                else:
                    self.test_results.append(("Model Loading Simulation", False, "Model loading returned False"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Model Loading Simulation", False, f"Exception: {e}"))
            return False
    
    def test_alternative_post_generation_simulation(self):
        """Test alternative post generation (simulated)"""
        try:
            ai_generator = AIContentGenerator()
            
            # Mock the model loading and generation
            with patch.object(ai_generator, 'load_model') as mock_load, \
                 patch.object(ai_generator, 'model_loaded', True), \
                 patch.object(ai_generator, 'tokenizer') as mock_tokenizer, \
                 patch.object(ai_generator, 'model') as mock_model:
                
                mock_load.return_value = True
                mock_tokenizer.apply_chat_template.return_value = "test template"
                mock_tokenizer.return_tensors.return_value = {"input_ids": MagicMock()}
                mock_tokenizer.decode.return_value = "Generated alternative content"
                
                # Mock torch operations
                with patch('torch.no_grad'):
                    with patch('torch.cuda.is_available', return_value=False):
                        result = ai_generator.generate_alternative_post("Original content", "https://example.com")
                
                if result is not None:
                    self.test_results.append(("Alternative Post Generation Simulation", True, "Alternative post generation method works"))
                    return True
                else:
                    self.test_results.append(("Alternative Post Generation Simulation", False, "Alternative post generation returned None"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Alternative Post Generation Simulation", False, f"Exception: {e}"))
            return False
    
    def test_fact_generation_simulation(self):
        """Test fact generation (simulated)"""
        try:
            ai_generator = AIContentGenerator()
            
            # Mock the model loading and generation
            with patch.object(ai_generator, 'load_model') as mock_load, \
                 patch.object(ai_generator, 'model_loaded', True), \
                 patch.object(ai_generator, 'tokenizer') as mock_tokenizer, \
                 patch.object(ai_generator, 'model') as mock_model:
                
                mock_load.return_value = True
                mock_tokenizer.apply_chat_template.return_value = "test template"
                mock_tokenizer.return_tensors.return_value = {"input_ids": MagicMock()}
                mock_tokenizer.decode.return_value = "Generated fact content"
                
                # Mock torch operations
                with patch('torch.no_grad'):
                    with patch('torch.cuda.is_available', return_value=False):
                        result = ai_generator.generate_fact("https://example.com")
                
                if result is not None:
                    self.test_results.append(("Fact Generation Simulation", True, "Fact generation method works"))
                    return True
                else:
                    self.test_results.append(("Fact Generation Simulation", False, "Fact generation returned None"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Fact Generation Simulation", False, f"Exception: {e}"))
            return False
    
    def test_joke_generation_simulation(self):
        """Test joke generation (simulated)"""
        try:
            ai_generator = AIContentGenerator()
            
            # Mock the model loading and generation
            with patch.object(ai_generator, 'load_model') as mock_load, \
                 patch.object(ai_generator, 'model_loaded', True), \
                 patch.object(ai_generator, 'tokenizer') as mock_tokenizer, \
                 patch.object(ai_generator, 'model') as mock_model:
                
                mock_load.return_value = True
                mock_tokenizer.apply_chat_template.return_value = "test template"
                mock_tokenizer.return_tensors.return_value = {"input_ids": MagicMock()}
                mock_tokenizer.decode.return_value = "Generated joke content"
                
                # Mock torch operations
                with patch('torch.no_grad'):
                    with patch('torch.cuda.is_available', return_value=False):
                        result = ai_generator.generate_joke("https://example.com")
                
                if result is not None:
                    self.test_results.append(("Joke Generation Simulation", True, "Joke generation method works"))
                    return True
                else:
                    self.test_results.append(("Joke Generation Simulation", False, "Joke generation returned None"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Joke Generation Simulation", False, f"Exception: {e}"))
            return False
    
    def test_model_unloading(self):
        """Test model unloading functionality"""
        try:
            ai_generator = AIContentGenerator()
            
            # Set up mock loaded state
            ai_generator.model_loaded = True
            ai_generator.model = MagicMock()
            ai_generator.tokenizer = MagicMock()
            
            # Mock torch operations
            with patch('torch.cuda.is_available', return_value=False), \
                 patch('gc.collect') as mock_gc:
                
                ai_generator.unload_model()
                
                # Check if model was unloaded
                if (ai_generator.model is None and 
                    ai_generator.tokenizer is None and 
                    not ai_generator.model_loaded):
                    self.test_results.append(("Model Unloading", True, "Model unloaded successfully"))
                    return True
                else:
                    self.test_results.append(("Model Unloading", False, f"Model not properly unloaded - model: {ai_generator.model}, tokenizer: {ai_generator.tokenizer}, loaded: {ai_generator.model_loaded}"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Model Unloading", False, f"Exception: {e}"))
            return False
    
    def test_error_handling(self):
        """Test error handling in AI generation"""
        try:
            ai_generator = AIContentGenerator()
            
            # Test with model loading failure
            with patch.object(ai_generator, 'load_model') as mock_load:
                mock_load.return_value = False
                
                result = ai_generator.generate_alternative_post("test content")
                
                if result is None:
                    self.test_results.append(("Error Handling", True, "Properly handles model loading failure"))
                    return True
                else:
                    self.test_results.append(("Error Handling", False, "Should return None when model loading fails"))
                    return False
                    
        except Exception as e:
            self.test_results.append(("Error Handling", False, f"Exception: {e}"))
            return False
    
    def run_all_tests(self):
        """Run all AI content generation tests"""
        print("\n" + "="*60)
        print("🤖 RUNNING AI CONTENT GENERATION TESTS")
        print("="*60)
        
        # Run all tests
        tests = [
            self.test_ai_generator_initialization,
            self.test_model_loading_simulation,
            self.test_alternative_post_generation_simulation,
            self.test_fact_generation_simulation,
            self.test_joke_generation_simulation,
            self.test_model_unloading,
            self.test_error_handling
        ]
        
        passed = 0
        total = len(tests)
        
        for test in tests:
            print(f"\n🔍 Running {test.__name__}...")
            if test():
                passed += 1
                print(f"✅ {test.__name__} PASSED")
            else:
                print(f"❌ {test.__name__} FAILED")
        
        # Print summary
        print("\n" + "="*60)
        print("📊 AI CONTENT GENERATION TEST SUMMARY")
        print("="*60)
        print(f"Total Tests: {total}")
        print(f"Passed: {passed}")
        print(f"Failed: {total - passed}")
        print(f"Success Rate: {(passed/total)*100:.1f}%")
        
        # Print detailed results
        print("\n📋 Detailed Results:")
        for test_name, success, message in self.test_results:
            status = "✅ PASS" if success else "❌ FAIL"
            print(f"  {status} {test_name}: {message}")
        
        return passed == total


def run_ai_content_generation_tests():
    """Main function to run AI content generation tests"""
    tester = TestAIContentGeneration()
    return tester.run_all_tests()


if __name__ == "__main__":
    success = run_ai_content_generation_tests()
    exit(0 if success else 1)
