"""
AiGiS Safety Captioning — Hugging Face Space (Gradio UI).

Same pipeline as the FastAPI service in ../src/api.py, wrapped in a simple
Gradio interface for HF Spaces:

    image in -> YOLO (detector.py) -> structured detections + compliance
             -> Claude (vlm_caption.py) -> safety observation
             -> annotated image + caption shown in the UI

The YOLO weights are pulled once at startup from the Hugging Face model repo
below, so no local .pt file is needed in this Space.

Every Gradio app also exposes an automatic API (see the "Use via API" link at
the bottom of the Space page), so this doubles as a callable endpoint too.

Requires the ANTHROPIC_API_KEY secret to be set in the Space settings.
"""
from __future__ import annotations

import os

import cv2
import gradio as gr
import numpy as np
import spaces
from huggingface_hub import hf_hub_download

import vlm_caption
from detector import SafetyDetector

HF_REPO_ID = os.environ.get("AIGIS_HF_REPO_ID", "omarkashif/hockey-yolo-models")
HF_FILENAME = os.environ.get("AIGIS_HF_FILENAME", "aigis/yolov8s.pt")
DATA_YAML = os.environ.get("AIGIS_DATA_YAML", "sh17.yaml")
CONF = float(os.environ.get("AIGIS_CONF", "0.25"))

# Load the model on CPU at startup. Actual GPU attachment (ZeroGPU) only
# happens inside the @spaces.GPU-decorated function below, at inference time.
weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=HF_FILENAME)
detector = SafetyDetector(weights_path, DATA_YAML, conf=CONF)


@spaces.GPU(duration=60)
def run(image: np.ndarray, skip_vlm: bool):
    if image is None:
        return None, "Upload an image first.", {}

    # Gradio gives RGB; Ultralytics/OpenCV (and our detector) work in BGR.
    frame_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    det_result, annotated_bgr = detector.detect_and_annotate(frame_bgr, name="upload.jpg")
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)
    det_dict = det_result.to_dict()

    if skip_vlm:
        return annotated_rgb, "_(VLM captioning skipped)_", det_dict

    if "ANTHROPIC_API_KEY" not in os.environ:
        return annotated_rgb, "**Error:** ANTHROPIC_API_KEY is not set for this Space.", det_dict

    ok, buf = cv2.imencode(".jpg", frame_bgr)
    vlm_result = vlm_caption.caption_bytes(buf.tobytes(), "image/jpeg", det_dict)

    caption_md = (
        f"**Caption:** {vlm_result.get('caption', '(parse error)')}\n\n"
        f"**Safety observation:** {vlm_result.get('safety_observation', '')}\n\n"
        f"**Severity:** {vlm_result.get('severity', 'n/a')}"
    )
    return annotated_rgb, caption_md, det_dict


demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Image(type="numpy", label="Upload image"),
        gr.Checkbox(label="Skip VLM caption (detector only, no API cost)"),
    ],
    outputs=[
        gr.Image(label="Annotated detections"),
        gr.Markdown(label="Safety caption"),
        gr.JSON(label="Raw detections + compliance"),
    ],
    title="AiGiS Safety Captioning",
    description=(
        "Upload a workplace image. YOLO detects PPE compliance deterministically; "
        "Claude turns those findings into a grounded safety observation."
    ),
)

if __name__ == "__main__":
    demo.launch()