# app.py
import streamlit as st
import os
from dotenv import load_dotenv
from langchain.chat_models import ChatOpenAI

# Load HuggingFace API token
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Initialize the HuggingFace LLM
llm = ChatOpenAI(
    openai_api_key=OPENAI_API_KEY,
    model_name="gpt-4o-mini",
    temperature=0.7,
    max_tokens=50
)

# Streamlit UI setup
st.set_page_config(page_title="🧠 HuggingFace Chatbot", page_icon="🤖")
st.title("🤖 HuggingFace Chatbot")
st.caption("Built with Streamlit + LangChain (50-word max answers)")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "Hi there! Ask me anything."}
    ]

# Display chat messages
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Chat input
if prompt := st.chat_input("Type your message here..."):
    # Add user message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Construct prompt (only user + assistant, formatted)
    conversation = "You are a helpful assistant. Keep replies within 50 words.\n\n"
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            conversation += f"User: {msg['content']}\n"
        elif msg["role"] == "assistant":
            continue  # Don't include previous assistant replies

    conversation += "Assistant:"  # Prompt the model to continue

    # Generate model response
    with st.chat_message("assistant"):
        response = llm.invoke(conversation)
        st.markdown(response.content)
        st.session_state.messages.append({"role": "assistant", "content": response.content})