import os
import time
import streamlit as st
import google.generativeai as genai
from streamlit_extras.colored_header import colored_header
import markdown

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

# 모델 설정
generation_config = {
    "temperature": 0.75,
    "top_p": 0.8,
    "top_k": 40,
    "max_output_tokens": 8000,
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash",
    generation_config=generation_config,
)

# 에니어그램 유형 정보 (날개 제외)
enneagram_types = {
    "1": "개혁가",
    "2": "조력가",
    "3": "성취자",
    "4": "예술가",
    "5": "사색가",
    "6": "충실가",
    "7": "낙천가",
    "8": "도전가",
    "9": "평화주의자"
}

# 날개 유형 정보
wing_types = {
    "1w9": "1w9",
    "1w2": "1w2",
    "2w1": "2w1",
    "2w3": "2w3",
    "3w2": "3w2",
    "3w4": "3w4",
    "4w3": "4w3",
    "4w5": "4w5",
    "5w4": "5w4",
    "5w6": "5w6",
    "6w5": "6w5",
    "6w7": "6w7",
    "7w6": "7w6",
    "7w8": "7w8",
    "8w7": "8w7",
    "8w9": "8w9",
    "9w8": "9w8",
    "9w1": "9w1",
}


# 본능 유형 정보
instinct_variants = {
    "sp": "자기보존",
    "so": "사회적",
    "sx": "성적"
}

# 관계 유형
relationship_types = {
    "연인": "Romantic Couple",
    "부부": "Married Couple",
    "친구": "Friends",
    "가족": "Family",
    "직장 동료": "Coworkers",
    "기타": "Others"
}


def generate_relationship_scenario(my_name, my_gender, my_type, my_wing, my_instinct, partner_name, partner_gender, partner_type, partner_wing, partner_instinct, relationship, situation):

    # 에니어그램 유형 정보 문자열 생성 (날개/본능 포함) - 조건문 추가
    def create_type_info(name, gender, type, wing, instinct):
        type_info = f"{name} ({gender}): {type} ({enneagram_types.get(type, '알 수 없는 유형')})"
        if wing:  # wing이 None이 아니면 추가
            type_info += f" 날개: {wing}" # wing_types 딕셔너리 사용 삭제
        if instinct: # instinct가 None이 아니면 추가
            type_info += f", 본능: {instinct_variants.get(instinct)}"
        return type_info

    my_type_info = create_type_info(my_name, my_gender, my_type, my_wing, my_instinct)
    partner_type_info = create_type_info(partner_name, partner_gender, partner_type, partner_wing, partner_instinct)

    system_prompt = f"""
    당신은 에니어그램 전문가이자 관계 코치입니다. 아래 정보를 바탕으로 관계 시나리오를 작성해주세요.
    [대상자 정보]
    - {my_type_info}
    - {partner_type_info}
    [관계]
    - {relationship}
    [상황]
    - {situation}
    [요청 사항]
    1. **{my_name} ({my_gender})의 에니어그램 유형에 대한 상세 설명:**
       - 핵심 가치, 성격적 특징, 강점과 약점, 주의할 점 등을 포함하여 자세히 설명해주세요.
    2. **{partner_name} ({partner_gender})의 에니어그램 유형에 대한 상세 설명:**
       - 핵심 가치, 성격적 특징, 강점과 약점, 주의할 점 등을 포함하여 자세히 설명해주세요.
    3. **에니어그램 유형을 바탕으로 한 구체적인 데이트 시나리오 작성:**
       - 만남, 대화 시작, 식사, 영화 관람, 갈등 발생 및 해결 등 데이트의 각 단계를 구체적인 대화 장면으로 묘사해주세요.
       - 각 단계에서 {my_name}과 {partner_name}의 감정 변화와 반응을 상세히 묘사하고, 각 인물의 말과 행동이 에니어그램 특성과 어떻게 연결되는지 부가 설명해주세요.
       - 시나리오에 {my_name}과 {partner_name}의 이름, 성별, 에니어그램 유형이 모두 반영되어야 합니다.
    4. **상대방을 서로 이해하기 위한 실질적인 조언:**
       - 각 대상자가 상대방의 에니어그램 유형을 이해하고, 더 나은 관계를 형성하기 위한 구체적인 조언을 제시해주세요.
       - 조언에 {my_name}과 {partner_name}의 이름, 성별, 에니어그램 유형이 모두 반영되어야 합니다.
       - 마지막에 상대방에게 해주면 좋은 말과 행동의 구체적인 예와 왜 상대방에게 좋은지에 대한 이유도 설명해주세요.
    """

    full_text = ""
    try:
        response = model.generate_content([system_prompt], stream=True)
        for chunk in response:
            full_text += chunk.text
            html_text = markdown.markdown(full_text)
            scenario_output_area.markdown(html_text, unsafe_allow_html=True)
            time.sleep(0.05)  # 출력 속도 조절
    except Exception as e:
        st.error(f"Error: {str(e)}")
        return ""

    return full_text


# Streamlit UI
st.set_page_config(page_title="에니어그램 관계 시뮬레이터", page_icon="🤝")

st.title("🤝 에니어그램 관계 시뮬레이터")
st.write("두 사람의 에니어그램 유형과 관계, 상황을 입력하여 관계 시나리오를 생성해보세요.")

# 사이드바에 탭 배치
with st.sidebar:
    tabs = st.tabs(["👤 내 정보", "👥 상대방 정보"])

    with tabs[0]:  # 내 정보 탭
        my_name = st.text_input("내 이름/닉네임", "나")
        my_gender = st.radio("내 성별", ["남성", "여성"])
        my_type = st.selectbox("내 에니어그램 유형", list(enneagram_types.keys()))
        my_wing = st.selectbox("내 날개 유형 (선택)", [""] + list(wing_types.keys()))
        my_instinct = st.selectbox("내 본능 유형 (선택)", [""] + list(instinct_variants.keys()))

    with tabs[1]:  # 상대방 정보 탭
        partner_name = st.text_input("상대방 이름/닉네임", "상대방")
        partner_gender = st.radio("상대방 성별", ["남성", "여성"])
        partner_type = st.selectbox("상대방 에니어그램 유형", list(enneagram_types.keys()))
        partner_wing = st.selectbox("상대방 날개 유형 (선택)", [""] + list(wing_types.keys()))
        partner_instinct = st.selectbox("상대방 본능 유형 (선택)", [""] + list(instinct_variants.keys()))

    relationship = st.selectbox("관계 유형", list(relationship_types.keys()))
    situation = st.text_area("상황 (예: 갈등 상황, 특별한 이벤트 등)", height=150)

generate_button = st.button("시나리오 생성 ➡️")

# 메인 영역 - 출력
scenario_output_area = st.empty()

if generate_button:
    full_text = generate_relationship_scenario(my_name, my_gender, my_type, my_wing, my_instinct, partner_name, partner_gender, partner_type, partner_wing, partner_instinct, relationship, situation)
    if full_text:
        scenario_output_area.markdown(markdown.markdown(full_text), unsafe_allow_html=True)