import gradio as gr
from PIL import Image
from typing import Tuple, Any, Dict, List

# Lazy imports to speed cold start
clf = None
yolo_severity = None

def _load_models():
    global clf, yolo_severity
    if clf is None:
        from transformers import pipeline
        # Image classification (damage types)
        clf = pipeline("image-classification", model="beingamit99/car_damage_detection")
    if yolo_severity is None:
        from ultralytics import YOLO
        # YOLOv8 severity detector (Light/Moderate/Severe)
        yolo_severity = YOLO("nezahatkorkmaz/car-damage-level-detection-yolov8")

def analyze(img: Image.Image) -> Tuple[Dict[str, Any], Image.Image, Any]:
    """
    Returns:
      - JSON summary (gate + top labels + detections)
      - Overlay image with boxes
      - Raw YOLO JSON (string or dict)
    """
    _load_models()

    # --- Gate using classifier ---
    # If model has a 'no_damage' label use it; otherwise treat max score < 0.5 as "no damage"
    preds = sorted(clf(img), key=lambda x: x["score"], reverse=True)
    top = preds[0] if preds else {"label": "unknown", "score": 0.0}
    label_lower = top["label"].lower()
    if "no" in label_lower and "damage" in label_lower:
        gate = False
    else:
        gate = top["score"] >= 0.5

    if not gate:
        return {"gate": "No visible damage", "classification_top": top}, img, {"detections": []}

    # --- Top-3 labels for type ---
    top3 = [{"label": p["label"], "score": float(p["score"])} for p in preds[:3]]

    # --- YOLO severity boxes ---
    yres = yolo_severity.predict(img)
    result = yres[0]
    plotted = result.plot()  # numpy array with drawn boxes

    dets = []
    if result.boxes is not None and len(result.boxes) > 0:
        # class names if available
        names = result.names if hasattr(result, "names") else {}
        for i in range(len(result.boxes)):
            b = result.boxes[i]
            xyxy = b.xyxy[0].tolist()
            conf = float(b.conf[0].item())
            cls_id = int(b.cls[0].item())
            cls_name = names.get(cls_id, str(cls_id))
            dets.append({
                "bbox_xyxy": [float(x) for x in xyxy],
                "confidence": conf,
                "class_id": cls_id,
                "class_name": cls_name
            })

    summary = {"gate": "Damaged", "classification_top3": top3, "detections": dets}
    try:
        raw_json = result.tojson()  # string
    except Exception:
        raw_json = {"error": "tojson failed"}

    from PIL import Image as _Image
    return summary, _Image.fromarray(plotted), raw_json

demo = gr.Interface(
    fn=analyze,
    inputs=gr.Image(type="pil", label="Upload a car photo"),
    outputs=[
        gr.JSON(label="Results (gate + top labels + detections)"),
        gr.Image(label="Detections Overlay"),
        gr.JSON(label="Raw YOLO JSON")
    ],
    title="Car Damage Inspector",
    description=(
        "Fast, open-source car damage analysis.\n"
        "- Step 1: Classify damage type (ViT).\n"
        "- Step 2: Detect severity with YOLOv8 (boxes).\n"
        "Models: beingamit99/car_damage_detection, nezahatkorkmaz/car-damage-level-detection-yolov8."
    ),
)

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