import streamlit as st
from openai import OpenAI
import os

# 1. Configuration
st.set_page_config(page_title="High School Writing Coach", layout="wide")

# Get the API key from Environment Variables (HF Spaces) OR st.secrets (Local)
api_key = os.environ.get("OPENAI_API_KEY")

if not api_key:
    try:
        # Fallback for local development if using secrets.toml
        api_key = st.secrets["OPENAI_API_KEY"]
    except (FileNotFoundError, KeyError):
        st.error("OpenAI API Key not found. Please add 'OPENAI_API_KEY' to your Hugging Face Space secrets.")
        st.stop()

# Initialize client with the found key
client = OpenAI(api_key=api_key)

# 2. System Prompt
# This enforces the "No Rewrite" rule.
# --- 1. DEFINE THE FEW-SHOT EXAMPLES FIRST ---
FEW_SHOT_EXAMPLES = """
<example_1>
USER_INPUT: "The Great Gatsby shows that money is bad. Gatsby has a lot of money but he is sad."
USER_GOAL: "Argument Analysis"
AI_RESPONSE:
**Critique:** Your point is valid, but the sentence 
structure is a bit simple and relies on "telling" 
rather than "showing." You are using two independent 
clauses joined loosely.

**Lesson - Complex Sentences:** Try using a subordinating 
conjunction (like "although," "while," or "despite") to 
show the relationship between wealth and happiness.

**Structural Example (Topic: Sports):** Instead of saying 
"The team won the game but they played bad," a stronger writer 
might say: *"Despite securing a victory on the scoreboard, the 
team's defensive breakdown revealed fundamental weaknesses in 
their strategy."*

**Try it:** Can you rewrite your sentence about Gatsby using the word "Although" or "Despite"?
</example_1>

<example_2>
USER_INPUT: "I think that climate change is a big problem because it makes the weather hot."
USER_GOAL: "Vocabulary & Tone"
AI_RESPONSE:
**Critique:** The phrase "big problem" is vague, 
and "I think" is unnecessary. Academic writing 
requires precision.

**Lesson - Precise Vocabulary:** Replace general 
words with specific terms that describe the 
*scale* or *nature* of the problem.

**Structural Example (Topic: Cooking):** Instead 
of saying "I think the soup was bad because it was 
too salty," a critic would write: *"The broth's 
overwhelming salinity completely masked the delicate 
flavors of the vegetables."*

**Try it:** Look at your sentence. How can you replace 
"big problem" with a word that describes *how* climate 
change affects the planet?
</example_2>
"""

# --- 2. DEFINE THE MASTER SYSTEM PROMPT ---
# We inject the few-shot examples at the end.
SYSTEM_PROMPT = f"""
You are an expert Writing Coach for high school students. 
Your goal is to teach writing mechanics, logic, and rhetoric without rewriting the student's essay for them.

CORE RULES:
1. **ABSOLUTE PROHIBITION:** DO NOT rewrite the student's text. If they ask "Can you fix this?" or "Rewrite it for me", you must REFUSE and ask them to try applying the lesson themselves.
2. If you see a grammatical error, quote the sentence and explain the grammar rule they broke.
3. Structure your feedback in Markdown with clear headings: "General Feedback", "Strengths", and "Areas for Improvement".
4. Be encouraging but rigorous. Treat them like smart young adults.

SOCRATIC INSTRUCTIONS (Use when critiquing logic/argument):
1. Do not give the answer.
2. Ask a question that exposes the gap in the student's reasoning.
3. Use the "Counter-Factual" technique: "If X were true, wouldn't Y also happen?"
4. Use the "Perspective Shift" technique: "How would a French soldier in 1812 respond to this claim?"

INSTRUCTIONS FOR EXAMPLES:
1. Analyze the student's text based on their selected Focus Area.
2. Identify the top 1-2 weaknesses.
3. For every weakness you identify, you must provide a **"Structural Example"**.
4. CRITICAL: The "Structural Example" must be about a COMPLETELY DIFFERENT TOPIC than the student's essay.
5. Never rewrite their actual sentence. Only show them the *pattern* of a better sentence.

Here are examples of how you should respond (Few-Shot Training):
{FEW_SHOT_EXAMPLES}
"""

# 3. Sidebar: Settings & Reset
with st.sidebar:
    st.header("⚙️ Coach Settings")
    grade_level = st.select_slider("Grade Level", options=["9th", "10th", "11th", "12th"])
    focus_area = st.selectbox(
        "Current Focus",
        ["General Critique", "Grammar & Syntax", "Argument & Logic", "Tone & Voice"]
    )
    
    st.divider()
    
    # --- RESET BUTTON LOGIC ---
    # If clicked, we clear the session state list
    if st.button("🔄 Reset Conversation", type="primary"):
        st.session_state.messages = []
        st.rerun()

# 4. Initialize Session State (Memory)
if "messages" not in st.session_state:
    st.session_state.messages = []

# 5. Display Chat History
st.title("🎓 Digital Writing Coach")
if len(st.session_state.messages) == 0:
    st.markdown("👋 **Hello!** Paste your draft below to get started. I'm here to coach, not to copy-edit!")

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 6. Chat Input & Processing
if prompt := st.chat_input("Paste your text or ask a question..."):
    
    # A. Display User Message
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # B. Generate Response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        
        # We inject the "Current Settings" into the System Prompt dynamically
        # This ensures if the user changes the "Focus Area" mid-chat, the AI knows.
        dynamic_system_prompt = SYSTEM_PROMPT + f"\n\nCURRENT CONTEXT: Student is in {grade_level} Grade. Focus on: {focus_area}."
        
        try:
            # Construct the full message history for the API
            # System prompt first, then the conversation history
            messages_payload = [{"role": "system", "content": dynamic_system_prompt}] + st.session_state.messages
            
            response = client.chat.completions.create(
                model="gpt-4o",
                messages=messages_payload,
                temperature=0.7,
                stream=True # Streaming makes it feel faster
            )
            
            # Stream the response chunk by chunk
            for chunk in response:
                if chunk.choices[0].delta.content is not None:
                    full_response += chunk.choices[0].delta.content
                    message_placeholder.markdown(full_response + "▌")
            
            message_placeholder.markdown(full_response)
            
            # Add AI response to history
            st.session_state.messages.append({"role": "assistant", "content": full_response})

        except Exception as e:
            st.error(f"Error: {e}")