api/routes/document_routes.py

from fastapi import APIRouter, UploadFile, File, HTTPException, Depends
from typing import List, Dict, Any
from loguru import logger

from services.document_service import document_service
from services.ai_service import ai_service
from services.test_service import test_service
from services.automation_service import automation_service

router = APIRouter()

@router.post("/upload")
async def upload_document(
    file: UploadFile = File(...),
    process_type: str = "requirements"
) -> Dict[str, Any]:
    """
    Upload and process a document.
    
    Parameters:
    - file: The document file to upload
    - process_type: Type of processing to perform (requirements, test_cases, etc.)
    """
    try:
        # Save uploaded file
        file_path = await document_service.save_upload_file(file)
        
        # Process document
        result = await document_service.process_document(file_path)
        
        # Segment document if needed
        segments = await document_service.segment_document(result["text"])
        
        return {
            "status": "success",
            "file_path": file_path,
            "segments": segments,
            "type": result["type"]
        }
    except Exception as e:
        logger.error(f"Error processing document: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/process-requirements")
async def process_requirements(
    file: UploadFile = File(...),
    ai_provider: str = "openai",
    model: str = "gpt-3.5-turbo"
) -> Dict[str, Any]:
    """
    Process requirements document and generate test cases.
    
    Parameters:
    - file: The requirements document
    - ai_provider: AI provider to use (openai, local, openrouter)
    - model: Model to use for processing
    """
    try:
        # Upload and process document
        file_path = await document_service.save_upload_file(file)
        result = await document_service.process_document(file_path)
        
        # Extract requirements using AI
        prompt = f"""
        Extract requirements from the following text. For each requirement, provide:
        1. ID
        2. Title
        3. Description
        4. Priority (High/Medium/Low)
        
        Text:
        {result["text"]}
        """
        
        ai_response = await ai_service.generate_response(
            prompt=prompt,
            provider=ai_provider,
            model=model
        )
        
        # Generate test cases
        requirements = _parse_requirements(ai_response["response"])
        test_cases = await test_service.generate_test_cases(
            requirements=requirements,
            ai_service=ai_service
        )
        
        return {
            "status": "success",
            "requirements": requirements,
            "test_cases": test_cases
        }
    except Exception as e:
        logger.error(f"Error processing requirements: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/generate-test-scripts")
async def generate_test_scripts(
    test_cases: List[Dict[str, Any]],
    framework: str = "pytest",
    language: str = "python",
    browser: str = "chrome"
) -> Dict[str, Any]:
    """
    Generate test scripts from test cases.
    
    Parameters:
    - test_cases: List of test cases
    - framework: Test framework to use (pytest, playwright)
    - language: Programming language (python)
    - browser: Browser to use (chrome, firefox, etc.)
    """
    try:
        # Generate test scripts
        scripts = await automation_service.generate_test_scripts(
            test_cases=test_cases,
            framework=framework,
            language=language,
            browser=browser
        )
        
        # Generate Gherkin feature file
        feature = await automation_service.generate_gherkin_feature(test_cases)
        
        return {
            "status": "success",
            "scripts": scripts,
            "feature": feature
        }
    except Exception as e:
        logger.error(f"Error generating test scripts: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

def _parse_requirements(text: str) -> List[Dict[str, Any]]:
    """Parse AI response into structured requirements."""
    requirements = []
    current_req = {}
    
    for line in text.split('\n'):
        line = line.strip()
        if not line:
            continue
            
        if line.startswith('ID:'):
            if current_req:
                requirements.append(current_req)
            current_req = {'id': line.split(':', 1)[1].strip()}
        elif line.startswith('Title:'):
            current_req['title'] = line.split(':', 1)[1].strip()
        elif line.startswith('Description:'):
            current_req['description'] = line.split(':', 1)[1].strip()
        elif line.startswith('Priority:'):
            current_req['priority'] = line.split(':', 1)[1].strip()
    
    if current_req:
        requirements.append(current_req)
    
    return requirements