#!/usr/bin/env python3
"""
Unit tests for description query functionality in BatchProcessor

Tests the desc query feature that returns function documentation and schemas.
"""

import unittest
import asyncio
import sys
import os
from pathlib import Path
from typing import List, Generator, Any

# Add the parent directory to the path to import sandai
sys.path.insert(0, str(Path(__file__).parent.parent))

from sandai.operator.processor import BatchProcessor
from sandai.operator.types import TaskInput, TaskOutput
from pydantic import BaseModel


class TestOptions(BaseModel):
    """Test options for the processor function"""
    name: str
    count: int = 10
    enabled: bool = True


class TestResults(BaseModel):
    """Test results from the processor function"""
    message: str
    processed_count: int
    success: bool = True


class TestDescQuery(unittest.TestCase):
    """Test cases for description query functionality"""

    def setUp(self):
        """Set up test fixtures"""
        self.processor = BatchProcessor(
            "test-desc-processor", "v1.0", ignore_nongil_check=True
        )

        @self.processor.on_batch(max_batch_size=2)
        def test_process_function(
            batch_inputs: List[TaskInput[TestOptions]], 
            config: Any, 
            context: Any
        ) -> Generator[TaskOutput[TestResults], None, None]:
            """
            Test processing function for unit tests
            
            This function processes a batch of inputs and generates outputs.
            It demonstrates the use of TestOptions and TestResults types.
            
            Args:
                batch_inputs: List of task inputs with TestOptions
                config: Configuration object
                context: Processing context
                
            Yields:
                TaskOutput with TestResults for each input
            """
            for task_input in batch_inputs:
                result = TestResults(
                    message=f"Processed {task_input.options.name}",
                    processed_count=task_input.options.count * 2,
                    success=True
                )
                
                yield TaskOutput[TestResults](
                    task_id=task_input.task_id,
                    status="success",
                    results=result,
                    artifacts=[]
                )

    def test_function_description_preparation(self):
        """Test that function description is prepared during on_batch decoration"""
        # Check that function description was prepared
        self.assertIsNotNone(self.processor._function_description)
        
        desc = self.processor._function_description
        
        # Check that all required fields are present
        self.assertIn('doc', desc)
        self.assertIn('options_schema', desc)
        self.assertIn('results_schema', desc)
        
        # Check that doc is not empty
        self.assertIsNotNone(desc['doc'])
        self.assertNotEqual(desc['doc'].strip(), '')
        
        # Check that schemas are present
        self.assertIsNotNone(desc['options_schema'])
        self.assertIsNotNone(desc['results_schema'])

    def test_docstring_extraction(self):
        """Test that docstring is correctly extracted"""
        desc = self.processor._function_description
        doc = desc['doc']
        
        # Check that the docstring contains expected content
        self.assertIn('Test processing function for unit tests', doc)
        self.assertIn('This function processes a batch of inputs', doc)
        self.assertIn('Args:', doc)
        self.assertIn('Yields:', doc)

    def test_options_schema_correctness(self):
        """Test that options schema is correctly generated"""
        desc = self.processor._function_description
        options_schema = desc['options_schema']
        
        # Check schema structure
        self.assertEqual(options_schema['type'], 'object')
        self.assertIn('properties', options_schema)
        
        properties = options_schema['properties']
        
        # Check that all TestOptions fields are present
        self.assertIn('name', properties)
        self.assertIn('count', properties)
        self.assertIn('enabled', properties)
        
        # Check field types
        self.assertEqual(properties['name']['type'], 'string')
        self.assertEqual(properties['count']['type'], 'integer')
        self.assertEqual(properties['enabled']['type'], 'boolean')
        
        # Check default values
        self.assertEqual(properties['count']['default'], 10)
        self.assertEqual(properties['enabled']['default'], True)
        
        # Check required fields
        required = options_schema.get('required', [])
        self.assertIn('name', required)  # name has no default, so it's required

    def test_results_schema_correctness(self):
        """Test that results schema is correctly generated"""
        desc = self.processor._function_description
        results_schema = desc['results_schema']
        
        # Check schema structure
        self.assertEqual(results_schema['type'], 'object')
        self.assertIn('properties', results_schema)
        
        properties = results_schema['properties']
        
        # Check that all TestResults fields are present
        self.assertIn('message', properties)
        self.assertIn('processed_count', properties)
        self.assertIn('success', properties)
        
        # Check field types
        self.assertEqual(properties['message']['type'], 'string')
        self.assertEqual(properties['processed_count']['type'], 'integer')
        self.assertEqual(properties['success']['type'], 'boolean')
        
        # Check default values
        self.assertEqual(properties['success']['default'], True)
        
        # Check required fields
        required = results_schema.get('required', [])
        self.assertIn('message', required)
        self.assertIn('processed_count', required)

    def test_handle_desc_query_direct(self):
        """Test direct handling of desc query"""
        # Create a desc query input
        desc_input = {
            "task_id": "test_desc_001",
            "task_type": "desc"
        }
        
        # Handle the desc query
        results = self.processor._handle_desc_query(desc_input)
        
        # Check result structure
        self.assertEqual(results['task_id'], 'test_desc_001')
        self.assertEqual(results['status'], 'success')
        
        # Check results content
        self.assertIn('doc', results)
        self.assertIn('options_schema', results)
        self.assertIn('results_schema', results)

    def test_schema_validation_with_real_data(self):
        """Test that generated schemas can validate real data"""
        desc = self.processor._function_description
        
        # Test options schema with valid data
        options_schema = desc['options_schema']
        
        # This is a basic validation - in a real scenario you might use jsonschema
        properties = options_schema['properties']
        
        # Valid options data
        valid_options = {
            "name": "test_task",
            "count": 5,
            "enabled": False
        }
        
        # Check that all fields in valid_options match schema types
        for field, value in valid_options.items():
            self.assertIn(field, properties)
            prop = properties[field]
            
            if prop['type'] == 'string':
                self.assertIsInstance(value, str)
            elif prop['type'] == 'integer':
                self.assertIsInstance(value, int)
            elif prop['type'] == 'boolean':
                self.assertIsInstance(value, bool)

    def test_schema_completeness(self):
        """Test that schemas contain all necessary information"""
        desc = self.processor._function_description
        
        # Check options schema completeness
        options_schema = desc['options_schema']
        self.assertIn('type', options_schema)
        self.assertIn('properties', options_schema)
        
        # Each property should have a type
        for prop_name, prop_def in options_schema['properties'].items():
            self.assertIn('type', prop_def)
        
        # Check results schema completeness
        results_schema = desc['results_schema']
        self.assertIn('type', results_schema)
        self.assertIn('properties', results_schema)
        
        # Each property should have a type
        for prop_name, prop_def in results_schema['properties'].items():
            self.assertIn('type', prop_def)

    def test_error_handling_in_desc_query(self):
        """Test error handling in desc query processing"""
        # Temporarily break the function description
        original_desc = self.processor._function_description
        self.processor._function_description = None
        self.processor.process_func = None
        
        desc_input = {
            "task_id": "error_test",
            "task_type": "desc"
        }
        
        result = self.processor._handle_desc_query(desc_input)
        
        # Should handle the error gracefully
        self.assertEqual(result['task_id'], 'error_test')
        self.assertEqual(result['status'], 'success')  # Still success, but with error info
        
        self.assertIn('error', result)
        
        # Restore original state
        self.processor._function_description = original_desc

    def test_multiple_desc_queries(self):
        """Test handling multiple desc queries"""
        desc_inputs = [
            {"task_id": "desc_1", "task_type": "desc"},
            {"task_id": "desc_2", "task_type": "desc"},
            {"task_id": "desc_3", "task_type": "desc"}
        ]
        
        results = []
        for desc_input in desc_inputs:
            result = self.processor._handle_desc_query(desc_input)
            results.append(result)
        
        # All should succeed and return the same schema info
        for i, result in enumerate(results):
            self.assertEqual(result['task_id'], f'desc_{i+1}')
            self.assertEqual(result['status'], 'success')
            
            # All should have the same schema content
            if i > 0:
                self.assertEqual(
                    result['options_schema'],
                    results[0]['options_schema']
                )
                self.assertEqual(
                    result['results_schema'],
                    results[0]['results_schema']
                )


class TestDescQueryWithoutTypes(unittest.TestCase):
    """Test desc query with a processor that has no type information"""

    def setUp(self):
        """Set up a processor without proper type annotations"""
        self.processor = BatchProcessor(
            "no-types-processor", "v1.0", ignore_nongil_check=True
        )

        @self.processor.on_batch(max_batch_size=1)
        def simple_function(batch_inputs, config, context):
            """Simple function without type annotations"""
            for task_input in batch_inputs:
                yield {
                    "task_id": task_input.task_id,
                    "status": "success",
                    "results": {"message": "processed"}
                }

    def test_desc_query_without_types(self):
        """Test desc query when no type information is available"""
        desc_input = {
            "task_id": "no_types_test",
            "task_type": "desc"
        }
        
        result = self.processor._handle_desc_query(desc_input)
        
        # Should still work
        self.assertEqual(result['status'], 'success')
        self.assertIn('doc', result)


def run_desc_query_tests():
    """Run all desc query tests"""
    # Create test loader
    loader = unittest.TestLoader()
    
    # Create test suite
    suite = unittest.TestSuite()
    
    # Add test cases
    suite.addTests(loader.loadTestsFromTestCase(TestDescQuery))
    suite.addTests(loader.loadTestsFromTestCase(TestDescQueryWithoutTypes))
    
    # Run tests
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    
    return result.wasSuccessful()


if __name__ == '__main__':
    # Run the tests
    success = run_desc_query_tests()
    
    if success:
        print("\n🎉 All desc query tests passed!")
        sys.exit(0)
    else:
        print("\n💥 Some desc query tests failed!")
        sys.exit(1)