app.py

import gradio as gr
import os
import warnings
from typing import Optional, Dict, Any

warnings.filterwarnings("ignore")

# Output directories
AUDIO_DIR = "outputs/audio"
os.makedirs(AUDIO_DIR, exist_ok=True)

# Mock implementations for story generation components
class StoryState:
    def __init__(self):
        self.scene_history = []
        self.facts = []
        self.world_meta = {}
    
    def get_context_window(self, n=3):
        return self.scene_history[-n:] if len(self.scene_history) > n else self.scene_history
    
    def update_facts(self, new_facts):
        self.facts.extend(new_facts)
    
    def append_scene(self, scene):
        self.scene_history.append(scene)
    
    def update_world_meta(self, meta):
        self.world_meta.update(meta)
    
    def to_dict(self):
        return {
            "scenes": len(self.scene_history),
            "facts": len(self.facts)
        }

class StoryGeneratorTool:
    def forward(self, context, initial_prompt=None, last_choice=None):
        if initial_prompt:
            return f"Story begins: {initial_prompt}\n\nThe adventure unfolds as mysterious events shape your path. You find yourself in an unfamiliar place, with shadows dancing around you and the sound of distant whispers filling the air. The path ahead splits into multiple directions, each promising different adventures and challenges."
        elif last_choice:
            return f"Following your choice to '{last_choice}', the story continues...\n\nYour decision leads you deeper into the mystery. New challenges emerge as the plot thickens. The environment around you shifts and changes, revealing hidden secrets and unexpected allies. What seemed like a simple choice now opens up entirely new possibilities for your adventure."
        else:
            return "The story continues with unexpected twists and turns. Ancient mysteries begin to unravel as you delve deeper into this strange world. Each step forward reveals new questions that demand answers, and you realize that your journey is far from over."

class ExtractFactsTool:
    def forward(self, scene_text):
        # Extract meaningful facts from scene text
        facts = []
        if "mysterious" in scene_text.lower():
            facts.append("There are mysterious elements in this world")
        if "path" in scene_text.lower():
            facts.append("Multiple paths and choices are available")
        if "adventure" in scene_text.lower():
            facts.append("This is an adventure-type story")
        if "challenge" in scene_text.lower():
            facts.append("Challenges and obstacles exist")
        
        return facts if facts else [f"Story element: {scene_text[:30]}..."]

def build_world(facts):
    return {
        "world_complexity": len(facts), 
        "narrative_depth": "high" if len(facts) > 5 else "medium",
        "story_themes": ["mystery", "adventure", "choice-driven"]
    }

def generate_choices(scene_text, facts):
    base_choices = [
        "Investigate the mysterious occurrence further",
        "Seek help from potential allies nearby", 
        "Explore the unknown path cautiously",
        "Take a bold and direct approach"
    ]
    
    # Customize choices based on scene content
    if "mystery" in scene_text.lower():
        base_choices[0] = "Solve the mystery using your wits"
    if "danger" in scene_text.lower():
        base_choices[3] = "Face the danger head-on"
    if "friend" in scene_text.lower() or "ally" in scene_text.lower():
        base_choices[1] = "Rally your allies for support"
    
    return base_choices

def advance_story(
    state: StoryState,
    initial_prompt: Optional[str] = None,
    last_choice: Optional[str] = None
) -> Dict[str, Any]:
    """
    Runs one step of the story pipeline focusing on NLP
    """
    # 1) Generate scene
    scene_tool = StoryGeneratorTool()
    scene_args = {
        "context": state.get_context_window(n=3),
        "initial_prompt": initial_prompt,
        "last_choice": last_choice,
    }
    scene_text = scene_tool.forward(**scene_args)
    
    # 2) Extract facts
    fact_tool = ExtractFactsTool()
    new_facts = fact_tool.forward(scene_text)
    
    # 3) Update state & build world metadata
    state.update_facts(new_facts)
    state.append_scene(scene_text)
    world_meta = build_world(state.facts)
    state.update_world_meta(world_meta)
    
    # 4) Generate next-step choices
    choices = generate_choices(scene_text, state.facts)
    
    # 5) Package and return
    return {
        "scene_text": scene_text,
        "choices": choices,
        "updated_state": state.to_dict(),
    }

# Global variables to store story state
story_state = StoryState()
current_story = ""
story_choices = []

def initialize_story(user_input):
    """Initialize the story with user input and generate first response"""
    global current_story, story_choices, story_state
    
    if not user_input.strip():
        # Hide everything if empty
        return (
            gr.update(visible=False, value=""),  # chat_output
            gr.update(visible=False),  # choice1
            gr.update(visible=False),  # choice2
            gr.update(visible=False),  # choice3
            gr.update(visible=False),  # choice4
            gr.update(visible=False),  # prev_btn
            gr.update(visible=False),  # next_btn
            gr.update(visible=True, value=""),  # keep user_input visible
            gr.update(visible=True),  # keep submit_btn visible
            gr.update(visible=False)  # choices_header
        )
    
    # Reset story state
    story_state = StoryState()
    
    # Generate story using pipeline
    result = advance_story(story_state, initial_prompt=user_input)
    
    # Update global state
    current_story = result["scene_text"]
    story_choices = result["choices"]
    
    print(f"Generated story: {current_story}")
    print(f"Generated choices: {story_choices}")
    
    # Return updates
    return (
        gr.update(visible=True, value=current_story),  # chat_output
        gr.update(visible=True, value=story_choices[0]),  # choice1
        gr.update(visible=True, value=story_choices[1]),  # choice2
        gr.update(visible=True, value=story_choices[2]),  # choice3
        gr.update(visible=True, value=story_choices[3]),  # choice4
        gr.update(visible=True),  # prev_btn
        gr.update(visible=True),  # next_btn
        gr.update(visible=False),  # hide user_input
        gr.update(visible=False),  # hide submit_btn
        gr.update(visible=True)  # show choices_header
    )

def make_choice(choice_num):
    """Handle story choice selection and generate next scene"""
    global current_story, story_choices, story_state
    
    if choice_num < len(story_choices):
        selected_choice = story_choices[choice_num]
        
        print(f"User selected choice {choice_num}: {selected_choice}")
        
        # Generate next story scene
        result = advance_story(story_state, last_choice=selected_choice)
        
        # Update global state
        current_story = result["scene_text"]
        story_choices = result["choices"]
        
        print(f"Generated new story: {current_story}")
        print(f"Generated new choices: {story_choices}")
        
        return (
            gr.update(value=current_story),           # chat_output
            gr.update(value=story_choices[0]),        # choice1
            gr.update(value=story_choices[1]),        # choice2
            gr.update(value=story_choices[2]),        # choice3
            gr.update(value=story_choices[3])         # choice4
        )    
    return tuple([gr.update()]*5)

def go_previous():
    """Navigate to previous story segment"""
    global story_state, current_story
    if len(story_state.scene_history) > 1:
        # Get previous scene
        prev_scene = story_state.scene_history[-2]
        return gr.update(value=f"📖 Previous Scene:\n\n{prev_scene}")
    return gr.update(value="📖 No previous story segment available")

def go_next():
    """Navigate back to current story segment"""
    global current_story
    return gr.update(value=current_story)

def restart_story():
    """Restart the story from beginning"""
    global story_state, current_story, story_choices
    story_state = StoryState()
    current_story = ""
    story_choices = []
    
    return (
        gr.update(visible=False, value=""),  # chat_output
        gr.update(visible=False),  # choice1
        gr.update(visible=False),  # choice2
        gr.update(visible=False),  # choice3
        gr.update(visible=False),  # choice4
        gr.update(visible=False),  # prev_btn
        gr.update(visible=False),  # next_btn
        gr.update(visible=True, value=""),  # show user_input
        gr.update(visible=True),  # show submit_btn
        gr.update(visible=False)  # hide choices_header
    )

# Create the Gradio interface
with gr.Blocks(title="AI Story Generation Platform", theme=gr.themes.Soft(), css="""
    .enter-button { margin-top: 20px !important; }
    .story-output { font-family: 'Georgia', serif !important; line-height: 1.6 !important; }
    .choice-button { margin: 5px !important; padding: 10px !important; }
    .nav-button { margin: 10px 5px !important; }
    .restart-button { margin-top: 20px !important; background: #ff6b6b !important; }
""") as demo:
    
    gr.Markdown("# 📚 AI Interactive Story Generation Platform")
    gr.Markdown("Enter your story idea to begin an interactive text-based adventure!")
    
    # Input section
    user_input = gr.Textbox(
        placeholder="Enter your story beginning (e.g., 'I wake up in a mysterious forest...')", 
        label="Story Input", 
        lines=3
    )
    submit_btn = gr.Button("🚀 Start Adventure", variant="primary", elem_classes="enter-button")

    # Story output
    chat_output = gr.Textbox(
        label="📖 Your Story", 
        lines=8, 
        interactive=False, 
        visible=False,
        elem_classes="story-output"
    )
    
    # Choices section
    choices_header = gr.Markdown("### 🎯 Choose Your Next Action", visible=False)
    
    with gr.Row():
        choice1_btn = gr.Button("Choice 1", visible=False, variant="secondary", elem_classes="choice-button")
        choice2_btn = gr.Button("Choice 2", visible=False, variant="secondary", elem_classes="choice-button")
    
    with gr.Row():
        choice3_btn = gr.Button("Choice 3", visible=False, variant="secondary", elem_classes="choice-button")
        choice4_btn = gr.Button("Choice 4", visible=False, variant="secondary", elem_classes="choice-button")
    
    # Navigation section
    with gr.Row():
        prev_btn = gr.Button("⬅️ Previous Scene", visible=False, variant="outline", elem_classes="nav-button")
        next_btn = gr.Button("Current Scene ➡️", visible=False, variant="outline", elem_classes="nav-button")
        restart_btn = gr.Button("🔄 Restart Story", visible=False, variant="stop", elem_classes="restart-button")
    
    # Story statistics
    story_stats = gr.Markdown("", visible=False)
    
    # Event handlers
    submit_btn.click(
        fn=initialize_story,
        inputs=[user_input],
        outputs=[
            chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn,
            prev_btn, next_btn, user_input, submit_btn, choices_header
        ]
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=[restart_btn]
    )
    
    # Choice event handlers
    choice1_btn.click(
        fn=lambda: make_choice(0), 
        outputs=[chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn]
    )
    choice2_btn.click(
        fn=lambda: make_choice(1), 
        outputs=[chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn]
    )
    choice3_btn.click(
        fn=lambda: make_choice(2), 
        outputs=[chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn]
    )
    choice4_btn.click(
        fn=lambda: make_choice(3), 
        outputs=[chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn]
    )
    
    # Navigation event handlers
    prev_btn.click(fn=go_previous, outputs=[chat_output])
    next_btn.click(fn=go_next, outputs=[chat_output])
    
    # Restart handler
    restart_btn.click(
        fn=restart_story,
        outputs=[
            chat_output, choice1_btn, choice2_btn, choice3_btn, choice4_btn,
            prev_btn, next_btn, user_input, submit_btn, choices_header
        ]
    ).then(
        fn=lambda: gr.update(visible=False),
        outputs=[restart_btn]
    )

if __name__ == "__main__":
    demo.launch()