import os
import time
import streamlit as st
import google.generativeai as genai
from google.auth import default
from dotenv import load_dotenv

# ✅ Set page config FIRST
st.set_page_config(page_title="Grammar & Spelling Assistant", page_icon="🧠")

# Load environment variables
load_dotenv()

# ✅ Use Application Default Credentials
if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
    credentials, _ = default()
    genai.configure(credentials=credentials)
else:
    st.error("❌ GOOGLE_APPLICATION_CREDENTIALS is not set. Please add it to .env or Hugging Face Secrets.")
    st.stop()

# Initialize Gemini model
model = genai.GenerativeModel("gemini-2.0-pro")

# Optional styling
st.markdown("""
<style>
body, .stApp { background-color: #121212 !important; color: #e0e0e0 !important; }
.stChatInput { background: #222 !important; border: 1px solid #555 !important; }
</style>
""", unsafe_allow_html=True)

# App header
st.title("📚 Grammar Guardian")
st.caption("Correct grammar, get explanations, and improve your writing!")

# Initialize chat history
if "history" not in st.session_state:
    st.session_state.history = []

# Display chat history
for message in st.session_state.history:
    role, content = message["role"], message["content"]
    with st.chat_message(role):
        st.markdown(content)

# Chat input
prompt = st.chat_input("Type a sentence you'd like to improve...")

if prompt:
    # Show user message
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.history.append({"role": "user", "content": prompt})

    # Process input with Gemini
    with st.spinner("Analyzing..."):
        try:
            full_prompt = f"""
You are a grammar correction assistant.
When a user gives a sentence, do two things:
1. Correct the sentence.
2. Explain clearly why you corrected it.

Respond using this format:
**Correction:** <Corrected sentence>
**Explanation:** <Explanation of why it was corrected>

Sentence: {prompt}
"""
            response = model.generate_content(full_prompt)
            result = response.text

            # Show assistant message
            with st.chat_message("assistant"):
                st.markdown(result)
            st.session_state.history.append({"role": "assistant", "content": result})

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