app/controllers/groq_client.py

# ============================================================================
#  GROQ AI CLIENT
# ============================================================================

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

from fastapi import HTTPException
from groq import Groq

from app.config import Config
from app.utils.logger import logger
from app.utils.document_processor import (
    DocumentProcessor, TokenCounter, ProcessedDocument
)
from app.utils.api_key_pool import api_key_pool


class GroqClient:
    """Groq AI test case generation client"""
    
    # Token limits for different models
    MODEL_TOKEN_LIMITS = {
        "openai/gpt-oss-120b": 8192,
        "llama-3.3-70b-versatile": 8192,
        "llama-3.1-70b-versatile": 8192,
        "llama3-70b-8192": 8192,
        "llama3-8b-8192": 8192,
        "mixtral-8x7b-32768": 32768,
    }
    
    DEFAULT_MAX_TOKENS = 8192
    
    @staticmethod
    def get_model_token_limit(model: str) -> int:
        """Get token limit for model"""
        for key, limit in GroqClient.MODEL_TOKEN_LIMITS.items():
            if key in model.lower():
                return limit
        return GroqClient.DEFAULT_MAX_TOKENS
    
    @staticmethod
    def build_system_prompt() -> str:
        """System prompt for test case generation - Frontend Compatible Format"""
        return """You are a Senior QA Engineer specialized in creating comprehensive test cases.

Your task is to generate detailed, actionable test cases from task descriptions.

**OUTPUT FORMAT:**
Return ONLY a valid JSON object with this exact structure (matches PM Tool Task interface):

{
  "test_cases": [
    {
      "title": "Clear, descriptive test case title",
      "description": "Detailed description of what this test case validates",
      "status": "Draft",
      "priority": "Critical|High|Medium|Low",
      "severity": "Blocker|Critical|Major|Minor|Trivial",
      "environment": "Development|Testing|Staging|Production",
      "stepsToReproduce": "1. First step\n2. Second step\n3. Third step\n4. Fourth step",
      "inputs": "Test data and inputs needed (as text)",
      "expectedBehavior": "What should happen when steps are executed correctly",
      "actualBehavior": "",
      "rootCause": "",
      "resolution": "",
      "version": "",
      "url": ""
    }
  ]
}

**CRITICAL RULES:**
1. **stepsToReproduce** MUST be a STRING with numbered steps separated by \n:
   - Correct: "1. Navigate to login page\n2. Enter email\n3. Click submit"
   - WRONG: ["step 1", "step 2"] (NOT an array!)
   
2. **priority** values: Critical, High, Medium, Low
3. **severity** values: Blocker, Critical, Major, Minor, Trivial
4. **status** MUST be "Draft" for all generated test cases
5. **environment** SHOULD be "Staging" by default
6. Leave these EMPTY (filled by testers later):
   - actualBehavior
   - rootCause
   - resolution

7. Generate diverse test scenarios:
   - Positive/Happy path cases (50%)
   - Negative/Error cases (30%)
   - Edge/Boundary cases (20%)

8. Each test case must have at least 3-5 detailed steps
9. Be SPECIFIC to the task description - no generic cases
10. Include realistic test data in "inputs" field
11. Return ONLY JSON - no markdown, explanations, or extra text"""
    
    @staticmethod
    def build_user_prompt(
        task: Dict[str, Any], 
        attachments: List[Dict[str, Any]], 
        count: int,
        document_context: str = None,
        existing_test_cases: List[Dict[str, str]] = None,
        additional_context: str = None
    ) -> str:
        """Build user prompt with task context and optional document content"""
        
        # Build attachment summary
        att_summary = ""
        if attachments:
            att_summary = "\n\n📎 **AVAILABLE DOCUMENTATION:**\n"
            for att in attachments[:5]:  # Limit to avoid token overflow
                att_summary += f"  • {att.get('fileName', att.get('filename', 'N/A'))} ({att.get('fileType', att.get('file_type', 'N/A'))})"
                if att.get('description'):
                    att_summary += f" - {att['description']}"
                att_summary += "\n"
        
        # Build document context section
        doc_context_section = ""
        if document_context:
            doc_context_section = f"""

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📄 **DOCUMENT CONTENT (for reference):**
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{document_context}

**Note:** Use the above document content to understand requirements, specifications, 
and business rules. Generate test cases that validate these requirements.
"""
        
        # Build existing test cases section to avoid duplicates
        existing_tc_section = ""
        if existing_test_cases and len(existing_test_cases) > 0:
            existing_tc_list = "\n".join([
                f"  {i+1}. {tc.get('title', 'Untitled')}"
                for i, tc in enumerate(existing_test_cases[:20])  # Limit to 20 to save tokens
            ])
            existing_tc_section = f"""

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
⚠️ **EXISTING TEST CASES - DO NOT DUPLICATE:**
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{existing_tc_list}

**CRITICAL:** These test cases already exist. Generate NEW, DIFFERENT test cases 
that cover scenarios NOT already covered by the existing test cases above.
"""
        
        # Build additional context section from user input
        additional_context_section = ""
        if additional_context:
            additional_context_section = f"""

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📝 **ADDITIONAL CONTEXT (User Provided):**
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{additional_context}

**Note:** This additional context was provided by the user. Use it to understand 
specific requirements, edge cases, or business rules that should be tested.
"""
        
        return f"""Generate {count} comprehensive test cases for this task:

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📋 TASK INFORMATION
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
• Task ID        : {task.get('id', 'N/A')}
• Task Code      : {task.get('taskCode', 'TXXX')}
• Title          : {task.get('title', 'Untitled')}
• Description    : {task.get('description', 'No description')}
• Project        : {task.get('projectName', 'N/A')}
• Issue          : {task.get('issueName', 'N/A')}
• Feature        : {task.get('deliverableName', 'N/A')}
• Status         : {task.get('status', 'N/A')}
• Priority       : {task.get('priority', 'Medium')}
• Assigned To    : {task.get('assignedTo', 'N/A')}
• Created By     : {task.get('createdByName', 'N/A')}
{att_summary}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{doc_context_section}{existing_tc_section}{additional_context_section}
**CONTEXT:**
- Feature: {task.get('deliverableName', 'N/A')}
- Related documentation: {len(attachments)} file(s)
- Existing test cases: {len(existing_test_cases) if existing_test_cases else 0}
- Additional context provided: {'Yes' if additional_context else 'No'}

**REQUIREMENTS:**
Generate exactly {count} test cases that:
1. Cover positive scenarios (happy path)
2. Include negative scenarios (error handling)
3. Test edge cases and boundaries
4. Validate data inputs/outputs where applicable
5. Are SPECIFIC to the task/feature described above
6. Are DIFFERENT from existing test cases listed above (if any)
7. Incorporate insights from additional context if provided

Return ONLY the JSON object - no additional text."""
    
    @staticmethod
    async def generate_test_cases(
        task: Dict[str, Any],
        attachments: List[Dict[str, Any]],
        count: int,
        model: str,
        temperature: float,
        processed_documents: List[ProcessedDocument] = None,
        max_context_tokens: int = None,
        existing_test_cases: List[Dict[str, str]] = None,
        additional_context: str = None
    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
        """Generate test cases using Groq AI with document context
        
        Args:
            existing_test_cases: List of existing test case titles to avoid duplicates
            additional_context: Extra context provided by user (text or PDF content)
        
        Returns: (test_cases, generation_metadata)
        """
        
        if not api_key_pool:
            raise HTTPException(
                status_code=500,
                detail="Groq API key not configured"
            )
        
        # Get model token limit
        model_limit = GroqClient.get_model_token_limit(model)
        
        # Reserve tokens for system prompt, response, and overhead
        system_prompt_tokens = 500  # Approximate
        response_reserve = 3000  # Reserve for response
        overhead = 200
        
        available_for_context = model_limit - system_prompt_tokens - response_reserve - overhead
        
        # Build document context if provided
        document_context = ""
        doc_stats = {
            "documents_processed": 0,
            "documents_skipped": 0,
            "chunks_used": 0,
            "tokens_used": 0,
            "context_truncated": False
        }
        
        if processed_documents:
            max_doc_tokens = min(
                max_context_tokens or DocumentProcessor.MAX_CONTEXT_TOKENS,
                available_for_context - 1000  # Reserve for task info
            )
            
            document_context, doc_tokens, doc_stats = DocumentProcessor.build_context_from_documents(
                documents=processed_documents,
                max_total_tokens=max_doc_tokens
            )
        
        try:
            # Get available API key from pool (shuffle)
            estimated_tokens = 2500  # Approximate for test case generation
            api_key = await api_key_pool.get_available_key(estimated_tokens)
            api_key_alias = api_key_pool.get_key_alias(api_key)
            
            logger.info(f"🔑 Using API key: {api_key_alias}")
            
            client = Groq(api_key=api_key)
            
            system_prompt = GroqClient.build_system_prompt()
            user_prompt = GroqClient.build_user_prompt(
                task=task,
                attachments=attachments,
                count=count,
                document_context=document_context if document_context else None,
                existing_test_cases=existing_test_cases,
                additional_context=additional_context
            )
            
            # Estimate total input tokens
            input_tokens = TokenCounter.estimate(system_prompt + user_prompt)
            
            logger.info(f"🤖 Generating {count} test cases for task {task.get('taskCode', 'N/A')}")
            logger.info(f"📊 Input tokens: ~{input_tokens} (limit: {model_limit})")
            
            if input_tokens > model_limit:
                # Truncate user prompt
                excess = input_tokens - model_limit + 500  # Buffer
                logger.warning(f"⚠️ Input exceeds limit, truncating by ~{excess} tokens")
                
                # Truncate document context first
                if document_context:
                    truncate_chars = int(excess * TokenCounter.CHARS_PER_TOKEN)
                    if len(document_context) > truncate_chars:
                        document_context = document_context[:-truncate_chars]
                        doc_stats["context_truncated"] = True
                        user_prompt = GroqClient.build_user_prompt(
                            task=task,
                            attachments=attachments,
                            count=count,
                            document_context=document_context,
                            existing_test_cases=existing_test_cases,
                            additional_context=additional_context
                        )
            
            # Call Groq API (synchronous, but run in thread pool to avoid blocking)
            loop = asyncio.get_event_loop()
            completion = await loop.run_in_executor(
                None,
                lambda: client.chat.completions.create(
                    model=model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    temperature=temperature,
                    max_completion_tokens=Config.GROQ_MAX_TOKENS,
                    top_p=1,
                    response_format={"type": "json_object"},
                    stop=None
                )
            )
            
            raw_content = completion.choices[0].message.content
            tokens_used = completion.usage.total_tokens if completion.usage else 0
            prompt_tokens = completion.usage.prompt_tokens if completion.usage else 0
            completion_tokens = completion.usage.completion_tokens if completion.usage else 0
            
            logger.info(f"✅ Groq response received | Key: {api_key_alias} | Total: {tokens_used} | Prompt: {prompt_tokens} | Completion: {completion_tokens}")
            
            # Record usage in API key pool
            await api_key_pool.record_usage(api_key, tokens_used)
            
            # Parse JSON
            try:
                parsed = json.loads(raw_content)
            except json.JSONDecodeError as e:
                logger.error(f"JSON parse error: {e}\nRaw: {raw_content[:500]}")
                raise HTTPException(
                    status_code=500,
                    detail=f"Failed to parse AI response: {str(e)}"
                )
            
            # Extract test cases
            test_cases = parsed.get("test_cases", [])
            
            if not test_cases:
                logger.warning(f"No test cases generated for task {task.get('id')}")
                return [], {"error": "No test cases generated"}
            
            # Add metadata to each test case (Frontend-compatible format)
            for tc in test_cases:
                # Ensure required fields exist
                tc['type'] = 'Test Case'
                tc['issuesId'] = task.get('issuesId')  # Link to Feature/Issue
                tc['projectId'] = task.get('projectId')
                
                # Ensure stepsToReproduce is a string (not array)
                if 'stepsToReproduce' not in tc or not tc.get('stepsToReproduce'):
                    # Convert test_steps array to string if present
                    if 'test_steps' in tc and isinstance(tc['test_steps'], list):
                        steps_text = '\n'.join(
                            f"{step.get('step_number', i+1)}. {step.get('action', '')}"
                            for i, step in enumerate(tc['test_steps'])
                        )
                        tc['stepsToReproduce'] = steps_text
                    else:
                        tc['stepsToReproduce'] = "1. Step to be defined"
                
                # Ensure expectedBehavior exists
                if 'expectedBehavior' not in tc or not tc.get('expectedBehavior'):
                    tc['expectedBehavior'] = tc.get('expected_outcome', '') or tc.get('objective', '')
                
                # Set defaults for empty fields
                tc['actualBehavior'] = tc.get('actualBehavior', '')
                tc['rootCause'] = tc.get('rootCause', '')
                tc['resolution'] = tc.get('resolution', '')
                tc['version'] = tc.get('version', '')
                tc['url'] = tc.get('url', '')
                tc['environment'] = tc.get('environment', 'Staging')
                tc['status'] = tc.get('status', 'Draft')
                
                # Map priority/severity to frontend values
                priority = tc.get('priority', 'Medium')
                if priority not in ['Critical', 'High', 'Medium', 'Low']:
                    tc['priority'] = 'Medium'
                
                severity = tc.get('severity', 'Minor')
                if severity not in ['Blocker', 'Critical', 'Major', 'Minor', 'Trivial']:
                    tc['severity'] = 'Minor'
                
                # Add reference metadata (for display purposes)
                tc['related_task_id'] = task.get('id')
                tc['related_task_code'] = task.get('taskCode', 'TXXX')
                tc['projectName'] = task.get('projectName', 'N/A')
                tc['issueName'] = task.get('issueName', 'N/A')
                tc['deliverableName'] = task.get('deliverableName', 'N/A')
                tc['createdByName'] = task.get('createdByName', 'N/A')
                
                # Set dates
                tc['stateDate'] = datetime.now().isoformat()
                tc['endDate'] = None
                tc['assignedTo'] = None
                
                # Clean up legacy fields that shouldn't be in final output
                for legacy_field in ['test_case_id', 'objective', 'preconditions', 'test_steps', 
                                     'expected_outcome', 'test_type', 'category', 'task_id', 
                                     'task_code', 'issue_name', 'feature_name', 'task_name', 
                                     'task_description', 'created_by']:
                    tc.pop(legacy_field, None)
            
            # Build generation metadata
            gen_metadata = {
                "model": model,
                "temperature": temperature,
                "api_key_alias": api_key_alias,
                "tokens_used": tokens_used,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "document_stats": doc_stats,
                "test_cases_count": len(test_cases)
            }
            
            logger.info(f"✅ Generated {len(test_cases)} test cases")
            return test_cases, gen_metadata
            
        except Exception as e:
            logger.error(f"Groq API error: {e}\n{traceback.format_exc()}")
            raise HTTPException(
                status_code=500,
                detail=f"AI generation failed: {str(e)}"
            )