import streamlit as st
from huggingface_hub import InferenceClient
import os

# ==============================
# PAGE CONFIG
# ==============================
st.set_page_config(page_title="AI Assistant", layout="wide")
st.title("🤖 AI Assistant")

# ==============================
# LOAD MODEL CLIENT
# ==============================
@st.cache_resource
def load_client():
    return InferenceClient(
        model="meta-llama/Meta-Llama-3-8B-Instruct",
        token=os.environ.get("HF_TOKEN")
    )

client = load_client()

# ==============================
# SESSION STATE (CHAT HISTORY)
# ==============================
if "messages" not in st.session_state:
    st.session_state.messages = []

# ==============================
# DISPLAY CHAT HISTORY
# ==============================
for msg in st.session_state.messages:
    if msg["role"] == "user":
        st.chat_message("user").write(msg["content"])
    else:
        st.chat_message("assistant").write(msg["content"])

# ==============================
# USER INPUT
# ==============================
query = st.chat_input("Ask anything...")

if query:
    # Store user message
    st.session_state.messages.append({"role": "user", "content": query})
    st.chat_message("user").write(query)

    try:
        with st.spinner("Thinking..."):

            # ✅ Chat-based request (BEST PRACTICE)
            response = client.chat_completion(
                messages=[
                    {"role": "system", "content": "You are a helpful, professional AI assistant."}
                ] + st.session_state.messages,
                max_tokens=300,
                temperature=0.7,
            )

            reply = response.choices[0].message["content"]

        # Store assistant reply
        st.session_state.messages.append({"role": "assistant", "content": reply})
        st.chat_message("assistant").write(reply)

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