import os

import openai
import streamlit as st

# Set the OpenAI API key
os.environ["OPENAI_API_KEY"] = "sk-proj-yCa8xuZojVBCMM0r7q3hT3BlbkFJwjPh2CN3opBzjby3hg2J"
openai.api_key = os.getenv("OPENAI_API_KEY")

# Define the prompt template
def generate_poem(theme):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a poet."},
            {"role": "user", "content": f"Write a poem about {theme}"},
        ],
        max_tokens=150,
        temperature=0.7,
    )
    return response['choices'][0]['message']['content'].strip()

# Function to chain prompts
def chain_prompts(initial_theme):
    initial_output = generate_poem(initial_theme)
    refined_prompt = f"Expand on this idea: {initial_output}"
    refined_output = generate_poem(refined_prompt)
    return refined_output

# Streamlit app
st.title("Poetry Generator")
theme = st.text_input("Enter a theme or prompt for your poem:")

# Generate poem button
if st.button("Generate Poem"):
    poem = chain_prompts(theme)
    st.write(poem)