# app.py - Hugging Face ready
import os
import re
import json
import runpy
import nbformat
from nbconvert import PythonExporter

# Prefer dependencies via requirements.txt. Small one-off installs if needed:
# os.system("pip install seaborn --quiet")

import torch
import gradio as gr
from PIL import Image
from datasets import load_dataset

# Use the clean module instead of converting the notebook
from yolo_module import detect_and_classify

# --- Load class names (optional, cached to file) ---
try:
    ds = load_dataset("tanganke/stanford_cars")
    class_names = ds["train"].features["label"].names
    with open("class_names.json", "w", encoding="utf-8") as f:
        json.dump(class_names, f)
    print(f"✅ Loaded {len(class_names)} class names")
except Exception as e:
    print("⚠️ Could not load dataset class names:", e)
    class_names = None

# --- Gradio UI ---
def gradio_interface(image):
    if image is None:
        return "Please upload an image."
    temp_path = "temp_image.png"
    image.save(temp_path)
    try:
        results = detect_and_classify(temp_path)
    except Exception as e:
        return f"❌ Error running YOLO pipeline: {e}"
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)

    if not results:
        return "No cars detected."

    lines = [f"Cars detected: {len(results)}"]
    for i, item in enumerate(results, start=1):
        # item may be (crop, pred_idx, color) or (crop, pred_idx, color, conf)
        if isinstance(item, (list, tuple)) and len(item) == 4:
            _, pred, color, conf = item
        elif isinstance(item, (list, tuple)) and len(item) >= 3:
            _, pred, color = item[:3]
            conf = None
        else:
            lines.append(f"Car {i}: {item}")
            continue

        if isinstance(pred, int) and class_names and 0 <= pred < len(class_names):
            name = class_names[pred]
        else:
            name = str(pred)

        if conf is not None:
            lines.append(f"Car {i}: {color} {name} ({conf*100:.1f}% confident)")
        else:
            lines.append(f"Car {i}: {color} {name}")

    return "\n".join(lines)

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Detection & Classification Results"),
    title="Car Detector + Classifier (YOLO)",
    description="(Status: Beta) To Test, Upload any car image and get its color, model, and confidence score."
)

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