import os
import gradio as gr

from langchain import PromptTemplate
from langchain.chains import LLMChain
from langchain.llms import OpenAI

eleven = gr.load("elevenlabs/tts", src="spaces")

openai_api_key = os.environ.get("OPENAI_API_KEY")

llm = OpenAI(temperature=0.9)

def split_text(text, max_length):
    chunks = []
    current_chunk = ''
    words = text.split()

    for word in words:
        if len(current_chunk) + len(word) <= max_length:
            current_chunk += ' ' + word
        else:
            chunks.append(current_chunk.strip())
            current_chunk = word

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks

def generate_story(text):
    """Generate a story using the langchain library and OpenAI's GPT-3 model."""
    prompt = PromptTemplate(
        input_variables=["text"],
        template=""" 
        You are a fun and seasoned storyteller. 
        Generate a short story for a 5 years old audience about {text}.
                 """
    )
    story = LLMChain(llm=llm, prompt=prompt)
    story_result = story.run(text=text)
    print(story_result)
    print("""
    —
    Cutting text in chunks
    —
    """)
    max_length = 250
    text_chunks = split_text(large_text, max_length)
    for chunk in text_chunks:
        print(chunk)
    
    return story_result

def app(text):
    story = generate_story(text)
    return story

with gr.Blocks() as demo:
    with gr.Column():
        text = gr.Textbox()
        submit_btn = gr.Button('Submit')
        audio = gr.Audio()
        story = gr.Textbox()

    submit_btn.click(fn=app, inputs=[text], outputs=[story])
        
demo.launch()