import base64
import json
import os
from glob import glob

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

CLASS_LABELS = ["Egyptian Mau", "leonberger", "samoyed"]

MODEL_REPO = "vasanthi8134/oxford-pets-3class-vit"
CLIP_MODEL = "openai/clip-vit-base-patch32"
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

vit_classifier = pipeline(
    "image-classification",
    model=MODEL_REPO,
)

clip_classifier = pipeline(
    "zero-shot-image-classification",
    model=CLIP_MODEL,
)

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

def classify_with_openai(image_path):
    if openai_client is None:
        return {
            "error": "Missing OPENAI_API_KEY in Hugging Face Space Secrets."
        }

    prompt = (
        "Classify the pet in this image. "
        f"Choose exactly one label from this list: {CLASS_LABELS}. "
        'Return valid JSON with keys: "label", "confidence", "reasoning". '
        "Do not use markdown code fences. "
        "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}",
                    },
                ],
            }
        ],
    )

    text = response.output_text.strip()

    if text.startswith("```json"):
        text = text[len("```json"):].strip()
    if text.startswith("```"):
        text = text[len("```"):].strip()
    if text.endswith("```"):
        text = text[:-3].strip()

    try:
        return json.loads(text)
    except Exception:
        return {"raw_response": response.output_text}

def classify_pet(image_path):
    vit_results = vit_classifier(image_path)
    vit_output = {item["label"]: round(float(item["score"]), 4) for item in vit_results}

    clip_results = clip_classifier(image_path, candidate_labels=CLASS_LABELS)
    clip_output = {item["label"]: round(float(item["score"]), 4) for item in clip_results}

    openai_output = classify_with_openai(image_path)

    return {
        "your_model_vit": vit_output,
        "open_source_clip": clip_output,
        "closed_source_openai": openai_output,
    }

example_files = []
for ext in ["jpg", "jpeg", "png", "webp"]:
    example_files.extend(glob(f"example_images/*.{ext}"))
    example_files.extend(glob(f"example_images/*.{ext.upper()}"))

example_files = [[path] for path in sorted(example_files)]

iface = gr.Interface(
    fn=classify_pet,
    inputs=gr.Image(type="filepath", label="Upload pet image"),
    outputs=gr.JSON(label="Model comparison"),
    title="Pet Classification Comparison",
    description=(
        "Compare a fine-tuned ViT model, a zero-shot CLIP model, "
        "and an OpenAI vision model on 3 pet classes: "
        "Egyptian Mau, leonberger, samoyed."
    ),
    examples=example_files if example_files else None,
)

iface.launch()