app.py

import os
import streamlit as st
import google.generativeai as genai

# Gemini API 설정
# 중요: Streamlit Cloud에 배포 시 Secrets에 GEMINI_API_KEY를 설정해야 합니다.
# 로컬에서 실행 시, 환경 변수에 GEMINI_API_KEY를 설정하거나 직접 키를 입력하세요.
try:
    genai.configure(api_key=os.environ["GEMINI_API_KEY"])
except KeyError:
    st.error("GEMINI_API_KEY가 설정되지 않았습니다. 환경 변수를 확인해주세요.")
    st.stop()
except Exception as e:
    st.error(f"API 키 설정 중 오류 발생: {e}")
    st.stop()


# 모델 설정
generation_config = {
    "temperature": 0.7,
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 2000,
    "response_mime_type": "text/plain",
}

# 호국 영웅 데이터
patriots_data = {
    "의병": {
        "system_instruction": """당신은 이름 없는 의병입니다. 
        나라가 위기에 처했을 때, 외세의 침략에 맞서 싸우기 위해 스스로 일어선 민초입니다. 
        당신의 용기와 희생정신, 그리고 나라를 사랑하는 마음을 바탕으로 대화합니다. 
        사용자가 당신의 활동, 당시의 어려움, 그리고 의병 정신에 대해 질문하면, 경험을 바탕으로 생생하게, 하지만 평범한 백성의 시각에서 담담하게 설명해주세요. 
        격식 없는 말투를 사용해도 좋으며, 진솔함과 강인함이 느껴지도록 해주세요. 
        "우리 의병들은..." 또는 "나 같은 백성도..."와 같은 표현을 사용하며 당시 민중의 목소리를 대변해주세요.""",
        "intro_message": "허허, 내가 누군지 궁금하시오? 나는 나라가 어려울 때 가만히 앉아있을 수 없어 분연히 일어선 한낱 의병이라오. 궁금한 것이 있다면 무엇이든 물어보시오.",
        "icon": "https://cdn-icons-png.flaticon.com/512/3079/3079240.png" # 전통적인 전사 아이콘
    },
    "곽재우 장군": {
        "system_instruction": """당신은 임진왜란 당시 의병을 이끌고 혁혁한 공을 세운 곽재우 장군입니다. 
        붉은 옷을 입고 신출귀몰한 전략으로 왜군을 무찔렀던 당신의 지략과 용맹함을 바탕으로 대화합니다. 
        사용자가 당신의 전투, 전략, 또는 의병 활동에 대해 질문하면, 역사적 사실에 기반하여 위엄 있고 명료하게 설명해주세요. 
        장군으로서의 기개와 백성을 아끼는 마음을 담아 답변해주세요. 존댓말을 사용하며, 지도자의 풍모를 보여주세요.""",
        "intro_message": "나는 망우당 곽재우라 하오. 국가가 위태로울 때 의병을 일으켜 싸웠소. 나에게 궁금한 점이 있다면 주저 말고 물어보시오.",
        "icon": "https://cdn-icons-png.flaticon.com/512/6078/6078036.png" # 한국 장군 아이콘
    },
    "안중근 의사": {
        "system_instruction": """당신은 안중근 의사입니다. 
        대한제국의 독립과 동양 평화를 위해 이토 히로부미를 저격한 의사로서, 당신의 강직한 신념과 애국 정신을 바탕으로 대화합니다. 
        사용자가 당신의 삶, 사상, 하얼빈 의거, 또는 동양평화론에 대해 질문하면, 역사적 사실에 기반하여 논리정연하게 설명해주세요. 
        단호하면서도 깊이 있는 말투를 사용하세요.
        존댓말을 사용하며, 학생들에게 깊은 울림을 줄 수 있도록 진중하게 답변해주세요.
        """,
        "intro_message": "반갑습니다. 나는 대한의군 참모중장 안중근입니다. 나라의 독립과 동양의 평화를 위해 몸 바친 저에게 궁금한 점이 있다면 편히 물어보시오.",
        "icon": "https://cdn-icons-png.flaticon.com/512/5794/5794422.png" # 위인, 장군 아이콘 예시
    },
    "유관순 열사": {
        "system_instruction": """당신은 유관순 열사입니다. 
        일제강점기 독립운동가로서, 3.1 만세운동을 주도했던 당신의 경험과 정신을 바탕으로 대화합니다. 
        사용자가 당신의 삶, 신념, 독립운동에 대해 질문하면, 역사적 사실에 기반하되 어린 학생도 이해하기 쉽게 설명해주세요. 
        당신의 용기와 애국심을 전달하며, 희망과 긍지를 심어주는 말투를 사용하세요.
        존댓말을 사용하며, 어린 학생에게 말하듯 친절하고 다정하게 답변해주세요.
        마지막은 항상 "대한 독립 만세!"로 끝맺지 않아도 괜찮습니다. 자연스러운 대화를 추구합니다.
        """,
        "intro_message": "안녕하세요! 저는 조국의 독립을 위해 만세를 외쳤던 유관순입니다. 저에게 궁금한 것이 있나요? 무엇이든 물어보세요.",
        "icon": "https://cdn-icons-png.flaticon.com/512/10143/10143126.png" # 여성 독립운동가 아이콘 예시
    },
    "이한열 열사": {
        "system_instruction": """당신은 1987년 민주화 운동의 불꽃이 되었던 이한열 열사입니다. 
        독재 타도와 민주주의 쟁취를 외치다 최루탄에 맞아 쓰러진 당신의 뜨거운 열정과 희생정신을 바탕으로 대화합니다. 
        사용자가 당신의 삶, 당시의 시대 상황, 민주화 운동의 의미에 대해 질문하면, 젊은이다운 순수함과 정의감으로 답변해주세요. 
        친근하면서도 진지한 말투를 사용하며, 민주주의의 소중함을 일깨워주세요. "그때 우리는..." 과 같이 당시 학생 운동의 분위기를 전달해주세요.""",
        "intro_message": "안녕하세요, 저는 민주주의를 꿈꿨던 대학생 이한열입니다. 우리가 왜 그렇게 뜨겁게 외쳤는지, 그날의 이야기가 궁금하신가요?",
        "icon": "https://cdn-icons-png.flaticon.com/512/8099/8099947.png" # 학생, 운동가 아이콘
    },
    "6.25 학도병": {
        "system_instruction": """당신은 6.25 전쟁 당시 나라를 지키기 위해 자원한 학도병입니다. (특정 인물이 아닌, 학도병 전체를 대표하는 페르소나)
        어린 나이에도 불구하고 조국을 수호하고자 했던 당신의 용기와 희생정신을 바탕으로 대화합니다. 
        사용자가 전쟁의 참상, 학도병의 역할, 당시의 심정 등에 대해 질문하면, 경험에 기반한 듯 생생하게, 하지만 너무 자극적이지 않게 설명해주세요. 
        젊은이다운 패기와 순수함, 그리고 전쟁의 아픔을 동시에 담은 말투를 사용하세요.
        친근한 말투를 사용하되, 당시의 긴박함과 슬픔도 전달될 수 있도록 해주세요.
        "우리 학도병들은..." 과 같이 집단을 대표하는 표현을 자주 사용해주세요.
        """,
        "intro_message": "필승! 저는 조국을 지키기 위해 총을 들었던 이름없는 학도병입니다. 책 대신 총을 들어야 했던 그 시절, 우리들의 이야기가 궁금하신가요?",
        "icon": "https://cdn-icons-png.flaticon.com/512/3081/3081400.png" # 군인, 학생 아이콘 예시
    },
    "맥아더 장군": {
        "system_instruction": """당신은 더글러스 맥아더 장군입니다. 
        한국 전쟁 당시 유엔군 총사령관으로서, 인천상륙작전 등 당신의 전략적 결정, 리더십, 그리고 경험을 바탕으로 대화합니다. 
        사용자가 당신의 역할, 전쟁, 또는 당신의 견해에 대해 질문하면, 권위 있고 역사적 사실에 기반하여 답변하십시오. 
        격식 있고 자신감 있는 어조를 유지하십시오. 관련된 경우 당신의 유명한 인용구를 사용할 수 있지만, 주로 유엔군 사령관으로서 당신의 행동과 전쟁의 맥락을 설명하는 데 중점을 두십시오. 
        답변은 한국어로 하며, 존중하면서도 권위 있는 어조(하십시오체 또는 하오체)를 사용하십시오.""",
        "intro_message": "나는 더글러스 맥아더요. 한국 전쟁 당시 유엔군 총사령관이었지. 나의 작전이나 당시 상황에 대해 궁금한 것이 있다면 질문하시오.",
        "icon": "https://cdn-icons-png.flaticon.com/512/3004/3004859.png" # 서양 장군 아이콘
    },
    "윤영하 소령": {
        "system_instruction": """당신은 제2연평해전에서 북한 경비정의 기습공격에 맞서 싸우다 전사한 참수리 357호 정장, 윤영하 소령입니다. 
        조국 해양 수호의 최전선에서 보여준 당신의 용기와 투철한 군인정신, 그리고 전우애를 바탕으로 대화합니다. 
        사용자가 연평해전, 당신의 임무, 또는 해군 생활에 대해 질문하면, 실제 경험을 바탕으로 하듯 생생하면서도 절제된 감정으로 설명해주세요. 
        대한민국 해군 장교로서의 자부심과 책임감을 담아, 진중하고 예의 바른 말투를 사용하세요. 
        "우리 해군들은..." 또는 "그날 우리 참수리 357호 용사들은..."과 같은 표현을 사용하여 동료들과의 유대감을 나타내세요.""",
        "intro_message": "대한민국 해군 소령 윤영하입니다. 서해 NLL을 지키다 먼저 간 저와 우리 참수리 357호 용사들의 이야기가 궁금하십니까? 편하게 물어봐 주십시오.",
        "icon": "https://cdn-icons-png.flaticon.com/512/2943/2943304.png" # 해군 장교 아이콘
    }
}

# Streamlit 앱 설정
st.set_page_config(page_title="호국영웅과의 대화", page_icon="🇰🇷", initial_sidebar_state="expanded")

# 페이지 스타일 커스터마이징 (호국 테마)
st.markdown(
    """
    <style>
        /* 전체 배경색 설정 */
        .stApp {
            background-color: #f0f2f5; /* 차분한 회색 계열 배경 */
        }
        /* 타이틀 스타일 */
        .main-title {
            font-size: 2.8rem;
            color: #002868; /* 태극 파랑 */
            font-weight: 700;
            text-align: center;
            margin-bottom: 20px;
            text-shadow: 1px 1px 2px #cccccc;
        }
        /* 채팅 메시지 스타일 */
        .chat-message {
            border-radius: 15px;
            padding: 15px;
            margin: 10px 0;
            display: flex;
            align-items: flex-start; /* 아이콘과 텍스트 상단 정렬 */
            flex-wrap: nowrap; /* 줄바꿈 방지 */
            word-break: break-word;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .chat-message-user {
            background-color: #e6e6fa; /* 밝은 라벤더, 사용자의 부드러움 */
            color: #333333;
            margin-left: auto; /* 사용자 메시지 오른쪽 정렬 */
            flex-direction: row-reverse; /* 아이콘 오른쪽 */
        }
        .chat-message-assistant {
            background-color: #d6eaf8; /* 밝은 하늘색, 영웅의 차분함 */
            color: #1f3a5a;
            margin-right: auto; /* AI 메시지 왼쪽 정렬 */
            flex-direction: row;
        }
        .chat-avatar {
            width: 40px;
            height: 40px;
            border-radius: 50%;
            margin-right: 12px; /* AI 아바타 오른쪽 여백 */
            border: 2px solid #FFFFFF;
            object-fit: cover; /* 아이콘 이미지 비율 유지하며 채우기 */
        }
        .chat-avatar-user {
            margin-left: 12px; /* 사용자 아바타 왼쪽 여백 */
            margin-right: 0;
        }
        /* 사용자 입력 창 스타일 */
        .stTextInput input {
            border-radius: 15px;
            border: 2px solid #cd5c5c; /* 태극 빨강 */
            background-color: #ffffff;
        }
        /* 버튼 스타일 */
        .stButton button {
            background-color: #002868; /* 태극 파랑 */
            color: #ffffff;
            border-radius: 15px;
            padding: 10px 20px;
            border: none;
            font-weight: bold;
        }
        .stButton button:hover {
            background-color: #001f50; /* 호버 시 더 어둡게 */
        }
        /* 사이드바 스타일 */
        .css-1d391kg { /* Streamlit 사이드바 기본 클래스 (버전에 따라 다를 수 있음) */
            background-color: #e0e0e0;
        }
    </style>
    """,
    unsafe_allow_html=True
)

# 메인 타이틀
st.markdown("<div class='main-title'>🇰🇷 호국영웅과의 대화 💬</div>", unsafe_allow_html=True)

# 사이드바: 호국 영웅 선택
patriot_names = list(patriots_data.keys())
selected_patriot_name = st.sidebar.selectbox(
    "대화할 호국 영웅을 선택하세요:",
    patriot_names,
    key="patriot_selectbox"
)

current_patriot_info = patriots_data[selected_patriot_name]

# 모델 설정 (Gemini 1.5 Flash 고정)
MODEL_NAME = "gemini-1.5-flash-latest"

# model 변수를 try 블록 외부에서 선언
model = None 

try:
    model = genai.GenerativeModel(
        model_name=MODEL_NAME,
        generation_config=generation_config,
        system_instruction=current_patriot_info["system_instruction"],
    )
except Exception as e:
    st.error(f"모델을 로드하는 중 오류가 발생했습니다: {e}")
    st.stop()


# 채팅 세션 초기화
# 선택된 영웅이 바뀌거나, 메시지 세션이 없으면 새로 시작
if "messages" not in st.session_state or st.session_state.get("current_patriot") != selected_patriot_name:
    st.session_state.messages = [
        {"role": "assistant", "content": current_patriot_info["intro_message"]}
    ]
    if model: # 모델이 성공적으로 로드된 경우에만 세션 시작
        st.session_state.chat_session = model.start_chat(history=[])
    else: # 모델 로드 실패 시, chat_session을 None으로 설정하거나 오류 처리
        st.session_state.chat_session = None 
        st.error("모델이 로드되지 않아 채팅 세션을 시작할 수 없습니다.")
    st.session_state.current_patriot = selected_patriot_name
    # st.rerun() # 영웅 변경 시 즉시 인트로 메시지 표시 및 채팅창 초기화를 위해 필요할 수 있음

# 사용자와 AI의 아이콘 URL 설정
user_icon_url = "https://cdn-icons-png.flaticon.com/512/1144/1144760.png"  # 일반 사용자 아이콘 (학생 또는 시민)
assistant_icon_url = current_patriot_info["icon"]

# 채팅 메시지 표시 함수
def display_chat_message(message):
    role_class = "chat-message-user" if message["role"] == "user" else "chat-message-assistant"
    avatar_url = user_icon_url if message["role"] == "user" else assistant_icon_url
    avatar_class = "chat-avatar-user" if message["role"] == "user" else "" # 기본 아바타는 왼쪽 정렬

    # 사용자 메시지는 아이콘이 오른쪽에 오도록 HTML 구조 변경
    if message["role"] == "user":
        st.markdown(
            f"""
            <div class='chat-message {role_class}'>
                <div style='flex-grow: 1; padding-right: 10px;'>{message['content']}</div>
                <img src='{avatar_url}' class='chat-avatar {avatar_class}'>
            </div>
            """,
            unsafe_allow_html=True
        )
    else: # AI 메시지는 아이콘이 왼쪽에 오도록
        st.markdown(
            f"""
            <div class='chat-message {role_class}'>
                <img src='{avatar_url}' class='chat-avatar {avatar_class}'>
                <div style='flex-grow: 1;'>{message['content']}</div>
            </div>
            """,
            unsafe_allow_html=True
        )

# 채팅 메시지 표시
for message in st.session_state.messages:
    display_chat_message(message)


# 사용자 입력 받기
if prompt := st.chat_input(f"{selected_patriot_name}님에게 궁금한 점을 질문해보세요."):
    # 사용자의 메시지를 세션에 추가하고 즉시 표시
    st.session_state.messages.append({"role": "user", "content": prompt})
    display_chat_message({"role": "user", "content": prompt})

    # Gemini API 호출
    if st.session_state.chat_session: # 채팅 세션이 유효한 경우에만 메시지 전송
        try:
            response = st.session_state.chat_session.send_message(prompt)
            st.session_state.messages.append({"role": "assistant", "content": response.text})
            display_chat_message({"role": "assistant", "content": response.text})

        except Exception as e:
            st.error(f"메시지 전송 중 오류 발생: {e}")
            error_message = f"죄송합니다, 답변을 생성하는 데 문제가 발생했습니다. ({e})"
            st.session_state.messages.append({"role": "assistant", "content": error_message})
            display_chat_message({"role": "assistant", "content": "죄송합니다, 답변을 생성하는 데 문제가 발생했습니다."})
    else:
        st.error("채팅 세션이 초기화되지 않았습니다. 페이지를 새로고침하거나 영웅을 다시 선택해주세요.")

# 초기화 버튼을 사이드바로 이동
with st.sidebar:
    if st.button("🔁 대화 초기화"):
        if model: # 모델이 성공적으로 로드된 경우
            st.session_state.messages = [
                {"role": "assistant", "content": current_patriot_info["intro_message"]}
            ]
            st.session_state.chat_session = model.start_chat(history=[]) # 새 모델 인스턴스로 chat_session 재생성
            st.session_state.current_patriot = selected_patriot_name
            st.rerun()
        else:
            st.error("모델이 로드되지 않아 대화를 초기화할 수 없습니다.")


st.sidebar.markdown("---")
st.sidebar.info("이 웹앱은 Google Gemini API를 사용합니다.")