app.py

# storyverse_weaver_streamlit/app_st.py (example name)
import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import os
import time
import random

# --- Assuming your core logic is in a sibling 'core' directory ---
# You might need to adjust sys.path if running locally vs. deployed
# import sys
# sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # If core is one level up

from core.llm_services import initialize_text_llms, is_gemini_text_ready, is_hf_text_ready, generate_text_gemini, generate_text_hf
from core.image_services import initialize_image_llms, is_dalle_ready, is_hf_image_api_ready, generate_image_dalle, generate_image_hf_model, ImageGenResponse
from core.story_engine import Story, Scene # Your existing Story and Scene classes
from prompts.narrative_prompts import get_narrative_system_prompt, format_narrative_user_prompt
from prompts.image_style_prompts import STYLE_PRESETS, COMMON_NEGATIVE_PROMPTS, format_image_generation_prompt
from core.utils import basic_text_cleanup

# --- Initialize Services ONCE ---
# Use Streamlit's caching for resource-heavy initializations if they don't depend on session state
@st.cache_resource # Caches the result across sessions/reruns if inputs don't change
def load_ai_services():
    print("--- Initializing AI Services (Streamlit Cache Resource) ---")
    initialize_text_llms()
    initialize_image_llms()
    # Return status flags to be stored in session_state or used directly
    return {
        "gemini_text_ready": is_gemini_text_ready(),
        "hf_text_ready": is_hf_text_ready(),
        "dalle_image_ready": is_dalle_ready(),
        "hf_image_ready": is_hf_image_api_ready()
    }

ai_services_status = load_ai_services()

# --- Application Configuration (Models, Defaults) ---
# (Similar logic to your Gradio app.py for populating TEXT_MODELS, IMAGE_PROVIDERS etc.)
TEXT_MODELS = {}
UI_DEFAULT_TEXT_MODEL_KEY = None
# ... (Populate based on ai_services_status["gemini_text_ready"], ai_services_status["hf_text_ready"]) ...
if ai_services_status["gemini_text_ready"]: TEXT_MODELS["✨ Gemini 1.5 Flash (Narrate)"] = {"id": "gemini-1.5-flash-latest", "type": "gemini"} # etc.
if ai_services_status["hf_text_ready"]: TEXT_MODELS["Mistral 7B (Narrate via HF)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"} # etc.
if TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0] # Simplified default

IMAGE_PROVIDERS = {}
UI_DEFAULT_IMAGE_PROVIDER_KEY = None
# ... (Populate based on ai_services_status["dalle_image_ready"], ai_services_status["hf_image_ready"]) ...
if ai_services_status["dalle_image_ready"]: IMAGE_PROVIDERS["🖼️ DALL-E 3"] = "dalle_3" #etc.
if ai_services_status["hf_image_ready"]: IMAGE_PROVIDERS["🎡 HF - SDXL Base"] = "hf_sdxl_base" #etc.
if IMAGE_PROVIDERS: UI_DEFAULT_IMAGE_PROVIDER_KEY = list(IMAGE_PROVIDERS.keys())[0] # Simplified default


# --- Helper: Placeholder Image (can be same as before) ---
@st.cache_data # Cache placeholder images
def create_placeholder_image_st(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
    # ... (same PIL logic as before) ...
    img = Image.new('RGB', size, color=color); draw = ImageDraw.Draw(img)
    try: font_path = "arial.ttf" if os.path.exists("arial.ttf") else None 
    except: font_path = None 
    try: font = ImageFont.truetype(font_path, 40) if font_path else ImageFont.load_default()
    except IOError: font = ImageFont.load_default()
    if hasattr(draw, 'textbbox'): bbox = draw.textbbox((0,0), text, font=font); tw, th = bbox[2]-bbox[0], bbox[3]-bbox[1]
    else: tw, th = draw.textsize(text, font=font) 
    draw.text(((size[0]-tw)/2, (size[1]-th)/2), text, font=font, fill=text_color); return img

# --- Initialize Session State ---
if 'story_object' not in st.session_state:
    st.session_state.story_object = Story()
if 'current_log' not in st.session_state:
    st.session_state.current_log = ["Welcome to StoryVerse Weaver (Streamlit Edition)!"]
if 'latest_scene_image' not in st.session_state:
    st.session_state.latest_scene_image = None
if 'latest_scene_narrative' not in st.session_state:
    st.session_state.latest_scene_narrative = "Describe your first scene to begin!"
if 'processing_scene' not in st.session_state:
    st.session_state.processing_scene = False

# --- Page Configuration (Do this ONCE at the top) ---
st.set_page_config(
    page_title="✨ StoryVerse Weaver ✨",
    page_icon="🌌",
    layout="wide", # "wide" or "centered"
    initial_sidebar_state="expanded" # "auto", "expanded", "collapsed"
)

# --- Custom CSS for Dark Theme "WOW" ---
# (You'd inject this using st.markdown(..., unsafe_allow_html=True) or a separate CSS file)
streamlit_omega_css = """
<style>
    /* Base dark theme */
    body { color: #D0D0E0; background-color: #0F0F1A; }
    .stApp { background-color: #0F0F1A; }
    h1, h2, h3, h4, h5, h6 { color: #C080F0; }
    .stTextInput > div > div > input, .stTextArea > div > div > textarea, .stSelectbox > div > div > select {
        background-color: #2A2A4A; color: #E0E0FF; border: 1px solid #4A4A6A; border-radius: 8px;
    }
    .stButton > button {
        background: linear-gradient(135deg, #7F00FF 0%, #E100FF 100%) !important;
        color: white !important; border: none !important; border-radius: 8px !important;
        padding: 0.5em 1em !important; font-weight: 600 !important;
        box-shadow: 0 4px 8px rgba(0,0,0,0.15) !important;
    }
    .stButton > button:hover { transform: scale(1.03) translateY(-1px); box-shadow: 0 8px 16px rgba(127,0,255,0.3) !important; }
    /* Add more specific styles for sidebar, expanders, image display etc. */
    .main .block-container { padding-top: 2rem; padding-bottom: 2rem; padding-left: 3rem; padding-right: 3rem; max-width: 1400px; margin: auto;}
    .stImage > img { border-radius: 12px; box-shadow: 0 6px 15px rgba(0,0,0,0.25); max-height: 600px;}
    .stExpander { background-color: #1A1A2E; border: 1px solid #2A2A4A; border-radius: 12px; margin-bottom: 1em;}
    .stExpander header { font-size: 1.1em; font-weight: 500; color: #D0D0FF;}
    .important-note { background-color: rgba(127,0,255,0.1); border-left: 5px solid #7F00FF; padding: 15px; margin-bottom:20px; color: #E0E0FF; border-radius: 6px;}
</style>
"""
st.markdown(streamlit_omega_css, unsafe_allow_html=True)


# --- Main App UI & Logic ---
st.markdown("<div align='center'><h1>✨ StoryVerse Weaver ✨</h1>\n<h3>Craft Immersive Multimodal Worlds with AI</h3></div>", unsafe_allow_html=True)
st.markdown("<div class='important-note'><strong>Welcome, Worldsmith!</strong> Describe your vision, choose your style, and let Omega help you weave captivating scenes with narrative and imagery. Ensure API keys (<code>STORYVERSE_...</code>) are correctly set in your environment/secrets!</div>", unsafe_allow_html=True)


# --- Sidebar for Inputs & Configuration ---
with st.sidebar:
    st.header("🎨 Scene Weaver Panel")
    
    with st.form("scene_input_form"):
        scene_prompt_text = st.text_area(
            "Scene Vision (Description, Dialogue, Action):", 
            height=200,
            placeholder="e.g., Amidst swirling cosmic dust, Captain Eva pilots her damaged starfighter..."
        )
        
        st.subheader("Visual Style")
        col_style1, col_style2 = st.columns(2)
        with col_style1:
            image_style_dropdown = st.selectbox("Style Preset:", options=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), index=0)
        with col_style2:
            artist_style_text = st.text_input("Artistic Inspiration (Optional):", placeholder="e.g., Moebius")
        
        negative_prompt_text = st.text_area("Exclude from Image (Negative Prompt):", value=COMMON_NEGATIVE_PROMPTS, height=100)

        with st.expander("⚙️ Advanced AI Configuration", expanded=False):
            text_model_key = st.selectbox("Narrative AI Engine:", options=list(TEXT_MODELS.keys()), index=0 if UI_DEFAULT_TEXT_MODEL_KEY in TEXT_MODELS else (list(TEXT_MODELS.keys()).index(UI_DEFAULT_TEXT_MODEL_KEY) if UI_DEFAULT_TEXT_MODEL_KEY else 0) )
            image_provider_key = st.selectbox("Visual AI Engine:", options=list(IMAGE_PROVIDERS.keys()), index=0 if UI_DEFAULT_IMAGE_PROVIDER_KEY in IMAGE_PROVIDERS else (list(IMAGE_PROVIDERS.keys()).index(UI_DEFAULT_IMAGE_PROVIDER_KEY) if UI_DEFAULT_IMAGE_PROVIDER_KEY else 0) )
            narrative_length = st.selectbox("Narrative Detail:", options=["Short (1 paragraph)", "Medium (2-3 paragraphs)", "Detailed (4+ paragraphs)"], index=1)
            image_quality = st.selectbox("Image Detail/Style:", options=["Standard", "High Detail", "Sketch Concept"], index=0)
        
        submit_scene_button = st.form_submit_button("🌌 Weave This Scene!", use_container_width=True, type="primary", disabled=st.session_state.processing_scene)

    if st.button("🎲 Surprise Me!", use_container_width=True, disabled=st.session_state.processing_scene):
        sur_prompt, sur_style, sur_artist = surprise_me_func() # Assuming this is defined as before
        # Need to update the actual input widget values; Streamlit doesn't directly map outputs to inputs like Gradio's Examples
        # This requires a more involved way to update widget states, or just display the suggestion.
        # For simplicity, we'll just show what it would generate. A real app might use st.experimental_rerun or callbacks.
        st.info(f"Surprise Idea: Prompt='{sur_prompt}', Style='{sur_style}', Artist='{sur_artist}'\n(Copy these into the fields above!)")


    if st.button("🗑️ New Story", use_container_width=True, disabled=st.session_state.processing_scene):
        st.session_state.story_object = Story()
        st.session_state.current_log = ["Story Cleared. Ready for a new verse!"]
        st.session_state.latest_scene_image = None
        st.session_state.latest_scene_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene!"
        st.experimental_rerun() # Rerun the script to refresh the UI

    with st.expander("🔧 AI Services Status", expanded=False):
        text_llm_ok, image_gen_ok = (ai_services_status["gemini_text_ready"] or ai_services_status["hf_text_ready"]), \
                                    (ai_services_status["dalle_image_ready"] or ai_services_status["hf_image_ready"])
        if not text_llm_ok and not image_gen_ok: st.error("CRITICAL: NO AI SERVICES CONFIGURED.")
        else:
            if text_llm_ok: st.success("Text Generation Service(s) Ready.")
            else: st.warning("Text Generation Service(s) NOT Ready.")
            if image_gen_ok: st.success("Image Generation Service(s) Ready.")
            else: st.warning("Image Generation Service(s) NOT Ready.")


# --- Main Display Area ---
st.markdown("---")
st.markdown("### 🖼️ **Your Evolving StoryVerse**", unsafe_allow_html=True) # For potential custom class via CSS

if st.session_state.processing_scene:
    st.info("🌌 Weaving your scene... Please wait.")
    # Could use st.spinner("Weaving your scene...")

# Display Latest Scene
if st.session_state.latest_scene_image or st.session_state.latest_scene_narrative != "Describe your first scene to begin!":
    st.subheader("🌠 Latest Scene")
    if st.session_state.latest_scene_image:
        st.image(st.session_state.latest_scene_image, use_column_width=True, caption="Latest Generated Image")
    st.markdown(st.session_state.latest_scene_narrative, unsafe_allow_html=True)
    st.markdown("---")


# Display Story Scroll (Gallery)
if st.session_state.story_object and st.session_state.story_object.scenes:
    st.subheader("📚 Story Scroll")
    # Streamlit doesn't have a direct "Gallery" like Gradio. We display images in columns.
    num_columns = 3
    cols = st.columns(num_columns)
    scenes_for_gallery = st.session_state.story_object.get_all_scenes_for_gallery_display() # Ensure this returns (PIL.Image or None, caption)

    for i, (img, caption) in enumerate(scenes_for_gallery):
        with cols[i % num_columns]:
            if img:
                st.image(img, caption=caption if caption else f"Scene {i+1}", use_column_width=True)
            elif caption: # If no image but caption (e.g. error)
                st.caption(caption) # Display caption as text
else:
    st.caption("Your story scroll is empty. Weave your first scene!")


# Interaction Log
with st.expander("⚙️ Interaction Log", expanded=False):
    st.markdown("\n\n".join(st.session_state.current_log), unsafe_allow_html=True)


# --- Logic for Form Submission ---
if submit_scene_button and scene_prompt_text.strip(): # Check if form submitted and prompt is not empty
    st.session_state.processing_scene = True
    st.session_state.current_log.append(f"**🚀 New Scene Request - {time.strftime('%H:%M:%S')}**")
    st.experimental_rerun() # Rerun to show "processing" state and disable button

    # ---- This is where the main generation logic happens ----
    # It's similar to add_scene_to_story_orchestrator but updates session_state
    
    # 1. Generate Narrative
    current_narrative = f"Narrative Error: Init failed for '{scene_prompt_text[:30]}...'"
    text_model_info = TEXT_MODELS.get(text_model_key)
    if text_model_info and text_model_info["type"] != "none":
        system_p = get_narrative_system_prompt("default")
        prev_narrative = st.session_state.story_object.get_last_scene_narrative()
        user_p = format_narrative_user_prompt(scene_prompt_text, prev_narrative)
        st.session_state.current_log.append(f"  Narrative: Using {text_model_key} ({text_model_info['id']}).")
        text_response = None
        if text_model_info["type"] == "gemini" and ai_services_status["gemini_text_ready"]: text_response = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
        elif text_model_info["type"] == "hf_text" and ai_services_status["hf_text_ready"]: text_response = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if narrative_length.startswith("Detailed") else 400)
        
        if text_response and text_response.success: current_narrative = basic_text_cleanup(text_response.text); st.session_state.current_log.append(f"  Narrative: Success.")
        elif text_response: current_narrative = f"**Narrative Error ({text_model_key}):** {text_response.error}"; st.session_state.current_log.append(f"  Narrative: FAILED - {text_response.error}")
        else: st.session_state.current_log.append(f"  Narrative: FAILED - No response from {text_model_key}.")
    else: current_narrative = "**Narrative Error:** Text model unavailable."; st.session_state.current_log.append(f"  Narrative: FAILED - Model '{text_model_key}' unavailable.")
    
    st.session_state.latest_scene_narrative = f"## Scene Idea: {scene_prompt_text}\n\n{current_narrative}"

    # 2. Generate Image
    generated_image_pil = None
    image_gen_error = None
    selected_image_provider_actual_type = IMAGE_PROVIDERS.get(image_provider_key)
    image_content_prompt = current_narrative if current_narrative and "Error" not in current_narrative else scene_prompt_text
    quality_kw = "ultra detailed, " if image_quality == "High Detail" else ("concept sketch, " if image_quality == "Sketch Concept" else "")
    full_img_prompt = format_image_generation_prompt(quality_kw + image_content_prompt[:350], image_style_dropdown, artist_style_text)
    st.session_state.current_log.append(f"  Image: Attempting with {image_provider_key} (type '{selected_image_provider_actual_type}').")

    if selected_image_provider_actual_type and selected_image_provider_actual_type != "none":
        img_response = None
        if selected_image_provider_actual_type.startswith("dalle_") and ai_services_status["dalle_image_ready"]:
            dalle_model = "dall-e-3" if selected_image_provider_actual_type == "dalle_3" else "dall-e-2"
            img_response = generate_image_dalle(full_img_prompt, model=dalle_model, quality="hd" if image_quality=="High Detail" else "standard")
        elif selected_image_provider_actual_type.startswith("hf_") and ai_services_status["hf_image_ready"]:
            hf_model_id = "stabilityai/stable-diffusion-xl-base-1.0"; iw,ih=768,768
            if selected_image_provider_actual_type == "hf_openjourney": hf_model_id="prompthero/openjourney";iw,ih=512,512
            img_response = generate_image_hf_model(full_img_prompt, model_id=hf_model_id, negative_prompt=negative_prompt_text or COMMON_NEGATIVE_PROMPTS, width=iw, height=ih)
        
        if img_response and img_response.success: generated_image_pil = img_response.image; st.session_state.current_log.append(f"  Image: Success from {img_response.provider}.")
        elif img_response: image_gen_error = f"**Image Error:** {img_response.error}"; st.session_state.current_log.append(f"  Image: FAILED - {img_response.error}")
        else: image_gen_error = "**Image Error:** No response/unknown issue."; st.session_state.current_log.append(f"  Image: FAILED - No response object.")
    else: image_gen_error = "**Image Error:** No valid provider."; st.session_state.current_log.append(f"  Image: FAILED - No provider configured.")

    st.session_state.latest_scene_image = generated_image_pil if generated_image_pil else create_placeholder_image("Image Gen Failed", color="#401010")

    # 3. Add to Story Object
    scene_err = None
    if image_gen_error and "**Narrative Error**" in current_narrative: scene_err = f"{current_narrative}\n{image_gen_error}"
    elif "**Narrative Error**" in current_narrative: scene_err = current_narrative
    elif image_gen_error: scene_err = image_gen_error
    
    st.session_state.story_object.add_scene_from_elements(
        user_prompt=scene_prompt_text, narrative_text=current_narrative, image=generated_image_pil,
        image_style_prompt=f"{image_style_dropdown}{f', by {artist_style_text}' if artist_style_text else ''}",
        image_provider=image_provider_key, error_message=scene_err
    )
    st.session_state.current_log.append(f"  Scene {st.session_state.story_object.current_scene_number} processed.")
    st.session_state.processing_scene = False
    st.experimental_rerun() # Rerun to update the main display with new scene and re-enable button

elif submit_scene_button and not scene_prompt_text.strip(): # If form submitted but prompt is empty
    st.warning("Please enter a scene vision/prompt!")