inference.py

"""
Inference module for LearningStudio Callout Detection wrapper.

This module:
1. Normalizes input to bytes (handles URLs, data URLs, raw base64)
2. Gets presigned S3 URL from API Gateway
3. Uploads image directly to S3 (bypasses API Gateway for large payloads)
4. Calls API Gateway to start detection job
5. Polls for completion
6. Transforms callouts to EMCO format
"""

import os
import base64
import time
import logging
from typing import Dict, Any, List, Optional, Tuple

import requests

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Environment variables (set in HF Inference Endpoint secrets)
API_GATEWAY_URL = os.environ.get("API_GATEWAY_URL", "")
API_KEY = os.environ.get("API_KEY", "")

# Polling configuration
MAX_WAIT_SECONDS = 900  # 15 minutes
POLL_INTERVAL_SECONDS = 5

# Default processing parameters (LearningStudio doesn't pass these)
DEFAULT_PARAMS = {
    "tiling": {"tile": 2048, "overlap": 0.30},
    "floodfill": {"erase_text": False, "min_fill_vs_text": 0.0},
    "preclean": {"denoise_sw": 8}
}


def normalize_to_bytes(image_input: str) -> Tuple[bytes, str]:
    """
    Normalize image input to bytes.

    Handles:
    - HTTP/HTTPS URLs: Downloads image
    - Data URLs (data:image/png;base64,...): Decodes base64
    - Raw base64: Decodes to bytes

    Args:
        image_input: Image URL, data URL, or base64 string

    Returns:
        Tuple of (image_bytes, filename)
    """
    # Check if it's a URL
    if image_input.startswith(("http://", "https://")):
        logger.info(f"Downloading image from URL: {image_input[:100]}...")
        response = requests.get(image_input, timeout=60)
        response.raise_for_status()

        # Try to get filename from URL
        from urllib.parse import urlparse
        parsed = urlparse(image_input)
        filename = os.path.basename(parsed.path) or "image.png"

        return response.content, filename

    # Check if it's a data URL
    if image_input.startswith("data:"):
        # Parse data URL: data:image/png;base64,<data>
        try:
            header, encoded = image_input.split(",", 1)
            # Extract extension from mime type
            mime_part = header.split(";")[0].replace("data:", "")
            ext = mime_part.split("/")[-1] if "/" in mime_part else "png"
            return base64.b64decode(encoded), f"image.{ext}"
        except ValueError:
            raise ValueError("Invalid data URL format")

    # Assume it's already base64
    try:
        return base64.b64decode(image_input), "image.png"
    except Exception as e:
        raise ValueError(f"Invalid base64 string: {e}")


def get_upload_url(filename: str = "image.png") -> Dict[str, str]:
    """
    Get presigned S3 URL for image upload.

    Args:
        filename: Original filename for the image

    Returns:
        Dict with job_id, upload_url, s3_url
    """
    if not API_GATEWAY_URL or not API_KEY:
        raise ValueError(
            "API_GATEWAY_URL and API_KEY must be set in environment variables. "
            "Configure these in your HF Inference Endpoint secrets."
        )

    url = f"{API_GATEWAY_URL.rstrip('/')}/upload-url"
    headers = {"x-api-key": API_KEY}
    params = {"filename": filename}

    logger.info(f"Getting upload URL from {url}")
    response = requests.get(url, headers=headers, params=params, timeout=30)
    response.raise_for_status()

    result = response.json()
    logger.info(f"Got upload URL for job_id={result.get('job_id')}")
    return result


def upload_to_s3(upload_url: str, image_bytes: bytes) -> None:
    """
    Upload image directly to S3 using presigned URL.

    Args:
        upload_url: Presigned PUT URL
        image_bytes: Image data to upload
    """
    logger.info(f"Uploading {len(image_bytes)} bytes to S3...")
    response = requests.put(
        upload_url,
        data=image_bytes,
        headers={"Content-Type": "image/png"},
        timeout=60
    )
    response.raise_for_status()
    logger.info("Upload complete")


def start_detection_job(job_id: str, s3_url: str, params: Optional[Dict] = None) -> str:
    """
    Start a detection job via API Gateway.

    Args:
        job_id: Job ID from get_upload_url
        s3_url: S3 URL from get_upload_url
        params: Optional processing parameters

    Returns:
        Job ID for polling
    """
    url = f"{API_GATEWAY_URL.rstrip('/')}/detect"
    headers = {
        "x-api-key": API_KEY,
        "Content-Type": "application/json"
    }
    payload = {
        "job_id": job_id,
        "s3_url": s3_url
    }
    if params:
        payload["params"] = params

    logger.info(f"Starting detection job {job_id}")
    response = requests.post(url, headers=headers, json=payload, timeout=30)
    response.raise_for_status()

    result = response.json()
    logger.info(f"Detection job started: {result.get('status')}")
    return job_id


def poll_for_completion(job_id: str) -> Dict[str, Any]:
    """
    Poll API Gateway for job completion.

    Args:
        job_id: Job ID to poll

    Returns:
        Final result with callouts
    """
    url = f"{API_GATEWAY_URL.rstrip('/')}/status/{job_id}"
    headers = {"x-api-key": API_KEY}

    elapsed = 0
    while elapsed < MAX_WAIT_SECONDS:
        logger.info(f"Polling job {job_id} (elapsed: {elapsed}s)")

        response = requests.get(url, headers=headers, timeout=30)
        response.raise_for_status()

        result = response.json()
        status = result.get("status")

        if status == "SUCCEEDED":
            logger.info(f"Job {job_id} completed successfully")
            return result

        if status in ("FAILED", "TIMED_OUT", "ABORTED"):
            error_msg = result.get("error", f"Job {status.lower()}")
            logger.error(f"Job {job_id} failed: {error_msg}")
            return {
                "status": status,
                "error": error_msg,
                "callouts": []
            }

        # Still running, wait and retry
        time.sleep(POLL_INTERVAL_SECONDS)
        elapsed += POLL_INTERVAL_SECONDS

    # Timeout
    logger.error(f"Job {job_id} timed out after {MAX_WAIT_SECONDS}s")
    return {
        "status": "TIMEOUT",
        "error": f"Timeout waiting for results after {MAX_WAIT_SECONDS}s",
        "callouts": []
    }


def transform_to_emco_format(
    callouts: List[Dict],
    image_base64: str,
    image_width: int = 0,
    image_height: int = 0
) -> Dict[str, Any]:
    """
    Transform callouts from Lambda format to EMCO format.

    Lambda format:
        {"bbox": [x, y, w, h], "score": 0.95, ...}  # xywh

    EMCO format:
        {"bbox": {"x1": x, "y1": y, "x2": x+w, "y2": y+h}, "confidence": 0.95, ...}  # xyxy

    Args:
        callouts: List of callouts from Lambda
        image_base64: Original image as base64
        image_width: Image width
        image_height: Image height

    Returns:
        EMCO-compatible response dict
    """
    predictions = []

    for i, callout in enumerate(callouts):
        bbox = callout.get("bbox", [0, 0, 0, 0])

        # Convert from [x, y, w, h] to {x1, y1, x2, y2}
        x, y, w, h = bbox[0], bbox[1], bbox[2], bbox[3]

        prediction = {
            "id": i + 1,
            "label": "callout",
            "class_id": 0,
            "confidence": callout.get("score", callout.get("confidence", 1.0)),
            "bbox": {
                "x1": int(x),
                "y1": int(y),
                "x2": int(x + w),
                "y2": int(y + h)
            }
        }

        # Include optional fields if present
        if "text" in callout:
            prediction["text"] = callout["text"]

        predictions.append(prediction)

    return {
        "predictions": predictions,
        "total_detections": len(predictions),
        "image": image_base64,
        "image_width": image_width,
        "image_height": image_height
    }


def inference(image_input: str, parameters: Optional[Dict] = None) -> Dict[str, Any]:
    """
    Run inference on an image.

    This is the main entry point for the HF wrapper.

    Flow:
    1. Normalize input to bytes
    2. Get presigned S3 URL
    3. Upload image directly to S3
    4. Start detection job (small JSON payload)
    5. Poll for completion
    6. Transform results to EMCO format

    Args:
        image_input: Image URL, data URL, or base64 string
        parameters: Optional processing parameters

    Returns:
        EMCO-compatible response with predictions
    """
    try:
        # 1. Normalize input to bytes
        logger.info("Normalizing input...")
        image_bytes, filename = normalize_to_bytes(image_input)

        # Keep base64 for response
        image_base64 = base64.b64encode(image_bytes).decode("utf-8")

        # 2. Get presigned upload URL
        logger.info("Getting upload URL...")
        upload_info = get_upload_url(filename)
        job_id = upload_info["job_id"]
        upload_url = upload_info["upload_url"]
        s3_url = upload_info["s3_url"]

        # 3. Upload image directly to S3
        logger.info("Uploading to S3...")
        upload_to_s3(upload_url, image_bytes)

        # 4. Start detection job (merge user params with defaults)
        logger.info("Starting detection job...")
        merged_params = {**DEFAULT_PARAMS, **(parameters or {})}
        start_detection_job(job_id, s3_url, merged_params)

        # 5. Poll for completion
        logger.info("Polling for completion...")
        result = poll_for_completion(job_id)

        # 6. Check for errors
        if result.get("status") in ("FAILED", "TIMED_OUT", "ABORTED", "TIMEOUT"):
            return {
                "error": result.get("error", "Unknown error"),
                "predictions": [],
                "total_detections": 0,
                "image": image_base64
            }

        # 7. Transform to EMCO format
        logger.info("Transforming results to EMCO format...")
        callouts = result.get("callouts", [])
        image_width = result.get("image_width", 0)
        image_height = result.get("image_height", 0)

        return transform_to_emco_format(
            callouts,
            image_base64,
            image_width,
            image_height
        )

    except requests.exceptions.RequestException as e:
        logger.error(f"Request error: {e}")
        return {
            "error": f"Request error: {str(e)}",
            "predictions": [],
            "total_detections": 0,
            "image": ""
        }
    except ValueError as e:
        logger.error(f"Validation error: {e}")
        return {
            "error": str(e),
            "predictions": [],
            "total_detections": 0,
            "image": ""
        }
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        return {
            "error": f"Unexpected error: {str(e)}",
            "predictions": [],
            "total_detections": 0,
            "image": ""
        }