import os
import streamlit as st
from groq import Groq

# Set up Streamlit page configuration
st.set_page_config(
    page_title="Writing Assistant",
    page_icon="✍️",
    layout="wide",
)

# Streamlit UI
st.title("✍️ Writing Assistant for IELTS, DET, and TOEFL")

# Prompt user to input the API key (Secure way)
api_key = st.text_input(
    "Enter your Groq API Key:", 
    type="password", 
    help="Enter your Groq API key to access the feedback model."
)

# Verify API key
if not api_key:
    st.warning("Please enter your Groq API key to proceed.")
    st.stop()

# Initialize Groq client
try:
    client = Groq(api_key=api_key)
except Exception as e:
    st.error(f"Failed to initialize Groq client: {e}")
    st.stop()

# Essay plans (30, 45, and 60 days)
plans = {
    "30 Days": ["Essay Topic 1", "Essay Topic 2", "Essay Topic 3"],
    "45 Days": ["Essay Topic 1", "Essay Topic 2", "Essay Topic 3", "Essay Topic 4"],
    "60 Days": ["Essay Topic 1", "Essay Topic 2", "Essay Topic 3", "Essay Topic 4", "Essay Topic 5"]
}

# Select plan
selected_plan = st.selectbox("Choose your plan:", list(plans.keys()), index=0)
topics = plans[selected_plan]

# Day selector dropdown
selected_day = st.slider("Select your current day:", 1, len(topics))
current_topic = topics[selected_day - 1]

# Show the current essay topic
st.subheader(f"📝 Today's Essay Topic: {current_topic}")

# Show upcoming topics
st.markdown("#### Upcoming Topics:")
for i, topic in enumerate(topics[selected_day:], start=selected_day + 1):
    st.write(f"Day {i}: {topic}")

# Essay submission
st.markdown("### Write Your Essay Below:")
essay_input = st.text_area("Your essay:", placeholder="Write your essay here...", height=300)

# Submit button
if st.button("Submit for Feedback"):
    if not essay_input.strip():
        st.error("Please write your essay before submitting.")
    else:
        # Define the system prompt for the Groq model
        system_prompt = """
        You are an experienced academic English instructor. You must provide feedback as an English teacher would, giving concise but effective advice to improve the student's writing. 
        Focus on specific, actionable feedback without lengthy explanations.

        For each error or area of improvement:
        1. Identify grammar or vocabulary mistakes by saying: "Replace 'X' with 'Y' (Type: Grammar/Vocabulary)".
        2. For cohesion, sentence structure, or other improvements, suggest: "Rephrase 'X' as 'Y' for better flow".
        3. Provide concise suggestions based on English proficiency levels (A1 to C1).
        4. Do not generate vague or lengthy answers; avoid unnecessary technical jargon.
        5. Limit feedback to 4-5 actionable points, focusing on the most important improvements.

        Essay Analysis Areas: Grammar, cohesion, sentence structure, vocabulary, and effective use of simple, compound, and complex sentences.
        """

        # Send the essay to the Groq model
        try:
            response = client.chat.completions.create(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": essay_input}
                ],
                model="llama3-8b-8192",
            )

            # Display feedback
            feedback = response.choices[0].message.content
            st.success("✅ Feedback Received!")
            st.markdown("### Model Feedback:")
            st.write(feedback)

        except Exception as e:
            st.error(f"Failed to get feedback from the model: {e}")