from openai import OpenAI
import os
import streamlit as st
from PIL import Image
import requests
from dotenv import load_dotenv
from transformers import pipeline
from datasets import load_dataset
import soundfile as sf
import torch

# Set the Streamlit app theme
st.set_page_config(
    page_title="Image to Story Generator",
    layout="wide",
    initial_sidebar_state="expanded",
)

from components import animation_html_a, animation_html, con_color
st.markdown(con_color,unsafe_allow_html=True)

load_dotenv()
# Load your OpenAI API key and other environment variables 
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")

client = OpenAI(api_key=OPENAI_API_KEY)

header = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}
# Logo image
@st.cache_data
def logo():
    logo_image = Image.open("images/background.png")
    return logo_image
st.image(logo(), use_column_width=True)

# Centered layout
col1, col2, col3 = st.columns([1, 2.2, 1])
co1, co2, co3 = st.columns([1,1.5,1])
c1, c2, c3, c4, c5 = st.columns([1,1,0.35,1,1])
st.divider()

def imagettext(url):
    response = requests.post("https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large", headers = header, data= url)
    text = response.json()
    text = text[0]['generated_text']
    return text 

def story_generator(scenario):
    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "system", "content": f"""
    You're a Story Teller. You wanna generate a story based on simple narrative, the story should be less than 60 words and should not exceed it.
    Context: {scenario}`
    STORY: 
    """}],
        max_tokens=70,  
        temperature=1,
    )
    story = response.choices[0].message.content.strip()
    
    current_sentence = []
    word_count = 0
    scenario_words = story.split()
    
    for word in scenario_words:
        current_sentence.append(word)
        word_count += 1
        if word_count >= 60 and word_count <= 100 and word == ".":
            break 
    storyy = " ".join(current_sentence)
    return storyy

@st.cache_resource
def text2speech():
    synthesiser = pipeline("text-to-speech", "microsoft/speecht5_tts")

    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
    speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)

    return synthesiser, speaker_embedding


with co2:
    with st.container(border=True):
            uploaded_file = st.file_uploader(":black[Choose an image]", type=["png", "jpg", "jpeg"])
            if uploaded_file is not None:
                with st.container(border=True):
                    st.image(uploaded_file, caption="Scroll Down", use_column_width=True)
                
                with c3:
                        animation_html_container = st.markdown(animation_html, unsafe_allow_html=True)
                        image_bytes = uploaded_file.read()  
                        scenario = imagettext(image_bytes) 
    
                story = story_generator(scenario)
                animation_html_container.empty()
                st.subheader("Generated Story:")
                st.info(story)

                x1 , x2, x3, x4, x5 = st.columns([1,1,1,1,1])
                with x3:
                    gene = st.markdown("######   Generating...    ")
                # Text to audio
                with c3:
                    animation_html_con = st.markdown(animation_html_a, unsafe_allow_html=True)
                    synethsiser,  speaker_embedding = text2speech()
                    speech = synethsiser(story, forward_params={"speaker_embeddings": speaker_embedding})
                    sf.write("speech.wav", speech["audio"], samplerate=speech["sampling_rate"])
                    gene.empty()
                    animation_html_con.empty()

                # Display audio with a title
                st.subheader("Story Audio:")
                st.audio("speech.wav", format="audio/wav", start_time=0)
                
r1, r2, r3 = st.columns([1,1.6,1])
t1 , t2, t3, t4, t5 = st.columns([1,1,0.8,1,1])
colz1, colz2, colz3, colz4, colz5, colz6,colz7,colz8,colz9,colz10 = st.columns([1,1,1,1,1,1,2.2,1,1,1])

with r2:
    st.markdown("### ♾ Designed and Developed by Harish Kumar ♾")
with t3:
    st.markdown("##### Connect with Me!!!")

# Twitter
twitter_logo = Image.open("images/x.png")  
colz7.image(twitter_logo, width=50)
colz7.markdown("[Twitter](https://twitter.com/isharizh)")

# LinkedIn
linkedin_logo = Image.open("images/linkedin.png")  
colz5.image(linkedin_logo, width=50)
colz5.markdown("[LinkedIn](https://www.linkedin.com/in/harizh/)")

# GitHub
github_logo = Image.open("images/github.png")  
colz6.image(github_logo, width=50)
colz6.markdown("[GitHub](https://github.com/isharizh)")