import gradio as gr
from huggingface_hub import InferenceClient
from io import BytesIO
from os import path, unlink, getenv
from PIL.Image import Image, open as open_image
import pandas as pd
from pandas import DataFrame
import requests
from utils import save_image_to_temp_file


def image_classification(client: InferenceClient, image_url: str | None, image: Image | None) -> tuple[Image | None, DataFrame]:
    temp_file_path = None
    try:
        if image is not None and image_url and image_url.strip():
            raise gr.Error("Both an image URL and an uploaded image were provided. Please provide only one or the other.")
        elif image is not None:
            temp_file_path = save_image_to_temp_file(image)
            classifications = client.image_classification(temp_file_path, model=getenv("IMAGE_CLASSIFICATION_MODEL"))
            image = None
        elif image_url and image_url.strip():
            try:
                response = requests.get(image_url, timeout=int(getenv("REQUEST_TIMEOUT")))
                response.raise_for_status()
                image = open_image(BytesIO(response.content))
                temp_file_path = save_image_to_temp_file(image)
                classifications = client.image_classification(temp_file_path, model=getenv("IMAGE_CLASSIFICATION_MODEL"))
            except Exception as e:
                raise gr.Error(f"Failed to fetch image from URL: {str(e)}")
        else:
            raise gr.Error("Please either provide an image URL or upload an image.")
        df = pd.DataFrame({
                              "Label": classification.label,
                              "Probability": f"{classification.score:.2%}"
                          }
                          for classification
                          in classifications)
        return image, df
    finally:
        # Clean up temporary file.
        if temp_file_path and path.exists(temp_file_path):
            try:
                unlink(temp_file_path)
            except Exception:
                pass # Ignore clean-up errors.