import streamlit as st
from google import genai
from gtts import gTTS
from PIL import Image
import io

st.set_page_config(page_title="마법의 책 읽기 카메라", layout="centered", page_icon="📖")

# Custom CSS for Kids UI
st.markdown("""
    <style>
    .main { background-color: #FFF8F0; }
    h1 { color: #FF6B6B; font-family: 'Comic Sans MS', 'Chalkboard SE', sans-serif; text-align: center; }
    h3 { color: #4ECDC4; text-align: center; }
    .stCameraInput { border-radius: 15px; border: 4px dashed #FFD93D; padding: 10px; }
    .stAudio { margin-top: 20px; }
    .story-box { 
        background-color: white; 
        padding: 20px; 
        border-radius: 15px; 
        box-shadow: 0 4px 8px rgba(0,0,0,0.1); 
        border-left: 8px solid #FFD93D;
        font-size: 20px;
        line-height: 1.6;
        color: #333;
    }
    </style>
""", unsafe_allow_html=True)

st.title("📖 마법의 책 읽기 카메라 📷")
st.markdown("### 책을 찰칵! 찍어주면 마법사 요정이 글을 읽어줘요 ✨")

import os

# Sidebar for configuration
with st.sidebar:
    st.header("⚙️ 부모님 설정 메뉴")
    st.info("API 키를 한 번만 입력하고 '키 저장하기'를 눌러주세요!")
    
    # Check if key is in environment variables (for Hugging Face Secrets)
    env_key = os.environ.get("GEMINI_API_KEY", "")
    
    with st.form("api_key_form"):
        gemini_api_key = st.text_input("Gemini API Key", type="password", placeholder="AIzaSy...", value=env_key)
        submitted = st.form_submit_button("키 저장하기 (엔터 역할)")
        
    if submitted:
        st.success("키가 성공적으로 적용되었습니다! 이제 사진을 찍어보세요.")

# Initialize session state for storing the current generated text and audio
if "story_text" not in st.session_state:
    st.session_state.story_text = ""
if "audio_bytes" not in st.session_state:
    st.session_state.audio_bytes = None

# Camera or File Upload Input
tab1, tab2 = st.tabs(["📸 카메라 촬영", "📁 사진 업로드 (모바일 추천)"])

with tab1:
    camera_image = st.camera_input("카메라로 책을 찍어보세요!")
with tab2:
    uploaded_image = st.file_uploader("스마트폰 기본 카메라로 찍은 사진을 올려주세요", type=['jpg', 'jpeg', 'png'])

final_image = camera_image if camera_image else uploaded_image

if final_image is not None:
    if not gemini_api_key:
        st.warning("⚠️ 왼쪽 메뉴에서 Gemini API 키를 먼저 입력하고 저장해주세요!")
    else:
        st.success("사진 찰칵! 요정에게 사진을 전달했어요. 잠시만 기다려주세요 🧚‍♀️")
        
        # This prevents re-running the API call if the image hasn't changed
        if "last_image" not in st.session_state or st.session_state.last_image != final_image.getvalue():
            st.session_state.last_image = final_image.getvalue()
            
            with st.spinner("마법사 요정이 책을 읽고 있어요... 뾰로롱! ✨"):
                try:
                    # 1. OCR + Contextualization with Gemini 2.5 Flash
                    client = genai.Client(api_key=gemini_api_key)
                    image = Image.open(final_image)
                    
                    # 모바일 고해상도 사진 용량 줄이기 (크래시 방지)
                    image.thumbnail((1024, 1024))
                    
                    prompt = """
                    이 이미지에 있는 책의 내용을 정확하게 읽어줘.
                    아이에게 동화책을 읽어주는 것처럼 친절하고 다정한 말투로 자연스럽게 바꿔서 들려줘.
                    내용에 어울리는 귀여운 이모지도 중간중간 넣어줘.
                    이미지에 텍스트가 없다면 "앗, 글자가 안 보여요! 다시 한 번 찍어줄래?" 라고 말해줘.
                    """
                    
                    response = client.models.generate_content(
                        model='gemini-2.5-flash',
                        contents=[prompt, image]
                    )
                    
                    st.session_state.story_text = response.text

                    
                    # 2. Text-to-Speech with gTTS
                    tts = gTTS(text=st.session_state.story_text, lang='ko', slow=False)
                    audio_fp = io.BytesIO()
                    tts.write_to_fp(audio_fp)
                    st.session_state.audio_bytes = audio_fp.getvalue()
                    
                except Exception as e:
                    st.error(f"앗! 마법사 요정이 실수했어요: {e}")
                    st.session_state.story_text = ""
                    st.session_state.audio_bytes = None

# Display Output
if st.session_state.story_text and st.session_state.audio_bytes:
    st.markdown("### 🎙️ 마법사 요정의 목소리를 들어보세요!")
    st.audio(st.session_state.audio_bytes, format='audio/mp3', autoplay=True)
    
    st.markdown("### 📝 요정이 읽어준 내용")
    st.markdown(f'<div class="story-box">{st.session_state.story_text}</div>', unsafe_allow_html=True)