import streamlit as st
import os
from dotenv import load_dotenv
from groq import Groq

load_dotenv()

# Initialize Groq client
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Directory for poems
poems_directory = "./poems"
os.makedirs(poems_directory, exist_ok=True)

def generate_poems(num_poems, length_per_poem, output_dir, instructions):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for i in range(num_poems):
        chat_completion = client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a talented poet."},
                {"role": "user", "content": f"{instructions}\n Make sure all the poems have the same length of {length_per_poem} tokens."}
            ],
            model="llama3-8b-8192",
            temperature=0.7,
            max_tokens=length_per_poem,
            top_p=1,
            stop=None,
            stream=False,
        )
        poem = chat_completion.choices[0].message.content.strip()
        with open(os.path.join(output_dir, f'poem_{i + 1}.txt'), 'w') as f:
            f.write(poem)

def generate_poems_page():
    st.header("Generate Poems")

    # Input fields for poem generation parameters
    num_poems = st.number_input("Number of poems to generate", min_value=1, max_value=100, value=5)
    length_per_poem = st.number_input("Length of each poem in tokens", min_value=50, max_value=1000, value=150)
    instructions = st.text_area("Custom instructions", "You are a talented poet.")

    if st.button("Generate Poems"):
        generate_poems(num_poems, length_per_poem, poems_directory, instructions)
        st.success(f"Generated {num_poems} poems and saved to {poems_directory}")