app.py

import streamlit as st
import streamlit.components.v1 as components
import os
import io
import concurrent.futures
import datetime
import zipfile
import json
from PIL import Image
from google import genai

# 로직 모듈 임포트
import logic_image
import logic_seo
import logic_tts
import utils
from config_style import STYLE_DEFINITIONS, THUMBNAIL_STRATEGIES

# ======================================================
# 1. 페이지 설정 및 CSS 디자인
# ======================================================
st.set_page_config(
    page_title="Nano Banana Studio Pro (Chapter Mode)",
    page_icon="🍌",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ------------------------------------------------------
# [핵심] 세션 상태 초기화 (챕터 관리용)
# ------------------------------------------------------
if 'chapters' not in st.session_state:
    st.session_state['chapters'] = [
        {
            "id": 0, 
            "text": "", 
            "results": [] 
        }
    ]
if 'next_chapter_id' not in st.session_state:
    st.session_state['next_chapter_id'] = 1

if 'final_audio' not in st.session_state:
    st.session_state['final_audio'] = None
if 'thumbnail_results' not in st.session_state:
    st.session_state['thumbnail_results'] = []
if 'thumbnail_text' not in st.session_state:
    st.session_state['thumbnail_text'] = ""
if 'seo_result' not in st.session_state:
    st.session_state['seo_result'] = None

# [CSS] 다크모드 + 오렌지 포인트 디자인
st.markdown("""
<style>
    @import url('https://cdn.jsdelivr.net/gh/orioncactus/pretendard/dist/web/static/pretendard.css');
    
    html, body, [class*="css"] {
        font-family: 'Pretendard', sans-serif;
    }
    .stApp {
        background-color: #121212;
        color: #E0E0E0;
    }
    
    /* 사이드바 */
    section[data-testid="stSidebar"] {
        background-color: #0A0A0A;
        border-right: 1px solid #333;
    }
    
    /* 입력창 가독성 개선 */
    .stTextInput input, .stTextArea textarea, .stSelectbox div[data-baseweb="select"] > div {
        background-color: #252525 !important;
        color: #FFFFFF !important;
        border: 1px solid #404040 !important;
        border-radius: 8px !important;
    }
    .stTextArea textarea[disabled] {
        background-color: #1E1E1E !important;
        color: #AAAAAA !important;
        border: 1px solid #333 !important;
    }
    ::placeholder {
        color: #AAAAAA !important;
        opacity: 1;
    }
    
    /* 탭 디자인 */
    button[data-baseweb="tab"] {
        font-size: 16px;
        font-weight: 700;
        color: #888;
    }
    button[data-baseweb="tab"][aria-selected="true"] {
        color: #FFC107 !important;
        background-color: transparent !important;
        border-bottom: 2px solid #FFC107 !important;
    }
    
    /* 버튼 디자인 */
    div.stButton > button[kind="primary"] {
        background: linear-gradient(135deg, #FFC107 0%, #FF9800 100%);
        color: #000000;
        border: none;
        padding: 14px 24px;
        font-size: 16px;
        font-weight: 800;
        border-radius: 8px;
        transition: all 0.3s ease;
    }
    div.stButton > button[kind="primary"]:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 15px rgba(255, 193, 7, 0.4);
    }
    div.stButton > button[kind="secondary"] {
        border-color: #555;
        color: #ddd;
    }

    h1, h2, h3, h4, h5 { color: #FFFFFF !important; }
    .highlight { color: #FFC107; }
    
    div[data-testid="stExpander"] {
        background-color: #1E1E1E;
        border: 1px solid #333;
        border-radius: 8px;
    }
    /* 챕터 구분선 */
    hr { border-color: #333; margin: 30px 0; }
</style>
""", unsafe_allow_html=True)

# ======================================================
# 2. 사이드바 (통합 설정)
# ======================================================
with st.sidebar:
    st.markdown("## ⚙️ 기본 설정 (Basic)")
    api_key = st.text_input("Google API Key", type="password", placeholder="API 키를 입력하세요...", label_visibility="collapsed")
    st.divider()

    st.markdown("### 🎨 Model Engine")
    st.radio("사용할 모델", ["🚀 Pro (Imagen 3)"], label_visibility="collapsed")
    # [고정] Pro 모델만 사용
    image_model_id = "gemini-3-pro-image-preview"
    text_model_id = "gemini-3-pro-preview"

    st.divider()

    st.markdown("### 📐 Canvas Ratio")
    ar_radio = st.radio("비율", ["16:9", "9:16"], label_visibility="collapsed", horizontal=True)
    aspect_ratio = ar_radio

    st.divider()
    
    st.markdown("### 🛠️ 고급 설정 (TTS & Split)")
    
    st.caption("TTS Model ID")
    tts_model_id = st.text_input("TTS ID", value="gemini-2.5-pro-preview-tts", label_visibility="collapsed")
    
    st.markdown("<br>", unsafe_allow_html=True)
    
    st.caption("Image Scene Split (이미지 생성용)")
    duration_per_scene = st.slider("장면당 시간(초)", 5, 600, 5, step=5)
    split_criteria = duration_per_scene * 8
    st.info(f"💡 {duration_per_scene}초 (약 {split_criteria}자) 단위로 장면을 나눕니다.")
    st.caption("※ TTS는 설정과 무관하게 500자 단위로 최적화됩니다.")
        
    st.divider()
    
    st.markdown("## 🎨 스타일 및 캐릭터 설정")
    st.caption("영상 전체의 분위기와 캐릭터를 결정합니다.")
    st.markdown("<br>", unsafe_allow_html=True)
    st.markdown("#### 🖌️ 화풍(Style) 선택")
    
    style_options = list(STYLE_DEFINITIONS.keys()) + ["직접 입력"]
    selected_style = st.selectbox("스타일 선택", style_options, label_visibility="collapsed")
    
    custom_style_input = ""
    current_prompt = ""
    
    if selected_style == "직접 입력":
        custom_style_input = st.text_input("스타일 프롬프트", placeholder="예: 지브리 스타일, 수채화풍")
        current_prompt = custom_style_input if custom_style_input else "(프롬프트를 입력해주세요)"
    else:
        style_value = STYLE_DEFINITIONS[selected_style]
        if isinstance(style_value, dict):
            current_prompt = style_value.get('prompt', str(style_value))
        else:
            current_prompt = str(style_value)

    st.markdown("#### 📜 적용될 프롬프트 (미리보기)")
    st.text_area("프롬프트 미리보기", value=current_prompt, height=150, disabled=True, label_visibility="collapsed")
    
    st.divider()
    
    st.markdown("#### 👤 캐릭터 참조 (Reference)")
    uploaded_file = st.file_uploader("이미지 업로드", type=["png", "jpg", "jpeg"], label_visibility="collapsed")
    reference_image = None
    if uploaded_file:
        reference_image = Image.open(uploaded_file)
        st.image(reference_image, caption="참조 이미지", use_container_width=True)


# ======================================================
# 3. 메인 콘텐츠
# ======================================================

st.title("Nano Banana <span class='highlight'>Chapter Studio</span>", anchor=False)
st.markdown("<div style='margin-top: -15px; color: #888; margin-bottom: 20px;'>AI Powered All-in-One Workspace</div>", unsafe_allow_html=True)

# 헬퍼 함수: 전체 스크립트 합치기 (컨텍스트용)
def get_full_script():
    full_text = []
    for ch in st.session_state['chapters']:
        if ch['text'].strip():
            full_text.append(ch['text'])
    return "\n\n".join(full_text)

tab1, tab2, tab3 = st.tabs(["🎬 챕터별 장면 생성", "📈 기획/SEO", "🎙️ 성우/TTS"])

# ----------------------------------------------------------------
# [TAB 1] 챕터별 장면/이미지 생성 (메인 기능)
# ----------------------------------------------------------------
with tab1:
    st.markdown("<br>", unsafe_allow_html=True)
    
    # 챕터 관리 루프
    chapters_to_remove = []
    
    for idx, chapter in enumerate(st.session_state['chapters']):
        with st.container():
            c_head1, c_head2 = st.columns([8, 1])
            with c_head1:
                st.markdown(f"#### 📖 Chapter {idx + 1}")
            with c_head2:
                if st.button("🗑️ 삭제", key=f"del_ch_{chapter['id']}"):
                    chapters_to_remove.append(idx)
            
            # 대본 입력
            new_text = st.text_area(
                f"Chapter {idx + 1} 대본",
                value=chapter['text'],
                height=150,
                placeholder="이 챕터의 내용을 입력하세요.",
                key=f"text_ch_{chapter['id']}",
                label_visibility="collapsed"
            )
            # 입력값 즉시 동기화
            st.session_state['chapters'][idx]['text'] = new_text
            
            # 생성 버튼
            if st.button(f"🚀 Chapter {idx + 1} 장면 생성하기", key=f"gen_ch_{chapter['id']}", type="primary"):
                if not api_key: st.error("API Key가 필요합니다.")
                elif not new_text.strip(): st.warning("대본을 입력해주세요.")
                else:
                    # [수정] 300자 제한 적용 (API 부하 감소)
                    full_context = new_text[:300]
                    
                    # 텍스트 분할
                    raw_scenes = logic_tts.split_text_smartly(new_text, limit=split_criteria)
                    scenes_text = raw_scenes # 제한 없음
                    
                    st.toast(f"📜 챕터 {idx+1}: {len(scenes_text)}개 장면 생성 시작!")
                    
                    temp_results = [None] * len(scenes_text)
                    
                    def run_scene_task(task_index, task_text):
                        client = genai.Client(api_key=api_key)
                        return logic_image.process_scene_task(
                            task_index,
                            {'text': task_text, 'full_script': full_context}, 
                            selected_style,
                            custom_style_input,
                            client,
                            text_model_id,
                            image_model_id,
                            aspect_ratio,
                            reference_image
                        )
                    
                    p_bar = st.progress(0, text=f"Chapter {idx+1} 그리는 중...")
                    
                    # [수정] 동시성 2로 제한 (안전 모드)
                    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
                        future_to_idx = {executor.submit(run_scene_task, i, t): i for i, t in enumerate(scenes_text)}
                        completed_cnt = 0
                        for future in concurrent.futures.as_completed(future_to_idx):
                            f_idx = future_to_idx[future]
                            try:
                                # [수정] timeout 설정 (45초) - 무한 대기 방지
                                _, prompt, img, video_prompt = future.result(timeout=45)
                                temp_results[f_idx] = {
                                    "prompt": prompt,
                                    "img_bytes": img,
                                    "script_text": scenes_text[f_idx],
                                    "video_prompt": video_prompt
                                }
                                completed_cnt += 1
                                p_bar.progress(completed_cnt / len(scenes_text), text=f"Scene {f_idx+1} 완료!")
                            except concurrent.futures.TimeoutError:
                                st.error(f"Scene {f_idx+1} 시간 초과 (45초) - 서버 응답 없음")
                                completed_cnt += 1 # 진행바는 넘겨줌
                            except Exception as e:
                                st.error(f"Scene {f_idx+1} 실패: {e}")
                                # 에러 내용을 더 자세히 보기 위해 화면에 출력
                                st.write(f"🔍 [Debug] Scene {f_idx+1} Error Detail: {type(e).__name__}, {str(e)}")
                                completed_cnt += 1
                    
                    st.session_state['chapters'][idx]['results'] = temp_results
                    p_bar.empty()
                    st.rerun()

            # 결과물 표시
            results = chapter.get('results', [])
            if results:
                st.info(f"✅ {len(results)}개의 장면이 생성되었습니다.")
                with st.expander(f"🖼️ Chapter {idx + 1} 결과물 확인하기 (클릭)", expanded=True):
                    for r_idx, item in enumerate(results):
                        if not item: continue
                        col_img, col_desc = st.columns([1, 2])
                        snippet = utils.make_scene_snippet(item.get("script_text", ""))
                        base_filename = f"ch{idx+1}_sc{r_idx+1:02d}_{snippet}"
                        
                        with col_img:
                            if item.get("img_bytes"):
                                st.image(item["img_bytes"], use_container_width=True)
                                d1, d2 = st.columns(2)
                                with d1:
                                    st.download_button("📥 이미지", data=item["img_bytes"], file_name=f"{base_filename}.png", mime="image/png", key=f"dl_img_{chapter['id']}_{r_idx}")
                                with d2:
                                    if item.get("video_prompt"):
                                        st.download_button("📄 프롬프트", data=item["video_prompt"], file_name=f"{base_filename}.txt", mime="text/plain", key=f"dl_txt_{chapter['id']}_{r_idx}")
                        with col_desc:
                            st.markdown(f"**Scene {r_idx + 1}**")
                            st.caption(item.get("script_text", ""))
                            st.text_area("비디오 프롬프트", item.get("video_prompt", ""), height=80, key=f"vp_{chapter['id']}_{r_idx}")
            st.divider()

    if chapters_to_remove:
        for rm_idx in sorted(chapters_to_remove, reverse=True):
            del st.session_state['chapters'][rm_idx]
        st.rerun()

    if st.button("➕ 챕터 추가하기", use_container_width=True):
        st.session_state['chapters'].append({
            "id": st.session_state['next_chapter_id'],
            "text": "",
            "results": []
        })
        st.session_state['next_chapter_id'] += 1
        st.rerun()

    # 통합 다운로드
    st.markdown("<br><br>", unsafe_allow_html=True)
    st.subheader("📦 전체 결과물 다운로드")
    
    total_files = 0
    all_zip_buffer = io.BytesIO()
    has_content = False
    
    with zipfile.ZipFile(all_zip_buffer, "w") as zip_file:
        for ch_idx, ch in enumerate(st.session_state['chapters']):
            for sc_idx, res in enumerate(ch.get('results', [])):
                if not res: continue
                snippet = utils.make_scene_snippet(res.get("script_text", ""))
                base_name = f"ch{ch_idx+1:02d}_sc{sc_idx+1:02d}_{snippet}"
                
                if res.get("img_bytes"):
                    zip_file.writestr(f"{base_name}.png", res["img_bytes"])
                    has_content = True
                    total_files += 1
                
                if res.get("video_prompt"):
                    zip_file.writestr(f"{base_name}.txt", res["video_prompt"])

    all_zip_buffer.seek(0)
    if has_content:
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M")
        st.download_button(
            label=f"⬇️ 전체 챕터 다운로드 (ZIP) - 총 {total_files}장",
            data=all_zip_buffer.getvalue(),
            file_name=f"full_project_{timestamp}.zip",
            mime="application/zip",
            type="primary",
            use_container_width=True
        )
    else:
        st.caption("생성된 결과물이 없습니다.")


# ----------------------------------------------------------------
# [TAB 2] 유튜브 기획 (SEO)
# ----------------------------------------------------------------
with tab2:
    st.markdown("<br>", unsafe_allow_html=True)
    st.markdown("#### 📈 전체 대본 SEO 분석")
    
    full_script = get_full_script()
    st.text_area("전체 대본 (자동 합침)", value=full_script, height=200, disabled=True)
    
    if st.button("✨ 전체 대본으로 SEO 기획서 생성", key="seo_btn", type="primary"):
        if not api_key: st.error("API Key 필요")
        elif not full_script.strip(): st.warning("대본 내용이 없습니다.")
        else:
            with st.spinner("Brain 모델이 전체 내용을 분석 중입니다..."):
                client = genai.Client(api_key=api_key)
                seo_result = logic_seo.generate_seo_content(client, text_model_id, full_script)
                st.session_state["seo_result"] = seo_result
                st.success("분석 완료!")
                
                c_seo1, c_seo2 = st.columns(2)
                with c_seo1:
                    st.markdown("##### 📌 추천 제목")
                    for t in seo_result['titles']: st.info(f"{t}")
                with c_seo2:
                    st.markdown("##### 🏷️ 추천 태그")
                    st.code(", ".join(seo_result['tags']))
                
                st.markdown("##### 📝 설명란")
                st.text_area("설명란 결과", seo_result['description'], height=200)

    # 썸네일 섹션
    st.divider()
    st.subheader("🖼️ 썸네일 생성")
    if st.session_state.get("seo_result"):
        # (썸네일 로직은 기존 코드 사용)
        pass 

# ----------------------------------------------------------------
# [TAB 3] AI 성우 (TTS)
# ----------------------------------------------------------------
with tab3:
    st.markdown("<br>", unsafe_allow_html=True)
    st.markdown("#### 🎙️ 전체 대본 TTS")
    
    # 1. 전체 대본 합치기 (기본값)
    full_script = get_full_script()
    
    # 2. 사용자 입력창
    tts_script = st.text_area("낭독할 대본 입력 (수정 가능)", value=full_script, height=200)
    
    voice_map = {
        "Charon (남성/다큐)": "Charon", "Puck (남성/쾌활)": "Puck",
        "Kore (여성/차분)": "Kore", "Fenrir (남성/강함)": "Fenrir",
        "Aoede (여성/높음)": "Aoede", "Orus (남성/중후)": "Orus"
    }
    
    v_sel = st.selectbox("성우 선택", list(voice_map.keys()), key="tts_voice_sel_main")
    voice_opt = voice_map[v_sel]

    if st.button("🎙️ 전체 녹음 시작", key="tts_full_btn_main", type="primary"):
        if not api_key: st.error("API Key 필요")
        elif not tts_script.strip(): st.warning("대본 내용이 없습니다.")
        else:
            chunks = logic_tts.split_text_smartly(tts_script, limit=500)
            audio_res = [None] * len(chunks)
            
            with st.status("녹음 진행 중...", expanded=True):
                with concurrent.futures.ThreadPoolExecutor() as executor:
                    def run_tts_task(task_index, chunk):
                        client = genai.Client(api_key=api_key)
                        return logic_tts.process_tts_task(task_index, chunk, client, tts_model_id, voice_opt)

                    f_map = {executor.submit(run_tts_task, i, c): i for i, c in enumerate(chunks)}
                    for f in concurrent.futures.as_completed(f_map):
                        idx, dat = f.result()
                        if isinstance(dat, bytes): 
                            audio_res[idx] = dat
                            st.write(f"✅ Part {idx+1} 완료")
            
            final_wav = logic_tts.merge_wav_bytes(audio_res)
            if final_wav:
                st.success("오디오 생성 완료!")
                st.audio(final_wav, format="audio/wav")
                st.download_button("다운로드 (WAV)", final_wav, "full_audio.wav", "audio/wav")