import json
import pika
from loguru import logger
from services.queue_service import queue_service
from services.test_service import test_service
from services.database_service import db_service
from models.base import TestCase, TestScript

def process_test_generation(ch, method, properties, body):
    """Process test generation messages"""
    try:
        data = json.loads(body)
        document_id = data.get('document_id')
        
        with db_service.get_session() as session:
            # Generate test cases
            test_cases = test_service.generate_test_cases(
                document_id=document_id,
                session=session
            )
            
            # Save test cases to database
            for test_case in test_cases:
                db_test_case = TestCase(
                    document_id=document_id,
                    title=test_case['title'],
                    description=test_case['description'],
                    steps=test_case['steps'],
                    expected_results=test_case['expected_results'],
                    priority=test_case.get('priority', 1)
                )
                session.add(db_test_case)
            
            # Generate test scripts
            for test_case in test_cases:
                scripts = test_service.generate_test_scripts(test_case)
                for script in scripts:
                    db_script = TestScript(
                        test_case_id=db_test_case.id,
                        language=script['language'],
                        code=script['code'],
                        framework=script['framework']
                    )
                    session.add(db_script)
        
        ch.basic_ack(delivery_tag=method.delivery_tag)
        logger.info(f"Successfully processed test generation for document {document_id}")
    
    except Exception as e:
        logger.error(f"Error processing test generation: {str(e)}")
        ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

def process_test_execution(ch, method, properties, body):
    """Process test execution messages"""
    try:
        data = json.loads(body)
        test_case_id = data.get('test_case_id')
        
        with db_service.get_session() as session:
            test_case = session.query(TestCase).get(test_case_id)
            if not test_case:
                raise ValueError(f"Test case {test_case_id} not found")
            
            # Execute test case
            result = test_service.execute_test_case(test_case)
            
            # Update test case status
            test_case.status = result['status']
            session.commit()
        
        ch.basic_ack(delivery_tag=method.delivery_tag)
        logger.info(f"Successfully executed test case {test_case_id}")
    
    except Exception as e:
        logger.error(f"Error executing test case: {str(e)}")
        ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

def main():
    """Start the worker"""
    try:
        # Start consuming test generation messages
        queue_service.consume_messages(
            queue="test_generation",
            callback=process_test_generation
        )
        
        # Start consuming test execution messages
        queue_service.consume_messages(
            queue="test_execution",
            callback=process_test_execution
        )
    
    except KeyboardInterrupt:
        logger.info("Worker stopped by user")
    except Exception as e:
        logger.error(f"Worker error: {str(e)}")
    finally:
        queue_service.close()

if __name__ == "__main__":
    main()