import gradio as gr
from main import app, MODELS, SUPPORTED_CROPS

# Safe import for local fallback
try:
    import spaces
except ImportError:
    class MockSpaces:
        def GPU(self, func):
            return func
    spaces = MockSpaces()

# Decorating the top-level Gradio function allows ZeroGPU to detect it on startup
@spaces.GPU
def gradio_predict(crop, image):
    if crop not in MODELS:
        return {"error": f"Model for '{crop}' not loaded"}
    return MODELS[crop].predict(image)["all_scores"]

demo = gr.Interface(
    fn=gradio_predict,
    inputs=[gr.Dropdown(SUPPORTED_CROPS, label="Crop"), gr.Image(type="pil")],
    outputs=gr.Label(num_top_classes=5),
)

app = gr.mount_gradio_app(app, demo, path="/")