import gradio as gr
from handler import EndpointHandler
import io

# Model loading
handler = EndpointHandler(".")

def predict(image):
    try:
        # 1. PIL Image ko Bytes me convert karein
        buf = io.BytesIO()
        image.save(buf, format='PNG')
        image_bytes = buf.getvalue()

        # 2. Handler ko DICTIONARY bhejein (Jaise wo .pop() kar sake)
        # Hum wahi format bhej rahe hain jo Hugging Face Inference API bhejti hai
        result = handler({"inputs": image_bytes})
        
        # 3. Result extract karein
        if isinstance(result, dict):
            return result.get("prediction", "No Prediction Found")
        return str(result)
        
    except Exception as e:
        return f"Backend Error: {str(e)}"

# Gradio Interface
demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs="text",
    title="Master Brain Captcha Solver"
)

demo.launch()