modules/api.py

import time
import os
import json
import logging
import numpy as np
from fastapi import APIRouter, UploadFile, File, BackgroundTasks, HTTPException, Request
from fastapi.responses import JSONResponse
from modules.element_processing import process_screenshot
from pydantic import BaseModel
from typing import List, Dict, Optional, Any, Tuple

# Custom JSON encoder to handle numpy types
class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NumpyEncoder, self).default(obj)

logger = logging.getLogger(__name__)
router = APIRouter()

class ElementResponse(BaseModel):
    code: str
    type: str
    text_content: Optional[str] = None
    object_label: Optional[str] = None
    bbox_normalized: List[float]
    bbox_pixels: Optional[List[int]] = None
    center_x: int
    center_y: int

class ScreenshotRequest(BaseModel):
    screen_width: int
    screen_height: int

class ApiResponse(BaseModel):
    output: Any
    image_url: str

class ScreenshotDataURIRequest(BaseModel):
    image_data_uri: str
    screen_width: Optional[int] = None
    screen_height: Optional[int] = None

@router.post("/process_screenshot/")
async def process_screenshot_endpoint(
    request: Request,
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    screen_width: int = None,
    screen_height: int = None
) -> JSONResponse:
    """
    Process a screenshot image to identify UI elements, assign codes, and return element data.
    
    Args:
        request: FastAPI request object
        background_tasks: FastAPI background tasks
        file: The uploaded screenshot image file
        screen_width: Target screen width (for coordinate scaling)
        screen_height: Target screen height (for coordinate scaling)
    
    Returns:
        JSON with elements array and annotated image URL
    """
    start_time = time.time()
    logger.info(f"Processing screenshot: {file.filename}")
    
    try:
        # Read image data
        image_data = await file.read()
        
        if not image_data:
            raise HTTPException(status_code=400, detail="Empty image data")
        
        # Process the screenshot
        elements, image_path = await process_screenshot(image_data, background_tasks, screen_width, screen_height)
        
        # Create full URL for the image
        base_url = str(request.base_url).rstrip('/')
        image_url = f"{base_url}/{image_path}" if image_path else ""
        
        # Log processing time
        processing_time = time.time() - start_time
        logger.info(f"Screenshot processed in {processing_time:.2f} seconds")
        
        # Use the custom JSON encoder to handle numpy types
        response_data = {
            "output": elements,
            "image_url": image_url
        }
        
        return JSONResponse(content=json.loads(json.dumps(response_data, cls=NumpyEncoder)))
        
    except Exception as e:
        logger.error(f"Error processing screenshot: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/process_screenshot_string/")
async def process_screenshot_string_endpoint(
    request: Request,
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    screen_width: int = None,
    screen_height: int = None
) -> JSONResponse:
    """
    Process a screenshot and return elements as a formatted string.
    Each line represents an icon in the format:
    'icon CODE: {'type': 'text/object', 'centerX': x, 'centerY': y, 'content': 'Text'}'
    
    Args:
        request: FastAPI request object
        background_tasks: FastAPI background tasks
        file: The uploaded screenshot image file
        screen_width: Target screen width (for coordinate scaling)
        screen_height: Target screen height (for coordinate scaling)
    
    Returns:
        JSON with string output and annotated image URL
    """
    start_time = time.time()
    logger.info(f"Processing screenshot for string output: {file.filename}")
    
    try:
        # Read image data
        image_data = await file.read()
        
        if not image_data:
            raise HTTPException(status_code=400, detail="Empty image data")
        
        # Process the screenshot
        elements, image_path = await process_screenshot(image_data, background_tasks, screen_width, screen_height)
        
        # Create full URL for the image
        base_url = str(request.base_url).rstrip('/')
        image_url = f"{base_url}/{image_path}" if image_path else ""
        
        # Format elements as string
        result_lines = []
        for element in elements:
            # Determine content based on element type
            if element["type"] == "text":
                content_value = element.get("text_content", "")
            else:
                content_value = element.get("object_label", "")
            
            # Format the element string
            element_str = (f"icon {element['code']}: {{"
                          f"'type': '{element['type']}', "
                          f"'centerX': {element['center_x']}, "
                          f"'centerY': {element['center_y']}, "
                          f"'content': '{content_value}'}}")
            
            result_lines.append(element_str)
        
        # Join all lines
        result_string = "\n".join(result_lines)
        
        # Log processing time
        processing_time = time.time() - start_time
        logger.info(f"Screenshot processed for string output in {processing_time:.2f} seconds")
        
        # Use the custom JSON encoder to handle numpy types
        response_data = {
            "output": result_string,
            "image_url": image_url
        }
        
        return JSONResponse(content=response_data)
        
    except Exception as e:
        logger.error(f"Error processing screenshot for string output: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/process_screenshot_data_uri/")
async def process_screenshot_data_uri_endpoint(
    request: Request,
    background_tasks: BackgroundTasks,
    data: ScreenshotDataURIRequest
) -> JSONResponse:
    """
    Process a screenshot from data URI to identify UI elements, assign codes, and return element data.
    
    Args:
        request: FastAPI request object
        background_tasks: FastAPI background tasks
        data: JSON containing image data URI and screen dimensions
    
    Returns:
        JSON with string output and annotated image URL
    """
    start_time = time.time()
    logger.info("Processing screenshot from data URI")
    
    try:
        # Extract image data from data URI
        import base64
        if not data.image_data_uri.startswith('data:image'):
            raise HTTPException(status_code=400, detail="Invalid data URI format")
        
        # Split the header and the base64 data
        header, encoded = data.image_data_uri.split(",", 1)
        image_data = base64.b64decode(encoded)
        
        if not image_data:
            raise HTTPException(status_code=400, detail="Empty image data")
        
        # Process the screenshot
        elements, image_path = await process_screenshot(
            image_data, 
            background_tasks, 
            data.screen_width, 
            data.screen_height
        )
        
        # Create full URL for the image
        base_url = str(request.base_url).rstrip('/')
        image_url = f"{base_url}/{image_path}" if image_path else ""
        
        # Format elements as string
        result_lines = []
        for element in elements:
            # Determine content based on element type
            if element["type"] == "text":
                content_value = element.get("text_content", "")
            else:
                content_value = element.get("object_label", "")
            
            # Format the element string
            element_str = (f"icon {element['code']}: {{"
                          f"'type': '{element['type']}', "
                          f"'centerX': {element['center_x']}, "
                          f"'centerY': {element['center_y']}, "
                          f"'content': '{content_value}'}}")
            
            result_lines.append(element_str)
        
        # Join all lines
        result_string = "\n".join(result_lines)
        
        # Log processing time
        processing_time = time.time() - start_time
        logger.info(f"Screenshot processed for string output in {processing_time:.2f} seconds")
        
        # Use the custom JSON encoder to handle numpy types
        response_data = {
            "output": result_string,
            "image_url": image_url
        }
        
        return JSONResponse(content=response_data)
        
    except Exception as e:
        logger.error(f"Error processing screenshot from data URI: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))