import gradio as gr
import json
import requests
from PIL import Image
import os

# --- Constants for the API ---
API_URL = "https://predict.ultralytics.com"
# It's recommended to use Hugging Face secrets for API keys
API_KEY = os.environ.get("ULTRALYTICS_API_KEY") 
MODEL_ID = "https://hub.ultralytics.com/models/RsLHnWMhiBPqy3iFZAgr"

def classify_image(image):
    """
    Takes an image, sends it to the Ultralytics API, and returns the
    classification name and confidence score.
    """
    # Convert the Gradio image (numpy array) to a file-like object
    image_pil = Image.fromarray(image)
    image_path = "temp_image.jpg"
    image_pil.save(image_path)

    headers = {"x-api-key": API_KEY}
    data = {
        "model": MODEL_ID,
        "imgsz": 640,
        "conf": 0.25,
        "iou": 0.45
    }

    try:
        with open(image_path, "rb") as f:
            response = requests.post(API_URL, headers=headers, data=data, files={"file": f})
        
        # Check for a successful response
        response.raise_for_status()

        # Parse the JSON response
        result_data = response.json()
        
        # Extract the relevant information
        # The structure is data -> images -> [0] -> results -> [0]
        if result_data.get("images") and result_data["images"][0].get("results"):
            top_result = result_data["images"][0]["results"][0]
            name = top_result.get("name")
            confidence = top_result.get("confidence")
            
            # Return the extracted values for the two output components
            return name, f"{confidence:.2f}"
        else:
            return "No classification found", "N/A"

    except requests.exceptions.RequestException as e:
        return f"API Error: {e}", ""
    except (KeyError, IndexError):
        return "Could not parse API response.", ""
    finally:
        # Clean up the temporary image file
        if os.path.exists(image_path):
            os.remove(image_path)

# --- Gradio Interface ---

# Define the input and output components
image_input = gr.Image(type="numpy", label="Upload an Image or Use Webcam")

# Define separate outputs for name and confidence
name_output = gr.Label(label="Classification")
confidence_output = gr.Textbox(label="Confidence Score")


# List of example images
example_images = [
    ["images/img1.jpg"],
    ["images/img2.jpg"],
    ["images/img3.jpg"],
    ["images/img4.jpg"],
    ["images/img5.jpg"],
]

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=image_input,
    outputs=[name_output, confidence_output],
    title="Tile Classification: Proof of Concept",
    description="Upload a picture or use your camera to classify an image using a pre-trained model. The model's prediction and confidence score will be displayed.",
    examples=example_images
)

# Launch the application
iface.launch()