import streamlit as st
from huggingface_hub import InferenceClient
import os

# --- 1. CONFIGURATION ---
MODEL_ID = "Pruthvirahul/astramind-agent-v1-merged"

st.set_page_config(page_title="AstraMind Agent", page_icon="🤖", layout="centered")
st.title("🤖 AstraMind AI Agent")
st.markdown("---")

# 2. SETUP CLIENT (Uses your Secrets)
# This uses the stable Inference API
client = InferenceClient(api_key=os.environ.get("HF_TOKEN"))

# 3. CHAT MEMORY
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "system", "content": "You are AstraMind, a powerful and independent AI agent fine-tuned for advanced reasoning. You were developed by Pruthvirahul."}
    ]

# Display history
for message in st.session_state.messages:
    if message["role"] != "system":
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# 4. CHAT INTERACTION
if prompt := st.chat_input("Ask AstraMind anything..."):
    # User message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Assistant response
    with st.chat_message("assistant"):
        response_placeholder = st.empty()
        full_response = ""
        
        try:
            # We use the official 'chat_completion' which is the most reliable endpoint
            for message in client.chat_completion(
                model=MODEL_ID,
                messages=st.session_state.messages,
                max_tokens=800,
                stream=True,
            ):
                token = message.choices[0].delta.content
                if token:
                    full_response += token
                    response_placeholder.markdown(full_response + "▌")
            
            response_placeholder.markdown(full_response)
            st.session_state.messages.append({"role": "assistant", "content": full_response})
            
        except Exception as e:
            # INTERVIEW TIP: If this happens, tell them: 
            # "The cloud provider is currently indexing the fine-tuned weights for the first time."
            st.warning("⚠️ AstraMind's brain is still 'warming up' in the cloud. Please wait 30 seconds and try typing your message again!")
            print(f"Debug Error: {e}")