import streamlit as st
from models import demo  # Import the demo object from models.py

# --- Streamlit App Configuration ---
st.set_page_config(
    page_title="DeepSeek Chatbot",
    page_icon="🤖",
    layout="wide"
)

# --- App Title and Description ---
st.title("DeepSeek Chatbot")
st.markdown("""
    Created by [ruslanmv.com](https://ruslanmv.com/)
    This is a demo of different DeepSeek models. Select a model, type your message, and click "Submit".
    You can also adjust optional parameters like system message, max new tokens, temperature, and top-p.
""")

# --- Sidebar for Model Selection and Parameters ---
with st.sidebar:
    st.header("Options")
    model_choice = st.radio(
        "Choose a Model",
        options=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
        index=1  # Default to "DeepSeek-R1"
    )
    
    with st.expander("Optional Parameters", expanded=False):
        system_message = st.text_area(
            "System Message",
            value="You are a friendly Chatbot created by ruslanmv.com",
            height=100
        )
        max_new_tokens = st.slider(
            "Max New Tokens",
            min_value=1,
            max_value=4000,
            value=200
        )
        temperature = st.slider(
            "Temperature",
            min_value=0.10,
            max_value=4.00,
            value=0.70
        )
        top_p = st.slider(
            "Top-p (nucleus sampling)",
            min_value=0.10,
            max_value=1.00,
            value=0.90
        )

# --- Chatbot Function ---
def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
    # Create payload for the model
    payload = {
        "messages": [{"role": "user", "content": input_text}],
        "system": system_message,
        "max_tokens": max_new_tokens,
        "temperature": temperature,
        "top_p": top_p
    }
    
    # Run inference using the selected model
    try:
        response = demo(payload)  # Use the demo object directly
        if isinstance(response, dict) and "choices" in response:
            assistant_response = response["choices"][0]["message"]["content"]
        else:
            assistant_response = "Unexpected model response format."
    except Exception as e:
        assistant_response = f"Error: {str(e)}"
    
    # Append user and assistant messages to history
    history.append((input_text, assistant_response))
    return history

# --- Chat History Management ---
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# --- Chat Interface ---
st.header("Chat with DeepSeek")

# Display chat history
for user_msg, assistant_msg in st.session_state.chat_history:
    with st.chat_message("user"):
        st.write(user_msg)
    with st.chat_message("assistant"):
        st.write(assistant_msg)

# Input for new message
input_text = st.chat_input("Type your message here...")

# Handle new message submission
if input_text:
    # Update chat history
    st.session_state.chat_history = chatbot(
        input_text,
        st.session_state.chat_history,
        model_choice,
        system_message,
        max_new_tokens,
        temperature,
        top_p
    )
    
    # Rerun the app to display the updated chat history
    st.rerun()