import streamlit as st
import requests
import os

# Initialize chat history in session state
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

st.title("DeepSeek-V3 Chatbot")

# Display chat messages from history
for message in st.session_state.chat_history:
    st.write(f"**{message['role'].capitalize()}:** {message['content']}")

# Input for user message
user_input = st.text_input("You:", key="user_input")

if st.button("Send"):
    if user_input:
        # Append user message to chat history
        st.session_state.chat_history.append({"role": "user", "content": user_input})

        # Prepare payload for Hugging Face Inference API
        payload = {
            "inputs": {
                "past_user_inputs": [msg['content'] for msg in st.session_state.chat_history if msg['role'] == 'user'],
                "generated_responses": [msg['content'] for msg in st.session_state.chat_history if msg['role'] == 'assistant'],
                "text": user_input
            }
        }

        # Fetch API token from environment variables
        API_TOKEN = os.getenv("API_TOKEN")
        if not API_TOKEN:
            st.error("API_TOKEN not found. Please set it as an environment variable.")
        else:
            api_url = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-V3"
            headers = {"Authorization": f"Bearer {API_TOKEN}"}

            # Send request to Hugging Face Inference API
            response = requests.post(api_url, headers=headers, json=payload)

            if response.status_code == 200:
                assistant_reply = response.json().get('generated_text', '')
                # Append assistant's reply to chat history
                st.session_state.chat_history.append({"role": "assistant", "content": assistant_reply})
                st.write(f"**Assistant:** {assistant_reply}")
            else:
                st.error("Error: Unable to get response from DeepSeek-V3 model.")
    else:
        st.warning("Please enter a message.")

if st.button("Clear Chat"):
    st.session_state.chat_history = []