import os
import streamlit as st
import threading
import asyncio
import requests
import io
from PIL import Image
from langchain.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field, validator
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env file

GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
API_KEY = os.environ["API_KEY"]

# Define the Story model for structured output
class Story(BaseModel):
    title: str = Field(description="A captivating title for the story.")
    characters: list[str] = Field(
        description="""Important:no json format. Six elements mandatory, each formatted as: 
                       "[Character Name], [comma-separated adjectives], cartoon, style africa, painting".
                       Describe each character's appearance in detail. Be creative!"""
    )
    scenes: list[str] = Field(
        description="""Important:no json format.no json format. Six elements mandatory, each a string describing a character's action.very important:use charaters description apperance, use only action verbs.
                      Each scene must follow the previous one chronologically, creating a complete narrative when combined.
                      Develop your story by detailing what each character DOES in each scene.Instead to use only name of characters to write this part, use name and this key word 'painting bening style mushgot 'as description appearence, it's very import to do loke that.if it's a new characters in the story, instead use his name use his name and add the keyword 'painting benin style'it is mandatory.Use your imagination!"""
    )
    metadonne: list[str] = Field(
        description="""Important: no json format.Six elements mandatory, each a concise one-sentence description of the corresponding scene in the 'scenes' field.
                      Explain the action taking place in each scene. Come up with your own unique descriptions!"""
    )


# Function to generate images using Stable Diffusion XL
def generate_image(scene, width=300):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    headers = {"Authorization": f"Bearer {API_KEY}"}
    negative_prompt = "ugly, blurry, low-resolution, deformed, mutated, disfigured, missing limbs, disjointed, distorted, deformed, unnatural"

    payload = {
        "inputs": scene,
        "guidance_scale": 0.8,
        "num_inference_steps": 8,  # Adjust for different image quality/speed
        "eta": 0.5,
        "seed": 46,  # Optionally provide a different seed for variety
        "negative_prompt": negative_prompt,
        "width": width  # Adjust the width for display
    }

    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        return image
    else:
        st.error(f"Image generation failed: {response.text}")
        return None


# Function to run Google GenAI asynchronously 
def run_google_genai_async():
    global response
    llm = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_API_KEY)

    system = "All instructions must be follow is very important, all story related to african culture and history is mandatory.You are a storyteller who specializes in creating educational tales about African culture. Your mission is to craft a narrative that teaches African children about their rich heritage. Your story is based on real events from the past, incorporating historical references, myths, and legends. story size is short length. Your narrative will be presented in six panels.Very important, For each panel, you will provide: A description of the characters, using precise and unique descriptions each time, ending with the keywords 'high quality', 'watercolor painting', 'painting Benin style', and 'mugshot', 'cartoon africa style' in the scenes or characters is mandatory.For description, using only words or groups of words separated by commas, without sentences. Each sentence in the panel's text should start with the character's name, and each sentence should be no longer than two small sentences. Each story has only three characters. Your story must always revolve around African legends and kingdoms, splitting the scenario into six parts. Be creative in each story"

    st.title("Storytelling with AI")
    title = st.text_input("Discover a new story on africa, tape a topic !")

    story_query = system + title
    parser = PydanticOutputParser(pydantic_object=Story)
    prompt = PromptTemplate(
        template="Answer the user query.\n{format_instructions}\n{query}\n",
        input_variables=["query"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )

    chain = prompt | llm | parser

    if title:
        response = chain.invoke({"query": story_query})
    else:
        response = None  # Set response to None if no title is provided


# Main Streamlit function
def main():
    global response  # Declare response as global to access it from run_google_genai_async
    response = None  # Initialize response

    thread = threading.Thread(target=run_google_genai_async)
    thread.start()

    # Streamlit UI - Display story elements and generated images
    if response:
        st.write(response)

        st.title("Images générées avec métadonnées dans une grille 2x3")
        for i in range(0, len(response.scenes), 2):
            col1, col2 = st.columns(2)
            col1.write(f"**Scène {i+1}:** {response.metadonne[i]}")
            scene_image = generate_image(response.scenes[i])
            if scene_image:
                col1.image(scene_image, caption=f"Image de la scène {i+1}", width=300)
            
            if i+1 < len(response.scenes):
                col2.write(f"**Scène {i+2}:** {response.metadonne[i+1]}")
                scene_image = generate_image(response.scenes[i+1])
                if scene_image:
                    col2.image(scene_image, caption=f"Image de la scène {i+2}", width=300)

if __name__ == "__main__":
    main()