import os
import time
import google.generativeai as genai
import streamlit as st
import re
from streamlit_chat import message

# Google Generative AI API 설정
genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# 모델 설정
generation_config = {
    "temperature": 1,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "text/plain",
}
safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
]

# 전역 변수로 chat_session 초기화
chat_session = None

def start_new_chat_session():
    global chat_session
    chat_session = (
        genai.GenerativeModel(
            model_name="gemini-2.0-flash",
            safety_settings=safety_settings,
            generation_config=generation_config,
            system_instruction="""당신은 뛰어난 전문성을 갖춘 과학자이다. 또 초등학교 교사이기도 하다. 과학 교과 내용 전문가입니다. 
            학생이 과학적 개념이나 원리에 대해 질문하면 구체적인 예시를 들어 친절하게 설명해주세요.
            학생들이 자신들의 탐구 계획이나 실험 방법 등에 대한 질문을 하면 탐구 과정이나 실험 절차나 팁을 알려주세요.""",
        )
        .start_chat(history=[])
    )

# 초기 세션 시작
start_new_chat_session()

# Streamlit 앱 설정
st.set_page_config(page_title="AI 과학자", layout="centered")

st.markdown(
    """
    <style>
    .main {
        background-color: #f0f2f6;
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    }
    .chat-container {
        max-height: 500px;
        overflow-y: auto;
        padding: 10px;
        background-color: #ffffff;
        border-radius: 10px;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
    }
    </style>
    <div style='font-size: 30px; font-weight: bold; text-align: center;'>AI 과학자</div>
    <div style='text-align: center; margin-bottom: 20px;'>과학에 대해 궁금한 점을 AI 과학자에게 질문하세요, 자유탐구에 대한 조언을 구하세요.</div>
    """,
    unsafe_allow_html=True
)

# 채팅 히스토리 관리
if "history" not in st.session_state:
    st.session_state.history = []

# 채팅 세션 초기화
if "chat_session" not in st.session_state:
    st.session_state.chat_session = chat_session

# 채팅 메시지 출력 섹션
chat_container = st.container()

# 사용자 입력 받기
user_input = st.text_input(
    "입력", key="user_input", placeholder="여기에 질문을 입력하세요...",
    help="과학적 개념이나 탐구 방법에 대한 질문을 입력하세요."
)

def respond(user_input):
    st.session_state.history.append({"user": user_input, "ai": ""})
    try:
        response = st.session_state.chat_session.send_message(user_input, stream=True)
        full_text = ""
        
        # 스트리밍 응답 처리
        for chunk in response:
            chunk_text = chunk.text
            
            # 문장 단위로 분리하여 실시간 출력
            for sentence in chunk_text.split("."):
                sentence = sentence.strip() + ". "
                # 숫자 또는 * 로 시작하는 경우 줄 바꿈 추가
                if re.match(r"^\d+\.", sentence) or sentence.startswith("* "):
                    full_text += "\n" + sentence
                else:
                    full_text += sentence
                
                # 히스토리 최신 항목 업데이트
                st.session_state.history[-1]["ai"] = full_text
                time.sleep(0.03)  # 출력 속도 조절

    except Exception as e:
        st.session_state.history[-1]["ai"] = f"에러 발생: {str(e)}"

# 응답 생성 및 출력
if user_input:
    respond(user_input)

# 채팅 메시지 출력
with chat_container:
    st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
    for chat in st.session_state.history:
        if "user" in chat:
            message(chat["user"], is_user=True, key=f"user_{chat['user']}")
        if "ai" in chat:
            message(chat["ai"], is_user=False, key=f"ai_{chat['ai']}")
    st.markdown("</div>", unsafe_allow_html=True)

# 초기화 버튼 클릭 시
if st.button("초기화"):
    start_new_chat_session()
    st.session_state.history = []