main.py

import streamlit as st
import os
from dotenv import load_dotenv
from datetime import date

from db import (
    insert_user_image, get_today_image, get_next_available_date, get_fallback_image,
    get_video_by_category, insert_chat_report, insert_user_info, get_user_info
)

from llm_chat import call_llm_ibm_with_history
from llm_summary import summarize_chat
from llm_image_analysis import generate_image_caption, enrich_image_description

from create_pdf import generate_pdf_from_report
from gtts import gTTS
import speech_recognition as sr
import io
import re



load_dotenv()


st.set_page_config(
    page_title="이음이",
    page_icon="favicon.ico"
)


if "page" not in st.session_state:
    st.session_state.page = "home"

if "user_id" not in st.session_state or st.session_state.user_id is None:
    st.error("로그인이 필요합니다.")
    st.stop()

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

if "mic_status" not in st.session_state:
    st.session_state.mic_status = False


st.sidebar.header("👤 보호자 설정")
uploaded_file = st.sidebar.file_uploader("이미지 업로드", type=["jpg", "jpeg", "png"])
desc = st.sidebar.text_input("사진 설명 입력 (선택)", placeholder="예: 가족 여행 사진, 친구와의 산책")
submit = st.sidebar.button("사진 저장")

if uploaded_file and submit:
    save_path = f"static/user/{st.session_state.user_id}/{uploaded_file.name}"
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    with open(save_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    if not desc:
        with st.spinner("🧠 사진을 분석 중입니다..."):
            caption = generate_image_caption(save_path)

            desc = enrich_image_description(caption)

            st.success(f"📸 자동 설명 생성: '{desc}'")

    scheduled_date = get_next_available_date(st.session_state.user_id)
    insert_user_image(
        user_id=st.session_state.user_id,
        image_path=save_path,
        scheduled_date=scheduled_date.isoformat(),
        description=desc
    )
    st.sidebar.success("이미지가 저장되었습니다.")


with st.sidebar.expander("🧓 환자 정보 등록"):
    if "patient_registered" not in st.session_state:
        st.session_state.patient_registered = False

    if not st.session_state.patient_registered:
        st.markdown("### 📋 환자 정보 입력", unsafe_allow_html=True)

        name = st.text_input("이름", key="patient_name")
        age = st.number_input("나이", min_value=0, max_value=120, step=1, key="patient_age")
        triggers = st.text_area("금지 단어(트라우마)", key="patient_triggers")
        gender = st.selectbox("성별", ["남성", "여성"], key="patient_gender")
        dementia_level = st.selectbox("치매 정도", ["경도", "경증", "중등도", "중증"], key="patient_dementia")

        if st.button("✅ 환자 정보 등록"):
            if name and age:
                trigger_list = [t.strip() for t in triggers.split(",")] if triggers else []
                trigger_str = ",".join(trigger_list)

                user_id = st.session_state.user_id

                success = insert_user_info(user_id, name, age, trigger_str, gender, dementia_level)
                if success:
                    st.session_state.patient_registered = True 
                    st.success(f"환자 '{name}' 정보가 등록되었습니다.")
                else:
                    st.warning("이미 해당 사용자에 대한 환자 정보가 등록되어 있습니다.")
            else:
                st.warning("이름과 나이는 필수 항목입니다.")
    else:
        st.success("✅ 환자 정보가 이미 등록되었습니다.")

st.sidebar.markdown("---")
st.sidebar.subheader("🖍️ 시각 설정")
font_size = st.sidebar.slider("글자 크기", 12, 48, 20)
theme_color = st.sidebar.color_picker("포인트 색상", "#4CAF50")
high_contrast_mode = st.sidebar.checkbox("👵 어르신용 고대비 모드", value=False)


st.markdown(f"""
<style>
  .hero {{
    min-height: 70vh;
    display: flex;
    align-items: center;
    justify-content: center;
    padding: 40px 16px;
    background: linear-gradient(135deg, {theme_color}1a 0%, #ffffff 60%);
  }}
  .home-card {{
    background: #ffffff;
    max-width: 640px;
    width: 100%;
    border-radius: 20px;
    box-shadow: 0 8px 28px rgba(0,0,0,0.08);
    padding: 48px 40px;
    text-align: center;
  }}
  .home-title {{
    font-size: 32px;
    margin: 0 0 8px 0;
  }}
  .home-sub {{
    font-size: 16px;
    color: #5f6368;
    margin-bottom: 28px;
  }}
  .pill {{
    display: inline-block;
    background: {theme_color}1a;
    color: {theme_color};
    border-radius: 999px;
    padding: 6px 12px;
    font-size: 12px;
    font-weight: 600;
    margin-bottom: 16px;
  }}
  .feature-row {{
    display: grid;
    grid-template-columns: 1fr 1fr 1fr;
    gap: 12px;
    margin-top: 16px;
  }}
  .feature {{
    background: #fafafa;
    border: 1px solid #f0f0f0;
    border-radius: 12px;
    padding: 12px;
    font-size: 14px;
  }}
  @media (max-width: 640px) {{
    .feature-row {{ grid-template-columns: 1fr; }}
  }}
</style>
""", unsafe_allow_html=True)


if st.session_state.page == "home":
    st.markdown(
        """
        <div class="hero">
          <div class="home-card">
            <div class="pill">이음이와 함께 하는 회상 대화</div>
            <h2 class="home-title">👋 환영합니다!</h2>
            <p class="home-sub">오늘의 사진과 함께, 따뜻한 이야기를 시작해볼까요?</p>
            <div class="feature-row">
              <div class="feature">📷 오늘의 회상 이미지</div>
              <div class="feature">🗣️ 음성(STT) 대화 지원</div>
              <div class="feature">📝 대화 요약 & 리포트</div>
            </div>
          </div>
        </div>
        """,
        unsafe_allow_html=True
    )


    _, center, _ = st.columns([1, 2, 1])
    with center:
        start = st.button("시작하기", use_container_width=True)
        if start:
            st.session_state.page = "chat"
            st.rerun()



elif st.session_state.page == "chat":
    st.title("📷 오늘의 회상 이미지")
    user_id = st.session_state.user_id
    today = date.today().isoformat()
    image_data = get_today_image(user_id, today)
    fallback = get_fallback_image()
    image_description = ""

    if image_data:
        st.image(image_data["image_path"])
        image_description = image_data["description"]
    elif fallback:
        st.image(fallback["image_path"], caption=fallback["description"])
        st.info("기본 이미지가 대신 표시됩니다.")
        image_description = fallback["description"]


    if not st.session_state.chat_history:

        user_info = get_user_info(user_id)

        if user_info:
            patient_name = user_info["name"]
            patient_age = user_info["age"]
            patient_triggers = user_info["trigger_elements"]
            patient_gender = user_info["gender"]
            patient_dementia = user_info["dementia_level"]
        else:
            patient_name = "미등록"
            patient_age = "미등록"
            patient_triggers = "미등록"
            patient_gender = "미등록"
            patient_dementia = "미등록"

        st.session_state.chat_history.append({
            "role": "user",
            "content": (
                f"이 사진은 '{image_description}'입니다. 전송 완료.\n"
                f"\n**[환자 정보]**\n"
                f"- 이름: {patient_name}\n"
                f"- 나이: {patient_age}세\n"
                f"- 성별: {patient_gender}\n"
                f"- 치매 진행 정도: {patient_dementia}\n"
                f"- 트라우마 요소: {patient_triggers}\n"
                ),
            "display": False
        })
        response = call_llm_ibm_with_history(st.session_state.chat_history)
        st.session_state.chat_history.append({"role": "assistant", "content": response})
        st.session_state.llm_initialized = True
    
    elif not st.session_state.get("llm_initialized", False):

        response = call_llm_ibm_with_history(st.session_state.chat_history)
        st.session_state.chat_history.append({"role": "assistant", "content": response})
        st.session_state.llm_initialized = True


    for i, msg in enumerate(st.session_state.chat_history):
        if msg.get("display", True):
            with st.chat_message(msg["role"]):
                content = msg["content"]
                if high_contrast_mode:
                    background = "#000000" if msg["role"] == "user" else "#FFD700"  # 검정 vs 노랑
                    text_color = "#FFFFFF" if msg["role"] == "user" else "#000000"  # 흰 vs 검
                else:
                    background = "#f0f0f0" if msg["role"] == "user" else "#e0f7e9"
                    text_color = "black" if msg["role"] == "user" else theme_color

                styled_message = f"""
                    <div style='
                        font-size: {font_size}px;
                        color: {text_color};
                        background-color: {background};
                        padding: 14px;
                        border-radius: 10px;
                        margin-bottom: 10px;
                        width: fit-content;
                    '>
                        {content}
                    </div>
                """
                st.markdown(styled_message, unsafe_allow_html=True)
                
                if msg["role"] == "assistant":
                    if st.button("🔊", key=f"tts_{i}"):
                        tts = gTTS(msg["content"], lang='ko')
                        mp3_fp = io.BytesIO()
                        tts.write_to_fp(mp3_fp)
                        mp3_fp.seek(0)
                        st.audio(mp3_fp, format='audio/mp3')

    col1, col2 = st.columns([9, 1])
    with col1:
        user_input = st.chat_input("메시지를 입력하세요")
    with col2:
        if st.button("🎙️", key="stt_button"):
            st.session_state.mic_status = True


    if st.session_state.mic_status:
        st.markdown(
            """
            <div style='
                background-color: #e6f4ea;
                border: 1px solid #b2d8c5;
                border-radius: 10px;
                padding: 10px;
                margin-top: 5px;
                font-size: 16px;
                color: #2e7d32;
                width: fit-content;
            '>
                🎙️ 듣고 있어요... 말을 해주세요!
            </div>
            """,
            unsafe_allow_html=True,
        )

        recognizer = sr.Recognizer()
        with sr.Microphone() as source:
            try:
                audio = recognizer.listen(source, timeout=5)
                text = recognizer.recognize_google(audio, language='ko-KR')
                st.success(f"📝 변화된 텍스트: {text}")

                st.session_state.chat_history.append({"role": "user", "content": text})
                response = call_llm_ibm_with_history(st.session_state.chat_history)
                st.session_state.chat_history.append({"role": "assistant", "content": response})


                st.session_state.mic_status = False
                st.rerun()
            except sr.WaitTimeoutError:
                st.error("⏱️ 너무 오래 기다렸어요. 다시 시도해 주세요.")
            except sr.UnknownValueError:
                st.error("🌀 음성을 인식하지 못했어요.")
            except sr.RequestError as e:
                st.error(f"❌ STT API 오류: {e}")

    if user_input:
        if "운동" in user_input and "영상" in user_input:
            video_url = get_video_by_category("운동", index=date.today().day % 12)
            if video_url:
                st.session_state.chat_history.append({
                    "role": "assistant",
                    "content": "오늘의 운동 영상이에요!"
                })
                st.video(video_url)
            else:
                st.warning("운동 영상을 찾을 수 없어요.")

        elif "노래" in user_input or "음악" in user_input:
            video_url = get_video_by_category("음악")
            st.write("🔗 가져온 URL:", video_url)
            if video_url:
                st.session_state.chat_history.append({
                    "role": "assistant",
                    "content": "마음이 편해지는 음악이에요!"
                })
                st.video(video_url)
            else:
                st.warning("음악 영상을 찾을 수 없어요.")
                
        elif any(kw in user_input for kw in ["다른 사진", "사진 바꿔", "다른 걸로", "사진 바꾸자"]):
            fallback = get_fallback_image()
            st.image(fallback["image_path"], caption=fallback["description"])
            st.info("다른 기본 이미지를 보여드릴게요!")
            image_description = fallback["description"]
            st.session_state.chat_history = []
            st.session_state.chat_history.append({
                "role": "user",
                "content": f"이 사진은 '{image_description}'입니다. 이 사진에 대해 이야기해줘."
            })

        else:
            with st.spinner("답변을 생성 중입니다..."):
                st.session_state.chat_history.append({"role": "user", "content": user_input})
                response = call_llm_ibm_with_history(st.session_state.chat_history)
                st.session_state.chat_history.append({"role": "assistant", "content": response})
                st.rerun()
    
    if st.session_state.page == "chat":
        if not st.session_state.get("summary_done", False):
            st.markdown("---")
            if st.button("✅ 오늘 대화 종료 및 요약 분석"):
                with st.spinner("📄 요약 및 분석 중입니다..."):
                    full_chat_text = "\n".join([
                        f"{msg['role']}: {msg['content']}" for msg in st.session_state.chat_history
                    ])
                    result = summarize_chat(full_chat_text)
                    st.markdown("### 📦 리포트 미리보기:")
                    st.code(result)

                    match_summary = re.search(r"1\.\s*대화 요약\s*(.*?)\s*2\.\s*회상 평가", result, re.DOTALL)

                    match_memo = re.search(r"2\.\s*회상 평가\s*(.*)", result, re.DOTALL)

                    if match_summary and match_memo:
                        summary_block = "1. 대화 요약\n" + match_summary.group(1).strip()
                        memo_block = "2. 회상 평가\n" + match_memo.group(1).strip()
                    else:
                        st.error("❌ LLM 응답에서 요약/분석 항목을 찾지 못했습니다.")
                        st.code(result)
                        st.stop()


                    if image_data and "id" in image_data:
                        insert_chat_report(
                            user_id=user_id,
                            image_id=image_data["id"],
                            fallback_image_id=None,
                            chat_summary=summary_block,
                            memo=memo_block
                        )
                    elif fallback and "id" in fallback:
                        insert_chat_report(
                            user_id=user_id,
                            image_id=None,
                            fallback_image_id=fallback["id"],
                            chat_summary=summary_block,
                            memo=memo_block
                        )
                    else:
                        st.warning("이미지 ID가 없어 리포트를 저장하지 못했습니다.")

                    st.session_state["last_summary"] = summary_block
                    st.session_state["last_memo"] = memo_block
                    st.session_state["summary_done"] = True

if st.session_state.get("summary_done", False):
    if st.button("📄 PDF로 저장하기"):
        from create_pdf import generate_pdf_from_report
        generate_pdf_from_report(user_id)
        st.success("✅ PDF 저장 완료!")