import base64
import streamlit as st
from chatbot.core import get_chat_response

# Helper function to load and convert image to base64
def load_image_as_base64(image_path):
    with open(image_path, "rb") as img_file:
        img_bytes = img_file.read()
    encoded = base64.b64encode(img_bytes).decode("utf-8")
    return f"data:image/png;base64,{encoded}"

# Convert your local avatar image (make sure it's in the same directory or adjust the path)
avatar_base64 = load_image_as_base64("avatar.png")

st.set_page_config(page_title="Kumiko v1 Assistant", layout="wide")

st.markdown(
    """
    <style>
    .title {
        text-align: center;
        font-size: 2em;
        font-weight: bold;
    }
    </style>
    <h1 class="title">Kumiko v1 Assistant</h1>
    """,
    unsafe_allow_html=True
)

if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for message in st.session_state.messages:
    if message["role"] == "assistant":
        avatar = avatar_base64  # Use the base64 avatar for assistant messages
    else:
        avatar = None
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

# Chat input
user_input = st.chat_input("Nhập tin nhắn của bạn...")

if user_input:
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    response = get_chat_response(user_input)
    
    st.session_state.messages.append({"role": "assistant", "content": response})
    with st.chat_message("assistant", avatar=avatar_base64):
        st.markdown(response)