app/controllers/test_case_generator.py

# ============================================================================
#  MAIN TEST CASE GENERATION ORCHESTRATOR
# ============================================================================

import asyncio
import time
import traceback
import os
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple

from fastapi import HTTPException

from app.config import Config
from app.models.schemas import (
    GenerateRequest, GenerateResponse, TestCase, TestStep,
    DocumentProcessingConfig, AttachmentConfig
)
from app.models.enums import TestCaseType, TestCasePriority, TestCaseCategory, OutputFormat
from app.controllers.pm_tool_client import PMToolClient
from app.controllers.groq_client import GroqClient
from app.controllers.task_filter import TaskFilter
from app.utils.logger import logger
from app.utils.output_formatter import OutputFormatter
from app.utils.document_processor import (
    DocumentProcessor, ProcessedDocument
)


class TestCaseGenerator:
    """Main orchestrator for test case generation"""
    
    @staticmethod
    async def process_documents_for_task(
        task: Dict[str, Any],
        token: str,
        request: GenerateRequest,
        download_dir: str = None
    ) -> Tuple[List[ProcessedDocument], List[Dict[str, Any]], Dict[str, Any]]:
        """Process documents for a task
        
        Downloads attachments, extracts text, and builds context for LLM.
        
        Returns: (processed_documents, attachment_info, doc_stats)
        """
        task_id = task.get('id')
        project_name = task.get('projectName')
        
        # Get document processing config
        doc_config = request.document_processing or DocumentProcessingConfig()
        
        if not doc_config.enabled:
            return [], [], {"enabled": False}
        
        # Fetch attachments by project name using dynamic mapping
        attachments = await PMToolClient.fetch_attachments_by_project(token, project_name)
        
        if not attachments:
            return [], [], {"enabled": True, "attachments_found": 0}
        
        # Create temp directory for downloads
        if not download_dir:
            download_dir = os.path.join(Config.OUTPUT_DIR, "temp_attachments")
        os.makedirs(download_dir, exist_ok=True)
        
        # Download and process each attachment
        processed_docs = []
        attachment_summaries = []
        doc_stats = {
            "enabled": True,
            "attachments_found": len(attachments),
            "attachments_downloaded": 0,
            "attachments_skipped": 0,
            "documents_processed": 0,
            "total_tokens": 0,
            "errors": []
        }
        
        for att in attachments:
            att_id = att.get('attachmentId')
            filename = att.get('fileName')
            file_type = att.get('fileType', '')
            
            # Check if file type is supported
            if not DocumentProcessor.is_processable(filename):
                doc_stats["attachments_skipped"] += 1
                logger.info(f"⏭️ Skipping unsupported file type: {filename}")
                attachment_summaries.append({
                    "attachment_id": att_id,
                    "filename": filename,
                    "file_type": file_type,
                    "description": att.get('description', ''),
                    "skipped": True,
                    "skip_reason": "Unsupported file type"
                })
                continue
            
            # Download attachment
            try:
                file_path, content = await PMToolClient.download_attachment(
                    token=token,
                    attachment_id=att_id,
                    download_dir=download_dir
                )
                
                if not file_path and not content:
                    doc_stats["errors"].append({
                        "attachment_id": att_id,
                        "filename": filename,
                        "error": "Download failed"
                    })
                    continue
                
                # Process document (extract text)
                processed = await DocumentProcessor.process_document(
                    file_path=file_path,
                    document_id=str(att_id),
                    max_tokens=doc_config.max_chunk_tokens
                )
                
                if not processed.skipped:
                    processed_docs.append(processed)
                    doc_stats["documents_processed"] += 1
                    doc_stats["total_tokens"] += processed.total_tokens
                    doc_stats["attachments_downloaded"] += 1
                    
                    # Add to summaries with extracted text preview
                    text_preview = ""
                    if processed.chunks:
                        text_preview = processed.chunks[0].content[:500] + "..."
                    
                    attachment_summaries.append({
                        "attachment_id": att_id,
                        "filename": filename,
                        "file_type": file_type,
                        "description": att.get('description', ''),
                        "tokens": processed.total_tokens,
                        "chunks": len(processed.chunks),
                        "text_preview": text_preview,
                        "skipped": False
                    })
                else:
                    doc_stats["attachments_skipped"] += 1
                    attachment_summaries.append({
                        "attachment_id": att_id,
                        "filename": filename,
                        "file_type": file_type,
                        "description": att.get('description', ''),
                        "skipped": True,
                        "skip_reason": processed.skip_reason
                    })
                    
            except Exception as e:
                logger.error(f"Error processing attachment {filename}: {e}")
                doc_stats["errors"].append({
                    "attachment_id": att_id,
                    "filename": filename,
                    "error": str(e)
                })
        
        logger.info(f"📎 Processed {doc_stats['documents_processed']}/{len(attachments)} attachments for project: {project_name}")
        
        return processed_docs, attachment_summaries, doc_stats
    
    @staticmethod
    async def process_task(
        task: Dict[str, Any],
        token: str,
        request: GenerateRequest,
        download_dir: str = None
    ) -> Tuple[List[TestCase], Optional[Dict[str, Any]]]:
        """Process a single task and generate test cases"""
        
        task_id = task.get('id')
        task_code = task.get('taskCode', 'TXXX')
        
        try:
            # Get document processing config
            doc_config = request.document_processing or DocumentProcessingConfig()
            
            # Process documents if enabled
            processed_docs = []
            attachments = []
            doc_stats = {}
            
            if request.include_attachments and doc_config.enabled:
                processed_docs, attachments, doc_stats = await TestCaseGenerator.process_documents_for_task(
                    task=task,
                    token=token,
                    request=request,
                    download_dir=download_dir
                )
            
            # Generate test cases using AI (with document context)
            raw_test_cases, gen_metadata = await GroqClient.generate_test_cases(
                task=task,
                attachments=attachments,
                count=request.test_case_count,
                model=request.groq_model,
                temperature=request.temperature,
                processed_documents=processed_docs if processed_docs else None,
                max_context_tokens=doc_config.max_total_tokens
            )
            
            # Convert to TestCase models
            test_cases = []
            for tc_data in raw_test_cases:
                try:
                    # Parse test steps
                    steps = [
                        TestStep(**step) for step in tc_data.get('test_steps', [])
                    ]
                    
                    # Parse test_type - handle combined values like "Functional|Data Validation"
                    test_type_str = tc_data.get('test_type', 'Functional')
                    if '|' in test_type_str:
                        test_type_str = test_type_str.split('|')[0].strip()
                    try:
                        test_type = TestCaseType(test_type_str)
                    except ValueError:
                        test_type = TestCaseType.FUNCTIONAL
                    
                    # Parse priority
                    priority_str = tc_data.get('priority', 'Medium')
                    try:
                        priority = TestCasePriority(priority_str)
                    except ValueError:
                        priority = TestCasePriority.MEDIUM
                    
                    # Parse category - handle combined values
                    category_str = tc_data.get('category', 'Positive')
                    if '|' in category_str:
                        category_str = category_str.split('|')[0].strip()
                    try:
                        category = TestCaseCategory(category_str)
                    except ValueError:
                        category = TestCaseCategory.POSITIVE
                    
                    # Build metadata
                    metadata = None
                    if request.include_metadata:
                        metadata = {
                            "generated_at": datetime.now().isoformat(),
                            "groq_model": request.groq_model,
                            "temperature": request.temperature,
                            "attachments_count": len(attachments),
                            "documents_processed": len(processed_docs),
                            "document_tokens": doc_stats.get('total_tokens', 0),
                            "task_status": task.get('status'),
                            "gen_metadata": gen_metadata
                        }
                    
                    tc = TestCase(
                        test_case_id=tc_data.get('test_case_id', f"TC-{task_code}-000"),
                        task_id=task_id,
                        task_code=task_code,
                        project_name=tc_data.get('project_name', 'N/A'),
                        issue_name=tc_data.get('issue_name', 'N/A'),
                        feature_name=tc_data.get('feature_name', 'N/A'),
                        task_name=tc_data.get('task_name', 'Untitled'),
                        task_description=tc_data.get('task_description', ''),
                        created_by=tc_data.get('created_by', 'N/A'),
                        title=tc_data.get('title', 'Untitled'),
                        objective=tc_data.get('objective', ''),
                        preconditions=tc_data.get('preconditions', []),
                        test_steps=steps,
                        expected_outcome=tc_data.get('expected_outcome', ''),
                        test_type=test_type,
                        priority=priority,
                        category=category,
                        metadata=metadata
                    )
                    test_cases.append(tc)
                    
                except Exception as e:
                    logger.warning(f"Skipping malformed test case: {e}")
                    continue
            
            return test_cases, None, doc_stats
            
        except Exception as e:
            error_info = {
                "task_id": task_id,
                "task_code": task_code,
                "error": str(e),
                "traceback": traceback.format_exc()
            }
            logger.error(f"Failed to process task {task_id}: {e}")
            return [], error_info, {}
    
    @staticmethod
    async def generate(request: GenerateRequest) -> GenerateResponse:
        """Main generation pipeline"""
        
        start_time = time.time()
        
        # Step 1: Authenticate
        auth_data = await PMToolClient.login(request.credentials)
        token = auth_data['token']
        
        # Step 2: Resolve project if project_selection provided
        resolved_project = None
        project_id = request.project_id
        
        if request.project_selection:
            resolved_project = await PMToolClient.resolve_project(
                token=token,
                project_id=request.project_selection.project_id,
                project_name=request.project_selection.project_name
            )
            if resolved_project:
                project_id = resolved_project.get('id')
        
        # Step 3: Fetch tasks
        task_ids = request.task_ids or ([request.task_id] if request.task_id else None)
        
        if project_id and not task_ids:
            # Fetch all tasks for project
            all_tasks = await PMToolClient.fetch_tasks_by_project(
                token=token,
                project_id=project_id,
                task_type=request.task_type
            )
        else:
            all_tasks = await PMToolClient.fetch_tasks(
                token=token,
                task_type=request.task_type,
                task_ids=task_ids
            )
        
        if not all_tasks:
            raise HTTPException(
                status_code=404,
                detail="No tasks found with the given criteria"
            )
        
        # Step 4: Apply filters
        filtered_tasks = TaskFilter.apply_filters(all_tasks, request)
        
        if not filtered_tasks:
            raise HTTPException(
                status_code=404,
                detail="No tasks match the filter criteria"
            )
        
        # Step 5: Setup download directory for documents
        doc_config = request.document_processing or DocumentProcessingConfig()
        download_dir = None
        if doc_config.enabled and doc_config.auto_download:
            download_dir = doc_config.download_dir or os.path.join(
                Config.OUTPUT_DIR, 
                f"attachments_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            )
            os.makedirs(download_dir, exist_ok=True)
        
        # Step 6: Process tasks concurrently
        logger.info(f"🚀 Processing {len(filtered_tasks)} tasks concurrently (max {Config.MAX_CONCURRENT_TASKS})")
        
        all_test_cases = []
        failed_tasks = []
        total_doc_stats = {
            "attachments_found": 0,
            "documents_processed": 0,
            "total_tokens": 0
        }
        
        # Use semaphore to limit concurrency
        semaphore = asyncio.Semaphore(Config.MAX_CONCURRENT_TASKS)
        
        async def process_with_semaphore(task):
            async with semaphore:
                return await TestCaseGenerator.process_task(task, token, request, download_dir)
        
        # Process all tasks
        results = await asyncio.gather(
            *[process_with_semaphore(task) for task in filtered_tasks],
            return_exceptions=True
        )
        
        # Collect results
        for result in results:
            if isinstance(result, Exception):
                logger.error(f"Task processing exception: {result}")
                failed_tasks.append({
                    "error": str(result),
                    "traceback": traceback.format_exc()
                })
            else:
                test_cases, error_info, task_doc_stats = result
                all_test_cases.extend(test_cases)
                if error_info:
                    failed_tasks.append(error_info)
                # Aggregate doc stats
                if task_doc_stats:
                    total_doc_stats["attachments_found"] += task_doc_stats.get("attachments_found", 0)
                    total_doc_stats["documents_processed"] += task_doc_stats.get("documents_processed", 0)
                    total_doc_stats["total_tokens"] += task_doc_stats.get("total_tokens", 0)
        
        # Step 7: Build metadata first (needed for JSON output)
        metadata = None
        if request.include_metadata:
            metadata = {
                "generation_time_seconds": round(time.time() - start_time, 2),
                "total_tasks_found": len(all_tasks),
                "total_tasks_filtered": len(filtered_tasks),
                "groq_model": request.groq_model,
                "temperature": request.temperature,
                "timestamp": datetime.now().isoformat(),
                "user_email": request.credentials.email,
                "project_resolved": resolved_project.get('name') if resolved_project else None,
                "document_processing": {
                    "enabled": doc_config.enabled,
                    "stats": total_doc_stats
                }
            }
        
        # Step 8: Generate output file if requested
        download_url = None
        if request.output_format != OutputFormat.JSON or len(all_test_cases) > 0:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            if request.output_format == OutputFormat.EXCEL:
                filename = f"test_cases_{timestamp}.xlsx"
                filepath = OutputFormatter.to_excel(all_test_cases, filename)
                download_url = f"/download/{filename}"
            
            elif request.output_format == OutputFormat.JSON:
                filename = f"test_cases_{timestamp}.json"
                filepath = OutputFormatter.to_json(all_test_cases, filename, metadata)
                download_url = f"/download/{filename}"
        
        # Step 9: Build response
        response = GenerateResponse(
            success=len(failed_tasks) == 0,
            message=f"Generated {len(all_test_cases)} test cases from {len(filtered_tasks)} tasks",
            total_tasks_processed=len(filtered_tasks),
            total_test_cases_generated=len(all_test_cases),
            failed_tasks=failed_tasks,
            test_cases=all_test_cases,
            metadata=metadata,
            download_url=download_url
        )
        
        logger.info(f"✅ Generation complete: {len(all_test_cases)} test cases in {metadata['generation_time_seconds']}s")
        
        return response