import streamlit as st
from transformers import pipeline
import re

def replace_multiple_spaces(text):
    return re.sub(r'\s+', ' ', text)

pipe = pipeline("summarization", model="gangyeolkim/kobart-korean-summarizer-v2")


assistant_icon_path = "https://huggingface.co/spaces/randmimc/aitom/resolve/main/chat_icon/assistant.ico"
user_icon_path = "https://huggingface.co/spaces/randmimc/aitom/resolve/main/chat_icon/user01.png"

# 세션 상태에서 approval_state 및 messages 초기화
if "approval_state" not in st.session_state:
    st.session_state.approval_state = "require"
if "messages" not in st.session_state:
    st.session_state.messages = []

# Custom CSS to align messages and wrap text
st.markdown("""
    <style>
    .chat-container {
        display: flex;
        align-items: flex-start;
        margin-bottom: 10px;
    }
    .user-container {
        justify-content: flex-end;
    }
    .assistant-container {
        justify-content: flex-start;
        flex-direction: row; /* Avatar와 Bubble을 수평 정렬 */
    }
    .chat-bubble {
        padding: 10px;
        border-radius: 10px;
        max-width: 70%;
        word-wrap: break-word;
        white-space: pre-wrap;
    }
    .user-bubble {
        background-color: #e5e5e5; 
        text-align: left;
    }
    .assistant-bubble {
        font-size: 14px;
        text-align: left; /* 왼쪽 정렬 */
        background-color: #f8d7da;
    }
    .avatar {
        width: 30px;
        height: 30px;
        border-radius: 50%;
        margin-right: 10px; /* Bubble과 간격 */
    }
    </style>
""", unsafe_allow_html=True)


# Display chat messages from history
for message in st.session_state.messages:
    if message["role"] == "user":
        st.markdown(
            f"""
            <div class="chat-container user-container">
                <div class="chat-bubble user-bubble">{message["content"]}</div>
            </div>
            """, unsafe_allow_html=True)
    else:
        st.markdown(
            f"""
            <div class="chat-container assistant-container">
                <img src="{assistant_icon_path}" class="avatar">
                <div class="chat-bubble assistant-bubble">{message["content"]}</div>
            </div>
            """, unsafe_allow_html=True)

# Start with an assistant message
if not st.session_state.messages:
    st.session_state.messages.append({
        "role": "assistant", 
        "avatar": assistant_icon_path, 
        "content": """더아이엠씨에 aitom은 현재 공사중에 있습니다.
        하지만 ! 간단히 활용해 보실 수 있는 요약 모델을 제공하고 있습니다.
        짧은 신문기사를 말해 주시면 요약문을 제공합니다.
        계속 사용하고 싶다면 "yes" 를 입력해주세요.
        """})
    st.rerun()

if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "avatar": user_icon_path, "content": prompt})
    if st.session_state.approval_state == "require":
        if prompt.lower() == "yes":
            st.session_state.approval_state = "approved"
            st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": "이제 간단한 요약 모델을 사용하실 수 있습니다. 간단한 신문기사를 입력해 주세요."})
        else:
            st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": '승인하지 못하였습니다. 승인을 위해 "yes"를 입력해주세요.'})
            
    elif st.session_state.approval_state == "approved":
        result = replace_multiple_spaces(prompt)
        summarized = pipe(result)
        st.session_state.messages.append({"role": "assistant", "avatar": assistant_icon_path, "content": summarized[0]["summary_text"]})

    st.rerun()