import streamlit as st
import requests
import base64
from PIL import Image
from io import BytesIO

# ───────────────────────────────
# CONFIG
# ───────────────────────────────
st.set_page_config(page_title="동화 삽화 생성기", layout="wide")
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-Kontext-dev"
headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}

# ───────────────────────────────
# UTILS
# ───────────────────────────────
def query_flux(prompt, image_bytes=None):
    inputs = {"prompt": prompt}
    if image_bytes:
        inputs["image"] = base64.b64encode(image_bytes).decode("utf-8")

    response = requests.post(API_URL, headers=headers, json=inputs)
    if response.status_code == 200:
        image_data = base64.b64decode(response.json()["image"])
        return Image.open(BytesIO(image_data))
    else:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None

# ───────────────────────────────
# UI
# ───────────────────────────────
st.title("📚 동화 삽화 생성기 (FLUX.1 Kontext)")
st.markdown("각 장면의 설명을 입력하면, 이야기 흐름에 맞는 삽화를 자동으로 생성합니다.")

with st.expander("등장인물 및 스타일 정보 (선택)"):
    character_prompt = st.text_area("등장인물 설명 (ex. 빨간 망토를 쓴 소녀, 회색 늑대 등)", height=100)
    reference_image = st.file_uploader("참조 이미지 업로드 (선택)", type=["jpg", "png"])

st.markdown("### 📝 동화 내용을 5개 장면으로 나누어 입력하세요")

scene_prompts = []
for i in range(5):
    text = st.text_area(f"장면 {i+1} 설명", key=f"scene_{i}", height=80)
    scene_prompts.append(text)

if st.button("🎨 삽화 생성하기"):
    if not any(scene_prompts):
        st.warning("최소 하나 이상의 장면 설명이 필요합니다.")
    else:
        ref_bytes = reference_image.read() if reference_image else None

        with st.spinner("이미지를 생성 중입니다..."):
            cols = st.columns(5)
            for i, prompt in enumerate(scene_prompts):
                if not prompt.strip():
                    continue

                full_prompt = f"{character_prompt}. {prompt}" if character_prompt else prompt
                img = query_flux(full_prompt, ref_bytes)
                if img:
                    cols[i].image(img, caption=f"장면 {i+1}")