import streamlit as st
import requests

# -----------------------------------
# 1. Hugging Face API Configuration
# -----------------------------------
API_URL = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"

def query(payload):
    """
    Query the Hugging Face Inference API with the given payload.
    Keeps the original approach: payload = {"inputs": user_input}.
    """
    headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()


# -----------------------------------
# 2. Streamlit Page Settings
# -----------------------------------
st.set_page_config(
    page_title="DeepSeek Chatbot - ruslanmv.com",
    page_icon="🤖",
    layout="centered"
)

# -----------------------------------
# 3. Session State Initialization
# -----------------------------------
# We'll keep a chat history in st.session_state
if "messages" not in st.session_state:
    st.session_state.messages = []

# -----------------------------------
# 4. Sidebar Configuration
# -----------------------------------
with st.sidebar:
    st.header("Configuration")
    st.markdown("[Get your HuggingFace Token](https://huggingface.co/settings/tokens)")

    # Although these parameters are shown on the sidebar, we won't actually
    # pass them to the payload in `query()`, to strictly preserve the "original" approach.
    st.write("**NOTE:** These sliders do not affect the inference in this demo.")
    system_message = st.text_area(
        "System Message (display only)",
        value="You are a friendly Chatbot created by ruslanmv.com",
        height=100
    )
    max_tokens = st.slider("Max Tokens (not used here)", 1, 4000, 512)
    temperature = st.slider("Temperature (not used here)", 0.1, 4.0, 0.7)
    top_p = st.slider("Top-p (not used here)", 0.1, 1.0, 0.9)

# -----------------------------------
# 5. Main Chat Interface
# -----------------------------------
st.title("🤖 DeepSeek Chatbot")
st.caption("Powered by Hugging Face Inference API - Original Inference Approach")

# Display the chat history, message by message
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# -----------------------------------
# 6. Capture User Input
# -----------------------------------
if user_input := st.chat_input("Type your message..."):
    # 6.1 Append user message to chat history
    st.session_state.messages.append({"role": "user", "content": user_input})
    
    # Display user's message
    with st.chat_message("user"):
        st.markdown(user_input)

    # -----------------------------------
    # 7. Query the Model
    # -----------------------------------
    try:
        with st.spinner("Generating response..."):
            # Prepare payload with the original approach
            payload = {"inputs": user_input}
            output = query(payload)

            # Check if the output is valid
            if (
                isinstance(output, list) 
                and len(output) > 0 
                and "generated_text" in output[0]
            ):
                assistant_response = output[0]["generated_text"]
            else:
                assistant_response = (
                    "Error: Unable to generate a response. Please try again."
                )

        # Display the assistant's response
        with st.chat_message("assistant"):
            st.markdown(assistant_response)

        # Store assistant's response in chat history
        st.session_state.messages.append(
            {"role": "assistant", "content": assistant_response}
        )

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