import streamlit as st
from huggingface_hub import InferenceClient
import os

# Get token from Space secrets
API_TOKEN = os.getenv("HF_TOKEN")
client = InferenceClient(token=API_TOKEN)

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "corrected_sentence" not in st.session_state:
    st.session_state.corrected_sentence = ""

# Title of the app
st.title("Sentence Improver & Chat App")

# --- Sentence Correction Section ---
st.subheader("Improve a Sentence")
user_input = st.text_input("Enter a sentence to improve:", "I goed to the park and play.")

if st.button("Improve Sentence"):
    if user_input:
        prompt = f"Correct and improve this sentence: '{user_input}'"
        try:
            response = client.text_generation(
                prompt,
                model="mistralai/Mixtral-8x7B-Instruct-v0.1",  # Or your working model
                max_new_tokens=100,
                temperature=0.7,
            )
            st.session_state.corrected_sentence = response.strip()
            st.success(f"Improved Sentence: {st.session_state.corrected_sentence}")
        except Exception as e:
            st.error(f"Error: {str(e)}")
    else:
        st.warning("Please enter a sentence first!")

# --- Chat Section ---
st.subheader("Chat About the Corrected Sentence")
if st.session_state.corrected_sentence:
    # Display chat history
    for speaker, message in st.session_state.chat_history:
        st.write(f"**{speaker}:** {message}")

    # Chat input with Enter key submission
    chat_input = st.text_input(
        "Ask something about the corrected sentence (press Enter to send):",
        key="chat_input",
        value="",  # Clear after each submission
        on_change=lambda: submit_chat(),  # Trigger on Enter
    )

    # Function to handle chat submission
    def submit_chat():
        chat_text = st.session_state.chat_input
        if chat_text:
            # Build prompt with corrected sentence as context
            prompt = (
                f"The corrected sentence is: '{st.session_state.corrected_sentence}'. "
                f"User asks: '{chat_text}'. Respond naturally."
            )
            try:
                response = client.text_generation(
                    prompt,
                    model="mistralai/Mixtral-8x7B-Instruct-v0.1",  # Same model
                    max_new_tokens=150,
                    temperature=0.7,
                )
                # Add to chat history
                st.session_state.chat_history.append(("You", chat_text))
                st.session_state.chat_history.append(("LLM", response.strip()))
                # Clear the input field by resetting the key
                st.session_state.chat_input = ""
            except Exception as e:
                st.error(f"Error in chat: {str(e)}")

else:
    st.write("Improve a sentence first to start chatting!")