import gradio as gr
from PIL import Image
import os
import sys
import logging

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

logger.info("Starting UK Public Transport Assistant app")

# Add error handling for module imports
try:
    from llava_inference import LLaVAHelper
    logger.info("Successfully imported LLaVAHelper")
except Exception as e:
    logger.error(f"Failed to import LLaVAHelper: {e}")
    logger.error("Stack trace:", exc_info=True)
    model = None
else:
    # Initialize model
    try:
        logger.info("Initializing LLaVA model...")
        model = LLaVAHelper()
        logger.info("LLaVA model initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize LLaVA model: {e}")
        logger.error("Stack trace:", exc_info=True)
        model = None

def answer_question(image, question):
    if model is None:
        return "Model initialization failed. Please check server logs."
    
    if image is None or question.strip() == "":
        return "Please upload an image and enter a question."
    
    try:
        return model.generate_answer(image, question)
    except Exception as e:
        return f"Error processing request: {str(e)}"

# Create examples directory if it doesn't exist
os.makedirs("assets", exist_ok=True)

demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Image(type="pil", label="Upload Public Transport Signage"),
        gr.Textbox(label="Ask a question (e.g., 'When is the next train to London?')")
    ],
    outputs=gr.Textbox(label="Answer"),
    title="UK Public Transport Assistant",
    description="Upload an image of UK public transport signage (like train timetables or metro maps), and ask a question related to it. Powered by LLaVA-1.5.",
    # Fix examples format - either provide valid examples or set to None
    examples=None  # Remove examples to avoid format errors
)

if __name__ == "__main__":
    demo.launch(share=True)  # Added share=True to make it accessible on a public URL