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

# Load environment variables
load_dotenv()

# Check for credentials and configure Gemini
if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
    credentials, _ = default()
    genai.configure(credentials=credentials)
else:
    st.error("❌ GOOGLE_APPLICATION_CREDENTIALS not set in the environment.")
    st.stop()

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

# Custom Streamlit styling (optional)
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 Title
st.title("📚 Grammar Guardian")
st.caption("Correct grammar, get explanations, and improve your writing skills!")

# Initialize session state
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("Write a sentence you'd like to improve...")

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

    with st.spinner("Analyzing..."):
        try:
            full_prompt = f"""
You are an advanced grammar assistant. Correct the given sentence and explain the changes clearly.

Respond in the following format:
**Correction:** <Corrected Sentence>
**Explanation:** <Why you corrected it>

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

            st.session_state.history.append({"role": "assistant", "content": result})

            with st.chat_message("assistant"):
                st.markdown(result)

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