import base64
import json
import os

import gradio as gr
from openai import OpenAI
from transformers import pipeline

OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None

# Load models
vit_classifier = pipeline("image-classification", model="durovali/vit-motorcycle")
clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")

labels_motorcycle = ["bmw", "honda", "kawasaki", "suzuki", "triumph", "yamaha"]


def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def classify_with_openai(image_path):
    if openai_client is None:
        return {
            "error": "Missing OPENAI_API_KEY. Add it to your environment or .env file."
        }

    prompt = (
        "Classify the motorcycle brand in this image. Choose the best matching label from this list: "
        f"{', '.join(labels_motorcycle)}. "
        "Return valid JSON with exactly these keys: "
        "label, confidence, reasoning. "
        "The confidence must be a number between 0 and 1."
    )

    base64_image = encode_image(image_path)
    response = openai_client.responses.create(
        model=OPENAI_MODEL,
        input=[
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": prompt},
                    {
                        "type": "input_image",
                        "image_url": f"data:image/jpeg;base64,{base64_image}",
                    },
                ],
            }
        ],
    )

    try:
        parsed_response = json.loads(response.output_text)
    except json.JSONDecodeError:
        parsed_response = {
            "raw_response": response.output_text,
            "warning": "OpenAI response was not valid JSON.",
        }

    return parsed_response


def classify_motorcycle(image):
    vit_results = vit_classifier(image)
    vit_output = {result['label']: result['score'] for result in vit_results}

    clip_results = clip_detector(image, candidate_labels=labels_motorcycle)
    clip_output = {result['label']: result['score'] for result in clip_results}

    openai_output = classify_with_openai(image)

    return {
        "ViT Classification (fine-tuned)": vit_output,
        "CLIP Zero-Shot Classification": clip_output,
        "OpenAI Vision Classification": openai_output,
    }


example_images = [
    ["example_images/bmw.jpg"],
    ["example_images/honda.jpg"],
    ["example_images/kawasaki.jpg"],
    ["example_images/triumph.jpg"],
    ["example_images/yamaha.jpg"],
]

iface = gr.Interface(
    fn=classify_motorcycle,
    inputs=gr.Image(type="filepath"),
    outputs=gr.JSON(),
    title="🏍️ Motorcycle Brand Classification",
    description=(
        "Upload a motorcycle image and compare predictions from:\n"
        "- A fine-tuned ViT model trained on motorcycle brand images\n"
        "- Zero-shot CLIP (openai/clip-vit-large-patch14)\n"
        "- OpenAI GPT-4.1 Vision\n\n"
        "Supported brands: BMW, Honda, Kawasaki, Suzuki, Triumph, Yamaha"
    ),
    examples=example_images
)

iface.launch()