import os

import streamlit as st
from huggingface_hub import InferenceClient

chat_parameters = {
    "system_prompt": "You are a friendly Chatbot.",
    "max_tokens": 512,
    "temperature": 0.7,
    "top_p": 0.95,
}


def get_chat_parameters():
    with st.sidebar:
        chat_parameters["system_prompt"] = st.text_area(
            "System Prompt", value=chat_parameters["system_prompt"])
        chat_parameters["max_tokens"] = st.slider(
            "Max Tokens", min_value=64, max_value=2048,
            value=chat_parameters["max_tokens"], step=64)
        chat_parameters["temperature"] = st.slider(
            "Temperature", min_value=0.1, max_value=2.0,
            value=chat_parameters["temperature"], step=0.1)
        chat_parameters["top_p"] = st.slider(
            "Top-p", min_value=0.1, max_value=1.0,
            value=chat_parameters["top_p"], step=0.05)


def get_reponse(client, user_input):
    # Construct message list with system prompt and history
    messages = [{"role": "system",
                 "content": chat_parameters["system_prompt"]}]
    messages.extend(st.session_state.chat_history)
    messages.append({"role": "user", "content": user_input})

    with st.spinner("Generating response..."):
        response_text = ""
        placeholder = st.empty()
        for chunk in client.chat_completion(
            messages=messages,
            max_tokens=chat_parameters["max_tokens"],
            stream=True,
            temperature=chat_parameters["temperature"],
            top_p=chat_parameters["top_p"],
        ):
            if chunk.choices and chunk.choices[0].delta.content:
                token = chunk.choices[0].delta.content
                response_text += token
                placeholder.markdown(f"**Assistant:** {response_text}")

        # Update chat history
        st.session_state.chat_history.append(
            {"role": "user", "content": user_input})
        st.session_state.chat_history.append(
            {"role": "assistant", "content": response_text})


def chatbot():
    model_id = "openai/gpt-oss-20b"
    hf_token = os.getenv("HF_TOKEN")

    # Initialize InferenceClient
    client = InferenceClient(token=hf_token, model=model_id)

    # Initialize chat history
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    # Page title
    st.title("🤖 Chatbot")

    # Sidebar: system prompt and generation parameters
    get_chat_parameters()

    # User input field
    user_input = st.text_input("Type your message...")

    # Handle user input
    if st.button("✅ Send"):
        if user_input.strip():
            get_reponse(client, user_input)
        else:
            st.warning("Please enter a message.")

    # Display chat history
    if st.session_state.chat_history:
        st.markdown("---")
        st.subheader("Chat History")
        for msg in st.session_state.chat_history:
            role = "🧑 You" if msg["role"] == "user" else "🤖 Assistant"
            st.markdown(f"**{role}:** {msg['content']}")


if __name__ == "__main__":
    chatbot()