import gradio as gr
from gradio_client import Client
from PIL import Image
import requests
from io import BytesIO

# Connect to remote Gradio apps
grounded_client = Client("admin08077/veo")
image_client = Client("admin08077/picgenai")
chat_client = Client("admin08077/genai")

def generate_book(topic):
    # 1. Grounded response
    try:
        grounded_response = grounded_client.predict(
            prompt=topic,
            api_name="/generate_grounded_response"
        )
    except Exception as e:
        grounded_response = f"Error: {e}"

    # 2. Image and caption
    try:
        caption, image = image_client.predict(
            prompt=topic,
            api_name="/generate_image_and_caption"
        )

        # Image might come back as file path or PIL.Image — normalize it
        if isinstance(image, str):  # handle file path
            if image.startswith("http"):
                response = requests.get(image)
                image = Image.open(BytesIO(response.content))
            else:
                image = Image.open(image)

    except Exception as e:
        caption = f"Error generating image: {e}"
        image = None

    # 3. Gemini story
    try:
        narrative = chat_client.predict(
            user_input=f"Write a short story about: {topic}",
            api_name="/chat_with_gemini"
        )
    except Exception as e:
        narrative = f"Error: {e}"

    # Return all parts
    return grounded_response, caption, image, narrative

# Gradio UI
with gr.Blocks(title="📘 AI-Generated Mini Book") as demo:
    gr.Markdown("## 📘 Generate a Mini Illustrated Book with Gemini + Gradio Clients")

    with gr.Row():
        topic_input = gr.Textbox(label="Book Topic", placeholder="e.g. Who won Euro 2024?")
        generate_button = gr.Button("Generate Book")

    grounded_output = gr.Textbox(label="🌍 Grounded Info", lines=3)
    caption_output = gr.Textbox(label="🖼️ Image Caption")
    image_output = gr.Image(label="📷 Illustration", type="pil")
    story_output = gr.Textbox(label="📖 Story", lines=8)

    generate_button.click(
        fn=generate_book,
        inputs=topic_input,
        outputs=[grounded_output, caption_output, image_output, story_output]
    )

demo.launch()