import gradio as gr
import tensorflow as tf
import numpy as np
import os
import random
from tensorflow.keras.preprocessing import image

# -----------------------------
# Safe Model Loading
# -----------------------------
try:
    model = tf.keras.models.load_model("Fruit_Classification_Model.h5")
    print("Model loaded successfully")
except Exception as e:
    print("Model loading failed:", e)
    model = None

# -----------------------------
# Safe Dataset Loading
# -----------------------------
DATASET_PATH = "Fruit_Classification"

if os.path.exists(DATASET_PATH):
    class_names = sorted([
        folder for folder in os.listdir(DATASET_PATH)
        if os.path.isdir(os.path.join(DATASET_PATH, folder))
    ])
else:
    class_names = []

print("Classes:", class_names)


# -----------------------------
# Prediction Function
# -----------------------------
def classify_from_text(text_input):
    try:
        if model is None:
            return None, "Model not loaded properly."

        if not class_names:
            return None, "Dataset folder not found."

        if text_input is None or text_input.strip() == "":
            return None, "Please enter a fruit name."

        text_input = text_input.strip().capitalize()

        if text_input not in class_names:
            return None, f"Invalid fruit name.\nValid: {class_names}"

        folder = os.path.join(DATASET_PATH, text_input)

        images = [
            img for img in os.listdir(folder)
            if img.lower().endswith((".jpg", ".jpeg", ".png"))
        ]

        if len(images) == 0:
            return None, "No images found in folder."

        img_name = random.choice(images)
        img_path = os.path.join(folder, img_name)

        # Load and preprocess image
        img = image.load_img(img_path, target_size=(224, 224))
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0

        # Predict
        prediction = model.predict(img_array, verbose=0)
        predicted_index = np.argmax(prediction)
        predicted_class = class_names[predicted_index]
        confidence = float(np.max(prediction)) * 100

        result = f"Predicted: {predicted_class}\nConfidence: {confidence:.2f}%"

        return img, result

    except Exception as e:
        return None, f"Error occurred: {str(e)}"


# -----------------------------
# Gradio Interface
# -----------------------------
interface = gr.Interface(
    fn=classify_from_text,
    inputs=gr.Textbox(
        label="Enter Fruit Name",
        placeholder="Apple, Banana, Lemon, Orange"
    ),
    outputs=[
        gr.Image(label="Sample Image"),
        gr.Textbox(label="Prediction Result")
    ],
    title="CNN Fruit Classification System",
    description="Enter fruit name → CNN selects sample image → CNN predicts fruit"
)

# Launch normally (DO NOT use ssr_mode=False)
if __name__ == "__main__":
    interface.launch()