# 패키지 불러오기
import re
import streamlit as st
from transformers import pipeline

# 글로벌 변수 선언
## 본 코드에서 사용된 신문기사 요약용 모델(https://huggingface.co/gangyeolkim/kobart-korean-summarizer-v2)은 테스트용 입니다.
pipe = pipeline("summarization", model="gangyeolkim/kobart-korean-summarizer-v2")

## assitant 아바타 이미지 경로
assistant_icon_path = "https://huggingface.co/spaces/randmimc/aitom/resolve/main/chat_icon/assistant.ico"

## 세션 상태에서 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 = []

st.markdown("""
    <style>
    .chat-container { /* 채팅 메시지를 감싸는 컨테이너를 설명합니다. */
        display: flex; /* 채팅 요소(아바타, 말풍선 등)을 가로 방향으로 정렬합니다. */
        margin-bottom: 10px; /*메세지간 간격을 10px로 합니다. */
    }
    .user-container {
        justify-content: flex-end; /* 사용자의 메시지를 우측정렬합니다.*/
    }
    .assistant-container {
        justify-content: flex-start; /* 어시스턴스 메시지를 좌측 정렬합니다.*/
        flex-direction: row; /* Avatar와 Bubble을 수평으로 정렬합니다. */
    }
    .chat-bubble { /* 말풍선 스타일을 정의합니다. */
        padding: 10px; /* 말풍선 안에 메시지간 간격을 10px로 합니다. */
        border-radius: 10px; /*말풍선의 모서리를 동글게 만듭니다.*/
        max-width: 70%; /* 한줄이 너무 길지 않게 제한합니다. */
        word-wrap: break-word; /* 텍스트를 단어 단위로 줄내림합니다. */
        white-space: pre-wrap; /* 텍스트를 단어 단위로 줄내림합니다. */
    }
    .user-bubble { /* 사용자 메시지 스타일을 정의합니다. */
        font-size: 14px; /* 폰트 사이즈를 지정합니다. */
        color: #333;  /* 폰트 색상을 지정합니다. */
        background-color: #e5e5e5; /* 말풍선 배경색을 지정합니다. */
    }
    .assistant-bubble {
        font-size: 14px; /* 폰트 사이즈를 지정합니다. */
        color: #333;  /* 폰트 색상을 지정합니다. */
        text-align: left; /* 메세지를 왼쪽 정렬 합니다. */
        background-color: #feeddf; /* 말풍선 배경색을 지정합니다. */
    }
    .avatar { /* 프로필 이미지(아바타) 스타일을 지정합니다. */
        width: 30px; /* 아바타 폭을 지정합니다. */
        height: 30px; /* 아바타 너을 지정합니다. */
        border-radius: 50%; /* 아바타를 원경으로 만듭니다. */
        margin-right: 10px; /*메세지간 간격을 10px로 합니다. */
    }
    </style>
""", unsafe_allow_html=True)

def user_message_style(question):
    return f"""<div class="chat-container user-container">
        <div class="chat-bubble user-bubble">{question}</div>
    </div>"""

def assistant_message_style(assistant_icon_path, answer):
    return f"""<div class="chat-container assistant-container">
        <img src="{assistant_icon_path}" class="avatar">
        <div class="chat-bubble assistant-bubble">{answer}</div>
    </div>"""

for message in st.session_state.messages:
    if message["role"] == "user":
        st.markdown(user_message_style(message["content"]), unsafe_allow_html=True)
    else:
        st.markdown(assistant_message_style(assistant_icon_path, message["content"]), unsafe_allow_html=True)

if not st.session_state.messages:
    greeting = '신문기사 요약 모델 입니다. 사용을 위해서는 "yes"를 입력해 주세요'
    st.session_state.messages.append({"role": "assistant", "content": greeting})
    st.markdown(assistant_message_style(assistant_icon_path, greeting), unsafe_allow_html=True)

if prompt := st.chat_input():
    # 사용자 질문 세션에 저장
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 사용자 질문 화면에 표시
    st.markdown(user_message_style(prompt), unsafe_allow_html=True)
    
    # 인증 
    if st.session_state.approval_state == "require":
        if prompt.lower() == "yes":
            st.session_state.approval_state = "approved"
            response = "이제 사용이 가능합니다. 단, 모델 크기가 작은 모델이여서 성능이 기대에 못미칠 수 있습니다."
        else:
            response = "인증에 실패하였습니다. 다시 yes를 입력해 주세요"
    # 응답 생성
    else:
        result = re.sub(r'\s+', ' ', prompt)
        summarized = pipe(result)
        response = summarized[0]["summary_text"]
    
    # 어시스턴스의 메세지를 세션에 저장
    st.session_state.messages.append({"role": "assistant", "content": response})
    # 어시스턴스의 메세지를 화면에 표시
    st.markdown(assistant_message_style(assistant_icon_path, response), unsafe_allow_html=True)