src/ui/handlers.py

"""
src/ui/handlers.py — Gradio event-handler functions for the VLA pipeline.

These functions are wired to Gradio UI events (button clicks, etc.) and
are intentionally kept free of Gradio component construction so they can
be unit-tested independently.
"""

import json
import logging

import cv2
import numpy as np

from src.pipeline.vla_pipeline import VLAPipeline

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Shared pipeline instance (initialised once at import time)
# ---------------------------------------------------------------------------
_pipeline: VLAPipeline | None = None


def get_pipeline() -> VLAPipeline:
    """Return the shared VLAPipeline, initialising it on first call."""
    global _pipeline
    if _pipeline is None:
        logger.info("Initializing VLA Pipeline...")
        _pipeline = VLAPipeline()
        logger.info("VLA Pipeline ready.")
    return _pipeline


# ---------------------------------------------------------------------------
# Main Gradio handler
# ---------------------------------------------------------------------------
def process_pipeline(
    image: np.ndarray | None,
    command: str,
) -> tuple:
    """
    Run the VLA pipeline and return outputs for Gradio components.

    Args:
        image:   Input image as a NumPy array (from gr.Image).
        command: Natural-language instruction from the user.

    Returns:
        5-tuple:
            - detection_image (np.ndarray)
            - detections_json (str)
            - action_plan_md  (str)
            - action_image    (np.ndarray)
            - status_message  (str)
    """
    empty = _blank_image("Upload an image and\nenter a command to begin")

    if image is None:
        return empty, "[]", "*Please upload an image.*", empty, "⚠️ No image provided."

    if not command or not command.strip():
        return empty, "[]", "*Please enter a command.*", empty, "⚠️ No command entered."

    logger.info("Processing: '%s'", command)
    result = get_pipeline().run(image, command)

    detections_json = json.dumps(result.detections, indent=2) if result.detections else "[]"
    action_plan_md = _format_action_plan(result.action_plan, result.interpretation)

    # Status bar
    if result.error:
        status = f"⚠️ {result.error}"
    else:
        plan_status = result.action_plan.get("status", "unknown")
        emoji = {"success": "✅", "partial": "⚠️", "failed": "❌"}.get(plan_status, "❓")
        n = len(result.detections)
        steps = len(result.action_plan.get("steps", []))
        status = (
            f"{emoji} Pipeline complete — Plan status: **{plan_status}**"
            f" |  {n} object{'s' if n != 1 else ''} detected"
            f" |  {steps} action step{'s' if steps != 1 else ''} generated"
        )

    det_img = _ensure_rgb(result.detection_image, empty)
    act_img = _ensure_rgb(result.action_image, empty)
    return det_img, detections_json, action_plan_md, act_img, status


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _format_action_plan(action_plan: dict, interpretation: dict) -> str:
    """Format the action plan dict as a readable Markdown string."""
    if not action_plan:
        return "*No action plan generated.*"

    lines: list[str] = []

    if interpretation:
        lines += [
            "### 🧠 Command Interpretation",
            "| Field | Value |",
            "|-------|-------|",
            f"| **Action** | `{interpretation.get('action', 'N/A')}` |",
            f"| **Source** | `{interpretation.get('source', 'N/A')}` |",
            f"| **Destination** | `{interpretation.get('destination') or 'N/A'}` |",
            "",
        ]

    steps = action_plan.get("steps", [])
    if steps:
        lines.append("### 📋 Execution Steps")
        for i, step in enumerate(steps, 1):
            lines.append(f"**{i}.** {step}")
        lines.append("")

    status = action_plan.get("status", "unknown")
    emoji = {"success": "✅", "partial": "⚠️", "failed": "❌"}.get(status, "❓")
    lines += ["---", f"### {emoji} Status: `{status.upper()}`"]

    return "\n".join(lines)


def _ensure_rgb(image: np.ndarray | None, fallback: np.ndarray) -> np.ndarray:
    """Convert image to RGB format expected by Gradio."""
    if image is None:
        return fallback
    if len(image.shape) == 2:
        return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    if image.shape[2] == 4:
        return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
    return image


def _blank_image(text: str = "") -> np.ndarray:
    """Return a dark placeholder image with optional centred text."""
    img = np.full((400, 600, 3), 18, dtype=np.uint8)  # near-black
    if text:
        for i, line in enumerate(text.split("\n")):
            y = 185 + i * 34
            (w, _), _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 1)
            x = max(0, (600 - w) // 2)
            cv2.putText(img, line, (x, y), cv2.FONT_HERSHEY_SIMPLEX,
                        0.7, (100, 100, 100), 1, cv2.LINE_AA)
    return img