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

# Check if the model file exists before loading
MODEL_PATH = "civil_tool_classifier.h5"
if not os.path.exists(MODEL_PATH):
    print(f"Error: Model file '{MODEL_PATH}' not found.")
    model = None
else:
    try:
        model = tf.keras.models.load_model(MODEL_PATH)
    except Exception as e:
        print(f"Error loading model: {e}")
        model = None

# Check if the CSV file exists before loading
CSV_PATH = "tools.csv"
if not os.path.exists(CSV_PATH):
    print(f"Error: CSV file '{CSV_PATH}' not found.")
    tools_df = None
    CLASS_NAMES = []
else:
    try:
        tools_df = pd.read_csv(CSV_PATH)
        CLASS_NAMES = list(tools_df["tool_name"].values)
    except Exception as e:
        print(f"Error loading CSV file: {e}")
        tools_df = None
        CLASS_NAMES = []

def predict_tool(img):
    if model is None or tools_df is None:
        return {"tool": "Error", "usage": "Model/CSV not loaded", "safety": "N/A"}
    
    # Pre-process the image
    img = image.array_to_img(img).resize((224, 224))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0) / 255.0
    
    # Predict
    predictions = model.predict(img_array)
    class_idx = np.argmax(predictions[0])
    
    # Get the tool info from the DataFrame
    if class_idx < len(CLASS_NAMES):
        tool_name = CLASS_NAMES[class_idx]
        try:
            row = tools_df[tools_df["tool_name"] == tool_name].iloc[0]
            usage = row.get("tool_usage", "Usage not available.")
            safety = row.get("tool_safety", "Safety not available.")
        except Exception:
            usage = "Usage information not available."
            safety = "Safety guidance not available."
    else:
        tool_name = "Unknown Tool"
        usage = "Not found in the list."
        safety = "Not found in the list."

    # 👇 Return JSON instead of multiple textboxes
    return {"tool": tool_name, "usage": usage, "safety": safety}

# Gradio interface with JSON output
demo = gr.Interface(
    fn=predict_tool,
    inputs=gr.Image(type="numpy"),
    outputs=gr.JSON(),  # 👈 JSON output for API
    title="Civil Tool Classifier",
    description="Upload an image of a civil tool to get its details."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)