from ultralytics import YOLO
import cv2
import torch
import gradio as gr
from PIL import Image
import numpy as np
import openai
import base64
import io
import os

# Load YOLOv8 model (Ensure best.pt is in the same directory)
model_path = "best.pt"
if not os.path.exists(model_path):
    raise FileNotFoundError(f"Model file '{model_path}' not found. Upload it to the Space.")

model = YOLO(model_path)  # Load model
model.to("cpu")  # Ensure it runs on CPU

# OpenAI API Key (Set as environment variable in Hugging Face Space)
openai_api_key = os.getenv('openai_api_key',"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")
print(f"DEBUG: OpenAI API Key - {openai_api_key[:5]}**********")  # Prints only first 5 characters for security
openai.api_key = openai_api_key



def encode_image(image):
    """Convert image to base64 for OpenAI API"""
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode()

def get_car_details(image):
    """Send the image to OpenAI to get car make, model, and year."""
    try:
        base64_image = encode_image(image)  # Convert image to Base64

        openai_api_key = os.getenv('openai_api_key',"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")
        if not openai_api_key:
            raise ValueError("❌ OpenAI API key is missing! Set it in the environment.")

        client = openai.OpenAI(api_key=openai_api_key)  # ✅ Pass API key explicitly

        response = client.chat.completions.create(
            model="gpt-4o-mini",  # 🔹 Use GPT-4o Mini
            messages=[
                {"role": "system", "content": "You are an AI that identifies car make, model, and year from images."},
                {"role": "user", "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
                    {"type": "text", "text": "Identify the make, model, and year of this car."}
                ]}
            ],
            max_tokens=100
        )

        return response.choices[0].message.content.strip()

    except Exception as e:
        return f"Error getting car details: {str(e)}"




def predict(input_img):
    # Convert PIL image to OpenCV format
    image = np.array(input_img)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Run inference
    results = model(image)
    
    # Confidence threshold
    threshold = 0.4  # Show only if confidence > 0.55

    # Draw bounding boxes
    for result in results:
        for box in result.boxes:
            conf = float(box.conf[0])  # Confidence score
            
            if conf > threshold:  # ✅ Only consider if confidence > 0.55
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                label = f"Damage: {conf:.2f}"
                
                # Draw red bounding box
                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 3)
                cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.5, (0, 0, 255), 2)
    
    # Convert back to PIL format
    output_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    
    # Get car details from OpenAI
    car_details = get_car_details(input_img)
    
    return output_img, car_details


# Gradio interface
gradio_app = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload a car image", sources=['upload', 'webcam'], type="pil"),
    outputs=[
        gr.Image(label="Detected Damage"),
        gr.Textbox(label="Car Make, Model & Year")
    ],
    title="Car Damage & Identification",
    description="Upload an image of a car to detect damage and identify the make, model, and year."
)

if __name__ == "__main__":
    gradio_app.launch(share=True)