api.py

"""
FastAPI Backend Service for UI Agent Testing

This service provides REST API endpoints to trigger and manage
UI testing with the agent_ui_rater_openai_complete module.
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from datetime import datetime
import os
import json
import sys
import threading
import traceback
from dotenv import load_dotenv

load_dotenv()

# Import the agent module
try:
    from agent_ui_rater_openai import (
        run_agent, 
        PlaywrightBridge, 
        RUNS_DIR,
        USER_PROFILES
    )
    import agent_ui_rater_openai as agent_module
    AGENT_AVAILABLE = True
except ImportError as e:
    print(f"Error: Could not import agent module: {e}")
    AGENT_AVAILABLE = False
    RUNS_DIR = os.path.join(os.path.dirname(__file__), "runs")

# Create FastAPI app
app = FastAPI(
    title="UI Agent Testing API",
    description="API for running automated UI testing with AI agents",
    version="1.0.0"
)

# Enable CORS for front-end communication
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify your front-end URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Ensure runs directory exists
if not os.path.exists(RUNS_DIR):
    os.makedirs(RUNS_DIR)

# Pydantic models for request/response validation
class TestRunRequest(BaseModel):
    """Request model for running a UI test"""
    url: str = Field(..., description="Target URL to test")
    profile: str = Field(default="AVERAGE", description="Tester profile (VERY_INEXPERIENCED, AVERAGE, TECH_SAVVY)")
    device_name: str = Field(default="iPhone 12", description="Device to emulate")
    browser_type: str = Field(default="Chrome", description="Browser type")
    resolution: str = Field(default="390x844", description="Screen resolution")
    language: str = Field(default="English", description="Testing language")
    ux_requirements: Optional[List[str]] = Field(default=None, description="UX flow requirements")
    ui_requirements: Optional[List[str]] = Field(default=None, description="UI consistency requirements")
    custom_mission: Optional[str] = Field(default=None, description="Custom testing objective")

class TestRunResponse(BaseModel):
    """Response model for test run initiation"""
    status: str
    message: str
    run_id: str
    estimated_duration: str = "5-10 minutes"

class TestResultSummary(BaseModel):
    """Summary of a test result"""
    run_id: str
    url: str
    test_date: str
    overall_score: int
    grade: str
    device: str
    passed_tests: int
    total_tests: int
    status: str = "completed"

class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    agent_available: bool
    runs_directory: str
    timestamp: str


# Global store for running tests (in production, use a proper database)
running_tests: Dict[str, Dict[str, Any]] = {}


def execute_agent_test(
    run_id: str,
    url: str,
    profile: str,
    device_name: str,
    browser_type: str,
    resolution: str,
    language: str,
    ux_requirements: Optional[List[str]],
    ui_requirements: Optional[List[str]],
    custom_mission: Optional[str]
):
    """Execute the agent test in a background thread"""
    try:
        print(f"[{run_id}] Starting agent test...")
        running_tests[run_id]["status"] = "running"
        
        # Reset agent globals
        agent_module.REPORT_BUFFER = []
        agent_module.SCREENSHOT_LIST = []
        agent_module.REPORT_TIMESTAMP = run_id
        
        # Create run directory
        run_dir = os.path.join(RUNS_DIR, run_id)
        if not os.path.exists(run_dir):
            os.makedirs(run_dir)
        
        # Set global current run directory
        agent_module.CURRENT_RUN_DIR = run_dir
        
        # Initialize Playwright bridge
        agent_module.BRIDGE = PlaywrightBridge(
            target_url=url,
            device_name=device_name,
            browser_type=browser_type,
            resolution=resolution,
            record_video_dir=run_dir
        )
        
        print(f"[{run_id}] Running agent with profile: {profile}")
        
        # Run the agent
        agent_module.run_agent(
            profile,
            target_url=url,
            language=language,
            ux_requirements=ux_requirements,
            ui_requirements=ui_requirements,
            device_name=device_name,
            browser_type=browser_type,
            resolution=resolution,
            custom_mission=custom_mission
        )
        
        # Cleanup
        if agent_module.BRIDGE:
            agent_module.BRIDGE.shutdown()
        
        # Rename video recording to standard name
        try:
            video_files = [f for f in os.listdir(run_dir) if f.endswith('.webm') and f != 'recording.webm']
            if video_files:
                latest_video = sorted(video_files)[-1]
                old_path = os.path.join(run_dir, latest_video)
                new_path = os.path.join(run_dir, "recording.webm")
                if os.path.exists(old_path):
                    os.rename(old_path, new_path)
                    print(f"[{run_id}] Renamed recording to: recording.webm")
        except Exception as e:
            print(f"[{run_id}] Error renaming recording: {e}")
        
        running_tests[run_id]["status"] = "completed"
        running_tests[run_id]["completed_at"] = datetime.now().isoformat()
        print(f"[{run_id}] Test completed successfully")
        
    except Exception as e:
        error_msg = f"Agent execution failed: {str(e)}"
        print(f"[{run_id}] Error: {error_msg}")
        traceback.print_exc()
        
        running_tests[run_id]["status"] = "failed"
        running_tests[run_id]["error"] = error_msg
        running_tests[run_id]["completed_at"] = datetime.now().isoformat()


@app.get("/", tags=["General"])
async def root():
    """Root endpoint"""
    return {
        "message": "UI Agent Testing API",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/api/health"
    }


@app.get("/api/health", response_model=HealthResponse, tags=["General"])
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy",
        agent_available=AGENT_AVAILABLE,
        runs_directory=RUNS_DIR,
        timestamp=datetime.now().isoformat()
    )


@app.post("/api/run-test", response_model=TestRunResponse, tags=["Testing"])
async def run_test(request: TestRunRequest, background_tasks: BackgroundTasks):
    """
    Trigger a new UI test run
    
    This endpoint initiates a new test run with the specified parameters.
    The test runs in the background and results can be retrieved via the
    /api/results/{run_id} endpoint.
    """
    if not AGENT_AVAILABLE:
        raise HTTPException(
            status_code=503,
            detail="Agent module not available. Please check server configuration."
        )
    
    # Validate profile
    if request.profile not in USER_PROFILES and not request.profile.startswith("Custom:"):
        raise HTTPException(
            status_code=400,
            detail=f"Invalid profile. Must be one of: {', '.join(USER_PROFILES.keys())} or start with 'Custom:'"
        )
    
    # Generate run ID (timestamp)
    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Store initial run info
    running_tests[run_id] = {
        "status": "pending",
        "url": request.url,
        "profile": request.profile,
        "device": request.device_name,
        "started_at": datetime.now().isoformat()
    }
    
    # Execute test in background
    background_tasks.add_task(
        execute_agent_test,
        run_id=run_id,
        url=request.url,
        profile=request.profile,
        device_name=request.device_name,
        browser_type=request.browser_type,
        resolution=request.resolution,
        language=request.language,
        ux_requirements=request.ux_requirements,
        ui_requirements=request.ui_requirements,
        custom_mission=request.custom_mission
    )
    
    return TestRunResponse(
        status="accepted",
        message="Test run initiated successfully",
        run_id=run_id
    )


@app.get("/api/results/{run_id}", tags=["Testing"])
async def get_results(run_id: str):
    """
    Get results for a specific test run
    
    Returns the complete results JSON for the specified run ID.
    """
    run_dir = os.path.join(RUNS_DIR, run_id)
    results_path = os.path.join(run_dir, "results.json")
    
    if not os.path.exists(results_path):
        # Check if test is still running
        if run_id in running_tests:
            return {
                "status": running_tests[run_id]["status"],
                "run_id": run_id,
                "message": "Test is still running. Please check back later."
            }
        
        raise HTTPException(status_code=404, detail="Run not found")
    
    try:
        with open(results_path, 'r') as f:
            results = json.load(f)
        
        # Add status from running_tests if available
        if run_id in running_tests:
            results["status"] = running_tests[run_id]["status"]
        else:
            results["status"] = "completed"
        
        return results
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error loading results: {str(e)}"
        )


@app.get("/api/history", response_model=List[TestResultSummary], tags=["Testing"])
async def get_history():
    """
    Get list of all test runs
    
    Returns a summary of all historical test runs.
    """
    history = []
    
    if not os.path.exists(RUNS_DIR):
        return history
    
    try:
        run_dirs = [d for d in os.listdir(RUNS_DIR) if os.path.isdir(os.path.join(RUNS_DIR, d))]
        
        for run_id in run_dirs:
            results_path = os.path.join(RUNS_DIR, run_id, "results.json")
            if os.path.exists(results_path):
                try:
                    with open(results_path, 'r') as f:
                        data = json.load(f)
                    
                    summary = data.get('summary', {})
                    report_info = data.get('report_info', {})
                    
                    # Determine status
                    status = "completed"
                    if run_id in running_tests:
                        status = running_tests[run_id]["status"]
                    
                    history.append(TestResultSummary(
                        run_id=run_id,
                        url=report_info.get('url', 'Unknown'),
                        test_date=report_info.get('test_date', 'Unknown'),
                        overall_score=summary.get('overall_score', 0),
                        grade=summary.get('grade', 'N/A'),
                        device=report_info.get('environment', 'Unknown'),
                        passed_tests=summary.get('passed_tests', 0),
                        total_tests=summary.get('total_tests', 0),
                        status=status
                    ))
                except:
                    continue
        
        # Sort by run_id (timestamp) descending
        history.sort(key=lambda x: x.run_id, reverse=True)
        
        return history
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error loading history: {str(e)}"
        )


@app.get("/api/status/{run_id}", tags=["Testing"])
async def get_status(run_id: str):
    """
    Get current status of a test run
    
    Useful for polling while a test is running.
    """
    if run_id in running_tests:
        return running_tests[run_id]
    
    # Check if results exist
    results_path = os.path.join(RUNS_DIR, run_id, "results.json")
    if os.path.exists(results_path):
        return {
            "status": "completed",
            "run_id": run_id
        }
    
    raise HTTPException(status_code=404, detail="Run not found")


@app.delete("/api/results/{run_id}", tags=["Testing"])
async def delete_results(run_id: str):
    """
    Delete results for a specific test run
    """
    run_dir = os.path.join(RUNS_DIR, run_id)
    
    if not os.path.exists(run_dir):
        raise HTTPException(status_code=404, detail="Run not found")
    
    try:
        import shutil
        shutil.rmtree(run_dir)
        
        # Remove from running_tests if present
        if run_id in running_tests:
            del running_tests[run_id]
        
        return {"status": "success", "message": f"Run {run_id} deleted successfully"}
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error deleting run: {str(e)}"
        )


if __name__ == "__main__":
    import uvicorn
    print("Starting FastAPI server for UI Agent Testing...")
    print(f"Agent available: {AGENT_AVAILABLE}")
    print(f"Runs directory: {RUNS_DIR}")
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")