import streamlit as st
import requests
from PIL import Image
from io import BytesIO
from gradio_client import Client
import json

def generate(outline, characters, settings, num_images):
    # Create a dynamic prompt with multiple image prompts based on the number of images
    prompt = f"""Hello! I would like to request a 4-paragraph story (700 words per paragraph) and a corresponding image prompt for each paragraph in JSON format described later in the prompt with the following detailed outline:\n\n{outline}\n\nCharacters: {characters}\n\nSettings: {settings}\n\nPlease generate the story with the following detailed JSON format:\n\np1, p2, p3, p4: Keys for story paragraphs; title: Key for story title; {', '.join([f'prompt{i+1}' for i in range(num_images)])}: Keys for image prompts for each paragraph (or more if additional images are needed). Please do not include any other text in the output. Only JSON is needed or it will break the system and cause loss. Please don't say 'Here is the requested output in JSON format' or anything similar. Thank you."""

    client = Client("Be-Bo/llama-3-chatbot_70b")
    hikaye = client.predict(
        message=prompt,
        api_name="/chat"
    )

    print("Debug: Generated hikaye:", hikaye)
    
    return hikaye

def cover(prompts):
    api_key = st.secrets['apikey']
    model = "mann-e/Mann-E_Turbo"
    headers = {"Authorization": f"Bearer {api_key}"}
    api_url = f"https://api-inference.huggingface.co/models/{model}"

    images = []
    for prompt in prompts:
        data = {"inputs": prompt}
        response = requests.post(api_url, headers=headers, json=data)

        if 'image' in response.headers.get('content-type', '').lower():
            image = Image.open(BytesIO(response.content))
            images.append(image)
        else:
            st.error(f"Failed to fetch image for prompt: {prompt}")
            images.append(None)

    return images

def parse_story_response(response, num_images):
    print("Debug: Raw response from API:", response)

    if not response:
        print("Debug: Response is empty or None.")
        return None, None, None, None, []

    title = response.get('title', '')
    p1 = response.get('p1', '')
    p2 = response.get('p2', '')
    p3 = response.get('p3', '')
    p4 = response.get('p4', '')

    prompts = [response.get(f'prompt{i+1}', '') for i in range(num_images)]
    
    return title, p1, p2, p3, p4, prompts

# Streamlit UI
st.title('Story AI By Ozi')

characters = st.text_area(label="Characters")
outline = st.text_area(label="Story Outline")
settings = st.text_area(label="Setting")

# Number chooser for images (up to 5)
num_images = st.slider('Number of Images to Generate', 1, 5, 4)

if st.button(label="Generate"):
    with st.spinner('Generating story and cover images...'):
        hikaye = generate(outline, characters, settings, num_images)
        print("Debug: Story generation response:", hikaye)

        if hikaye:
            try:
                hikaye_json = json.loads(hikaye)
            except json.JSONDecodeError as e:
                st.error(f"Failed to parse JSON response: {e}")
                st.stop()

            title, p1, p2, p3, p4, prompts = parse_story_response(hikaye_json, num_images)

            if title and p1 and p2 and p3 and p4:
                st.markdown(f'### {title}')

                # Generate and display images
                images = cover(prompts)
                for i, image in enumerate(images):
                    if image:
                        st.image(image, caption=f"Image {i+1}: {prompts[i]}")
                    else:
                        st.error(f"Failed to generate image {i+1}.")

                # Display paragraphs
                st.markdown(f'''
                {p1}
                {p2}
                {p3}
                {p4}
                ''')

            else:
                st.error("Failed to generate or parse story.")
        else:
            st.error("No story data received.")