# app.py (Streamlit version for StoryVerse Weaver - Robust Error Handling)
import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import os
import time
import random
import traceback

# --- Page Configuration (MUST BE THE VERY FIRST STREAMLIT COMMAND) ---
st.set_page_config(
    page_title="✨ StoryVerse Weaver ✨",
    page_icon="😻",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Core Logic Imports ---
from core.llm_services import initialize_text_llms, is_gemini_text_ready, is_hf_text_ready, generate_text_gemini, generate_text_hf, LLMTextResponse
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
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 AI Services (Cached) ---
@st.cache_resource
def load_ai_services_config():
    print("--- Initializing AI Services (Streamlit Cache Resource) ---")
    initialize_text_llms()
    initialize_image_llms()
    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_config()

# --- Application Configuration (Models, Defaults) ---
TEXT_MODELS = {}
UI_DEFAULT_TEXT_MODEL_KEY = None
if AI_SERVICES_STATUS["gemini_text_ready"]:
    TEXT_MODELS["✨ Gemini 1.5 Flash (Narrate)"] = {"id": "gemini-1.5-flash-latest", "type": "gemini"}
if AI_SERVICES_STATUS["hf_text_ready"]: # Used if Gemini not ready or as an option
    TEXT_MODELS["Mistral 7B (Narrate via HF)"] = {"id": "mistralai/Mistral-7B-Instruct-v0.2", "type": "hf_text"}

if TEXT_MODELS:
    if AI_SERVICES_STATUS["gemini_text_ready"] and "✨ Gemini 1.5 Flash (Narrate)" in TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = "✨ Gemini 1.5 Flash (Narrate)"
    elif AI_SERVICES_STATUS["hf_text_ready"] and "Mistral 7B (Narrate via HF)" in TEXT_MODELS: UI_DEFAULT_TEXT_MODEL_KEY = "Mistral 7B (Narrate via HF)"
    else: UI_DEFAULT_TEXT_MODEL_KEY = list(TEXT_MODELS.keys())[0]
else:
    TEXT_MODELS["No Text Models Configured"] = {"id": "dummy_text_error", "type": "none"}
    UI_DEFAULT_TEXT_MODEL_KEY = "No Text Models Configured"

IMAGE_PROVIDERS = {}
UI_DEFAULT_IMAGE_PROVIDER_KEY = None
if AI_SERVICES_STATUS["dalle_image_ready"]: 
    IMAGE_PROVIDERS["🖼️ OpenAI DALL-E 3"] = "dalle_3"
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🖼️ OpenAI DALL-E 3"
elif AI_SERVICES_STATUS["hf_image_ready"]: 
    IMAGE_PROVIDERS["🎡 HF - SDXL Base"] = "hf_sdxl_base" 
    IMAGE_PROVIDERS["🎠 HF - OpenJourney"] = "hf_openjourney"
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "🎡 HF - SDXL Base"
if not IMAGE_PROVIDERS: 
    IMAGE_PROVIDERS["No Image Providers Configured"] = "none"
    UI_DEFAULT_IMAGE_PROVIDER_KEY = "No Image Providers Configured"

# --- Custom CSS ---
streamlit_omega_css = """ /* ... Your full omega_css string here ... */ """
st.markdown(streamlit_omega_css, unsafe_allow_html=True)

# --- Helper: Placeholder Image Creation ---
@st.cache_data 
def create_placeholder_image_st(text="Processing...", size=(512, 512), color="#23233A", text_color="#E0E0FF"):
    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 ---
ss = st.session_state # Shortcut
if 'story_object' not in ss: ss.story_object = Story()
if 'current_log' not in ss: ss.current_log = ["Welcome to StoryVerse Weaver!"]
if 'latest_scene_image_pil' not in ss: ss.latest_scene_image_pil = None
if 'latest_scene_narrative' not in ss: ss.latest_scene_narrative = "## ✨ A New Story Begins ✨\nDescribe your first scene idea in the panel on the left!"
if 'status_message' not in ss: ss.status_message = {"text": "Ready to weave your first masterpiece!", "type": "processing"}
if 'processing_scene' not in ss: ss.processing_scene = False
if 'form_scene_prompt' not in ss: ss.form_scene_prompt = ""
if 'form_image_style' not in ss: ss.form_image_style = "Default (Cinematic Realism)"
if 'form_artist_style' not in ss: ss.form_artist_style = ""

# --- UI Definition ---
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 Weaver help you craft captivating scenes. Ensure API keys (<code>STORYVERSE_...</code>) are set in your environment/secrets!</div>", unsafe_allow_html=True)

# --- Sidebar for Inputs & Configuration ---
with st.sidebar:
    st.markdown("### 💡 **Craft Your Scene**")
    with st.form("scene_input_form_key", clear_on_submit=False):
        scene_prompt_text_val = st.text_area("Scene Vision:", value=ss.form_scene_prompt, height=150)
        st.markdown("#### 🎨 Visual Style")
        col_style1, col_style2 = st.columns(2)
        with col_style1: image_style_val = st.selectbox("Style Preset:", options=["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys())), index=(["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))).index(ss.form_image_style) if ss.form_image_style in (["Default (Cinematic Realism)"] + sorted(list(STYLE_PRESETS.keys()))) else 0)
        with col_style2: artist_style_val = st.text_input("Artistic Inspiration:", value=ss.form_artist_style)
        negative_prompt_val = st.text_area("Exclude from Image:", value=COMMON_NEGATIVE_PROMPTS, height=80)
        with st.expander("⚙️ Advanced AI Configuration", expanded=False):
            text_model_val = st.selectbox("Narrative Engine:", options=list(TEXT_MODELS.keys()), index=list(TEXT_MODELS.keys()).index(UI_DEFAULT_TEXT_MODEL_KEY) if UI_DEFAULT_TEXT_MODEL_KEY in TEXT_MODELS else 0)
            image_provider_val = st.selectbox("Visual Engine:", options=list(IMAGE_PROVIDERS.keys()), index=list(IMAGE_PROVIDERS.keys()).index(UI_DEFAULT_IMAGE_PROVIDER_KEY) if UI_DEFAULT_IMAGE_PROVIDER_KEY in IMAGE_PROVIDERS else 0)
            narrative_length_val = st.selectbox("Narrative Detail:", ["Short", "Medium", "Detailed"], index=1)
            image_quality_val = st.selectbox("Image Detail:", ["Standard", "High Detail", "Sketch"], index=0)
        submit_button_val = st.form_submit_button("🌌 Weave Scene!", use_container_width=True, type="primary", disabled=ss.processing_scene)

    col_btn_s, col_btn_c = st.columns(2)
    with col_btn_s:
        if st.button("🎲 Surprise Me!", use_container_width=True, disabled=ss.processing_scene, key="sidebar_surprise_btn"):
            sur_prompt, sur_style, sur_artist = surprise_me_func(); ss.form_scene_prompt = sur_prompt; ss.form_image_style = sur_style; ss.form_artist_style = sur_artist; st.rerun()
    with col_btn_c:
        if st.button("🗑️ New Story", use_container_width=True, disabled=ss.processing_scene, key="sidebar_clear_btn"):
            ss.story_object = Story(); ss.current_log = ["Cleared."]; ss.latest_scene_image_pil = None; ss.latest_scene_narrative = "## New Story"; ss.status_message = {"text": "Cleared", "type": "processing"}; ss.form_scene_prompt = ""; st.rerun()

    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 Ready.")
            else: st.warning("Text Generation NOT Ready.")
            if image_gen_ok: st.success("Image Generation Ready.")
            else: st.warning("Image Generation NOT Ready.")

# --- Main Display Area ---
st.markdown("---")
st.markdown("### 🖼️ **Your Evolving StoryVerse**", unsafe_allow_html=True) 

status_type = ss.status_message.get("type", "processing")
st.markdown(f"<p class='status-message status-{status_type}'>{ss.status_message['text']}</p>", unsafe_allow_html=True)

tab_latest, tab_scroll, tab_log = st.tabs(["🌠 Latest Scene", "📚 Story Scroll", "⚙️ Interaction Log"])

with tab_latest:
    if ss.processing_scene and ss.latest_scene_image_pil is None :
        st.image(create_placeholder_image_st("🎨 Conjuring visuals..."), use_container_width=True)
    elif ss.latest_scene_image_pil:
        st.image(ss.latest_scene_image_pil, use_container_width=True, caption="Latest Generated Image")
    else:
        st.image(create_placeholder_image_st("Describe a scene to begin!", size=(512,300), color="#1A1A2E"), use_container_width=True)
    st.markdown(ss.latest_scene_narrative, unsafe_allow_html=True)

with tab_scroll:
    if ss.story_object and ss.story_object.scenes:
        st.subheader("Story Scroll")
        num_cols_gallery = st.slider("Gallery Columns:", 1, 5, 3, key="gallery_cols_slider")
        gallery_cols_list = st.columns(num_cols_gallery)
        scenes_for_gallery_data = ss.story_object.get_all_scenes_for_gallery_display()
        for i, (img, caption) in enumerate(scenes_for_gallery_data):
            with gallery_cols_list[i % num_cols_gallery]:
                display_img_gallery = img
                if img is None: display_img_gallery = create_placeholder_image_st(f"S{i+1}\nNo Image", size=(180,180), color="#2A2A4A")
                st.image(display_img_gallery, caption=caption if caption else f"Scene {i+1}", use_container_width=True, output_format="PNG")
    else:
        st.caption("Your story scroll is empty. Weave your first scene!")

with tab_log:
    log_display_text = "\n\n---\n\n".join(ss.current_log[::-1][:50]) 
    st.markdown(log_display_text, unsafe_allow_html=True)

# --- Logic for Form Submission ---
if submit_button_val: 
    if not scene_prompt_text_val.strip():
        ss.status_message = {"text": "Scene prompt cannot be empty!", "type": "error"}
        st.rerun() 
    else:
        ss.processing_scene = True
        ss.status_message = {"text": f"🌌 Weaving Scene {ss.story_object.current_scene_number + 1}...", "type": "processing"}
        ss.current_log.append(f"**🚀 Scene {ss.story_object.current_scene_number + 1} - {time.strftime('%H:%M:%S')}**")
        st.rerun() # Rerun to show "processing" and disable button

        # ---- Main Generation Logic ----
        _scene_prompt = scene_prompt_text_val
        _image_style = image_style_val
        _artist_style = artist_style_val
        _negative_prompt = negative_prompt_val
        _text_model = text_model_val
        _image_provider = image_provider_val
        _narr_length = narrative_length_val
        _img_quality = image_quality_val
        
        current_narrative_text = f"Narrative Error: Init failed."
        generated_image_pil = None
        image_gen_error_msg = None # Specific to image generation part
        final_scene_error_overall = None # Overall error for the scene

        # 1. Generate Narrative
        try:
            with st.spinner("✍️ Crafting narrative... (This may take a moment)"):
                text_model_info = TEXT_MODELS.get(_text_model)
                if text_model_info and text_model_info["type"] != "none":
                    system_p = get_narrative_system_prompt("default"); prev_narr = ss.story_object.get_last_scene_narrative(); user_p = format_narrative_user_prompt(_scene_prompt, prev_narr)
                    ss.current_log.append(f"  Narrative: Using {_text_model} ({text_model_info['id']}).")
                    text_resp = None
                    if text_model_info["type"] == "gemini" and AI_SERVICES_STATUS["gemini_text_ready"]: text_resp = generate_text_gemini(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
                    elif text_model_info["type"] == "hf_text" and AI_SERVICES_STATUS["hf_text_ready"]: text_resp = generate_text_hf(user_p, model_id=text_model_info["id"], system_prompt=system_p, max_tokens=768 if _narr_length.startswith("Detailed") else 400)
                    if text_resp and text_resp.success: current_narrative_text = basic_text_cleanup(text_resp.text); ss.current_log.append("  Narrative: Success.")
                    elif text_resp: current_narrative_text = f"**Narrative Error ({_text_model}):** {text_resp.error}"; ss.current_log.append(f"  Narrative: FAILED - {text_resp.error}")
                    else: ss.current_log.append(f"  Narrative: FAILED - No response.")
                else: current_narrative_text = "**Narrative Error:** Model unavailable."; ss.current_log.append(f"  Narrative: FAILED - Model '{_text_model}' unavailable.")
            ss.latest_scene_narrative = f"## Scene Idea: {_scene_prompt}\n\n{current_narrative_text}"
        except Exception as e_narr:
            current_narrative_text = f"**Narrative Generation Exception:** {type(e_narr).__name__} - {str(e_narr)}"
            ss.current_log.append(f"  CRITICAL NARRATIVE ERROR: {traceback.format_exc()}")
            ss.latest_scene_narrative = f"## Scene Idea: {_scene_prompt}\n\n{current_narrative_text}"
            final_scene_error_overall = current_narrative_text

        # 2. Generate Image
        try:
            with st.spinner("🎨 Conjuring visuals... (This may take a moment)"):
                selected_img_prov_type = IMAGE_PROVIDERS.get(_image_provider)
                img_content_prompt = current_narrative_text if current_narrative_text and "Error" not in current_narrative_text else _scene_prompt
                quality_kw = "ultra detailed, " if _img_quality == "High Detail" else ("concept sketch, " if _img_quality == "Sketch Concept" else "")
                full_img_prompt_for_gen = format_image_generation_prompt(quality_kw + img_content_prompt[:350], _image_style, _artist_style)
                ss.current_log.append(f"  Image: Using {_image_provider} (type '{selected_img_prov_type}').")
                if selected_img_prov_type and selected_img_prov_type != "none":
                    img_resp = None
                    if selected_img_prov_type.startswith("dalle_") and AI_SERVICES_STATUS["dalle_image_ready"]:
                        dalle_model = "dall-e-3" if selected_img_prov_type == "dalle_3" else "dall-e-2"
                        img_resp = generate_image_dalle(full_img_prompt_for_gen, model=dalle_model)
                    elif selected_img_prov_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_img_prov_type == "hf_openjourney": hf_model_id="prompthero/openjourney";iw,ih=512,512
                        img_resp = generate_image_hf_model(full_img_prompt_for_gen, model_id=hf_model_id, negative_prompt=_negative_prompt, width=iw, height=ih)
                    if img_resp and img_resp.success: generated_image_pil = img_resp.image; ss.current_log.append("  Image: Success.")
                    elif img_resp: image_gen_error_msg = f"**Image Error:** {img_resp.error}"; ss.current_log.append(f"  Image: FAILED - {img_resp.error}")
                    else: image_gen_error_msg = "**Image Error:** No response."; ss.current_log.append("  Image: FAILED - No response.")
                else: image_gen_error_msg = "**Image Error:** No provider configured."; ss.current_log.append(f"  Image: FAILED - No provider.")
            ss.latest_scene_image_pil = generated_image_pil if generated_image_pil else create_placeholder_image_st("Image Gen Failed", color="#401010")
        except Exception as e_img:
            image_gen_error_msg = f"**Image Generation Exception:** {type(e_img).__name__} - {str(e_img)}"
            ss.current_log.append(f"  CRITICAL IMAGE ERROR: {traceback.format_exc()}")
            ss.latest_scene_image_pil = create_placeholder_image_st("Image Gen Exception", color="#401010")
            if not final_scene_error_overall: final_scene_error_overall = image_gen_error_msg
            else: final_scene_error_overall += f"\n{image_gen_error_msg}"

        # 3. Add Scene
        if not final_scene_error_overall: # If no major error from narrative already
            if image_gen_error_msg and "**Narrative Error**" in current_narrative_text: final_scene_error_overall = f"{current_narrative_text}\n{image_gen_error_msg}"
            elif "**Narrative Error**" in current_narrative_text: final_scene_error_overall = current_narrative_text
            elif image_gen_error_msg: final_scene_error_overall = image_gen_error_msg
        
        ss.story_object.add_scene_from_elements(
            user_prompt=_scene_prompt, narrative_text=current_narrative_text, image=generated_image_pil,
            image_style_prompt=f"{_image_style}{f', by {_artist_style}' if _artist_style else ''}",
            image_provider=_image_provider, error_message=final_scene_error_overall
        )
        ss.current_log.append(f"  Scene {ss.story_object.current_scene_number} processed.")
        
        # 4. Set final status message
        if final_scene_error_overall: 
            ss.status_message = {"text": f"Scene {ss.story_object.current_scene_number} generated with errors. Check log.", "type": "error"}
        else: 
            ss.status_message = {"text": f"🌌 Scene {ss.story_object.current_scene_number} woven successfully!", "type": "success"}
        
        ss.processing_scene = False
        st.rerun()