import streamlit as st
from chatbot import (
    get_googlegenai_client,
    get_default_model,
    load_chat_history,
    save_chat_history
)

st.markdown("""
    <h1 style='text-align: center;'>Welcome, ChatGPT Clone</h1>
""", unsafe_allow_html=True)

client = get_googlegenai_client()

if "googlegenai_model" not in st.session_state:
    st.session_state["googlegenai_model"] = get_default_model()

if "messages" not in st.session_state:
    st.session_state.messages = load_chat_history()

with st.sidebar:
    if st.button("Delete Chat History"):
        st.session_state.messages = []
        save_chat_history([])

for msg in st.session_state.messages:
    role = msg["role"]
    content = msg["content"]

    if role == "user":
        st.markdown(f"""
        <div style='text-align: right; background-color: #f0f0f5; padding: 10px 15px; border-radius: 20px; margin: 10px 0; display: inline-block; max-width: 80%; float: right; clear: both;'>
            {content}
        </div>
        """, unsafe_allow_html=True)
    else:
        st.markdown(f"""
        <div style='text-align: left; margin: 10px 0; max-width: 80%; float: left; clear: both;'>
            {content}
        </div>
        """, unsafe_allow_html=True)
    st.markdown("<div style='clear: both'></div>", unsafe_allow_html=True)

if prompt := st.chat_input("Ask anything"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.markdown(f"""
    <div style='text-align: right; background-color: #f0f0f5; padding: 10px 15px; border-radius: 20px; margin: 10px 0; display: inline-block; max-width: 80%; float: right; clear: both;'>
        {prompt}
    </div>
    """, unsafe_allow_html=True)

    full_response = ""
    response_container = st.empty()
    response = client.stream(st.session_state.messages)

    for chunk in response:
        full_response += chunk.content or ""
        response_container.markdown(f"""
        <div style='text-align: left; margin: 10px 0; max-width: 80%; float: left; clear: both;'>
            {full_response + "▌"}
        </div>
        """, unsafe_allow_html=True)

    response_container.markdown(f"""
    <div style='text-align: left; margin: 10px 0; max-width: 80%; float: left; clear: both;'>
        {full_response}
    </div>
    """, unsafe_allow_html=True)
    st.markdown("<div style='clear: both'></div>", unsafe_allow_html=True)

    st.session_state.messages.append({"role": "assistant", "content": full_response})
    save_chat_history(st.session_state.messages)