"""
Simple VQA using Hugging Face API with error handling
"""

import gradio as gr
import requests
from PIL import Image
import io
import os

HF_TOKEN = os.environ.get("HF_TOKEN", "YOUR_TOKEN_HERE")

def answer_question(image, question):
    if image is None:
        return "Please upload an image first."
    if not question.strip():
        return "Please ask a question."
    
    try:
        # Convert image to bytes
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        
        # API call with timeout
        headers = {
            "Authorization": f"Bearer {HF_TOKEN}",
            "Content-Type": "application/json"
        }
        
        # For VQA models, use this format
        payload = {
            "inputs": {
                "image": img_bytes.hex(),
                "question": question
            }
        }
        
        response = requests.post(
            "https://api-inference.huggingface.co/models/Salesforce/blip-vqa-base",
            headers=headers,
            json=payload,
            timeout=30  # 30 second timeout
        )
        
        if response.status_code == 200:
            result = response.json()
            return result.get("answer", "No answer found")
        elif response.status_code == 503:
            return "Model is loading. Please try again in a few seconds."
        else:
            return f"Error: {response.status_code} - {response.text[:100]}"
    
    except requests.exceptions.Timeout:
        return "Request timed out. Please try again."
    except requests.exceptions.ConnectionError:
        return "Connection error. Check your internet and try again."
    except Exception as e:
        return f"Error: {str(e)}"

# Create interface
with gr.Blocks(title="VQA") as demo:
    gr.Markdown("# 🖼️ Visual Question Answering")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image", type="pil", height=300)
            question_input = gr.Textbox(label="Question", placeholder="What's in this image?")
            submit_btn = gr.Button("Answer", variant="primary")
        with gr.Column():
            answer_output = gr.Textbox(label="Answer", lines=4, interactive=False)
    
    submit_btn.click(answer_question, [image_input, question_input], answer_output)

demo.launch(share=True)