app.py

import streamlit as st
import json
from dataclasses import dataclass
from typing import List, Dict, Optional
from pathlib import Path
import requests
from datetime import datetime
import re
from huggingface_hub import InferenceClient

# Creating an OpenAI call
from openai import OpenAI

@dataclass
class Character:
    """Represents a character in the visual novel, storing their attributes and state"""
    name: str
    description: str
    personality: str
    relationships: Dict[str, int]  # Character name -> relationship value
    current_mood: str = "neutral"
    background_story: str = ""

    def __init__(self, name: str, description: str, personality: str, relationships: Dict[str, int] = None,
                 current_mood: str = "neutral", background_story: str = ""):
        self.name = name
        self.description = description
        self.personality = personality
        self.relationships = relationships if relationships is not None else {}
        self.current_mood = current_mood
        self.background_story = background_story

@dataclass
class StoryState:
    """Maintains the current state of the visual novel, including scene and character information"""
    current_scene: str
    characters: Dict[str, Character]
    player_choices: List[str]
    story_branch: str
    mood: str
    setting: str
    genre: str
    chapter: int = 1
    scene_number: int = 1

class VisualNovelGenerator:
    def __init__(self):
        """Initialize the visual novel generator with API configuration"""
        self.setup_streamlit_config()
        self.initialize_session_state()
        self.setup_api()

    def setup_api(self):
        """Configure the OpenAI API client"""
        # Load API key from Streamlit secrets or environment variables
        self.api_key = st.secrets.get("API_KEY") or st.session_state.get("api_key")
        
        # Initialize OpenAI client
        self.client = OpenAI(
        base_url = "https://integrate.api.nvidia.com/v1",
        api_key = self.api_key
        )

    def api_request(self, messages: List[Dict[str, str]], max_tokens: int = 1000, temperature=0.7) -> str:
        """Send a request to the API and handle the response"""
        try:
            # temperature = st.slider("Creativity: ", 0.0, 1.0, 0.7, key='_tempslider')
            response = self.client.chat.completions.create(
                model="meta/llama-3.1-405b-instruct",
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature
            )
            return response.choices[0].message.content
        except Exception as e:
            st.error(f"API Error: {str(e)}")
            return None

    def setup_streamlit_config(self):
        """Configure the Streamlit page settings for optimal display"""
        st.set_page_config(
            page_title="Visual Novel Generator",
            page_icon="📚",
            layout="wide",
            initial_sidebar_state="expanded"
        )
    # def image_api_request(self, prompt: str) -> str:
    #     """Send a request to generate an image based on the prompt"""
    #     image_api_key = st.secrets.get("IMG_API_KEY") or st.session_state.get("IMG_API_KEY")
    #     client = InferenceClient("black-forest-labs/FLUX.1-dev", token=image_api_key)
    #     try:
    #         image = client.text_to_image(prompt, target_size=(512, 512))    
    #         return image
    #     except Exception as e:
    #         st.error(f"Image Generation Error: {str(e)}")
    #         return None

    def initialize_session_state(self):
        """Set up the initial session state variables"""
        if 'api_key' not in st.session_state:
            st.session_state.api_key = None
        if 'story_state' not in st.session_state:
            st.session_state.story_state = None
        if 'story_history' not in st.session_state:
            st.session_state.story_history = []
        if 'save_states' not in st.session_state:
            st.session_state.save_states = {}

    def create_api_key_input(self):
        """Create an input field for the API key if not already set"""
        if not self.api_key:
            st.sidebar.markdown("### API Configuration")
            api_key = st.sidebar.text_input("Enter your API key:", type="password")
            if api_key:
                st.session_state.api_key = api_key
                self.api_key = api_key
                self.setup_api()
                st.sidebar.success("API key set successfully!")
                st.experimental_experimental_rerun()
            return False
        return True

    def create_story_setup_interface(self):
        """Create the initial story setup interface with enhanced options"""
        if not self.create_api_key_input():
            return

        st.title("📚 Interactive Visual Novel Generator")
        st.markdown("### Create Your Story")
        
        with st.form("story_setup"):
            # Advanced story settings
            col1, col2, col3 = st.columns(3)

            with col1:
                genre = st.selectbox(
                    "Genre",
                    ["Romance", "Mystery", "Fantasy", "Science Fiction", "Horror", "Adventure"]
                )
                story_theme = st.text_area(
                    "Theme/Topic",
                    help="Main theme or concept of your story"
                )
                story_theme = st.selectbox(
                    "Theme/Topic",
                    ["Coming of Age", "Good vs Evil", "Power of Friendship", "Forbidden Love", 
                     "Quest for Identity", "Revenge", "Redemption", "Personal Growth", 
                     "Family Bonds", "Justice vs Mercy"],
                    help="Main theme or concept of your story"
                )
                
            with col2:
                setting = st.select_slider(
                    "Time Period",
                    options=["Prehistoric", "Ancient", "Medieval", "Renaissance", "Modern", "Future"],
                    value="Modern",
                    help="Choose the time period for your story"
                )
                tone = st.selectbox(
                    "Tone",
                    ["Light and Humorous", "Dark and Serious", "Dramatic", "Mysterious", "Romantic"]
                )
                
            with col3:
                pacing = st.selectbox(
                    "Story Pacing",
                    ["Fast-paced", "Moderate", "Slow-burn"]
                )
                narrative_style = st.selectbox(
                    "Narrative Style",
                    ["First Person", "Third Person Limited", "Third Person Omniscient"]
                )
            
            # Enhanced character creation
            st.markdown("### Create Your Characters")
            num_characters = st.number_input("Number of Characters", min_value=1, max_value=5, value=2)
            
            characters = {}
            for i in range(num_characters):
                with st.expander(f"Character {i+1}"):
                    char_name = st.text_input("Name", key=f"char_name_{i}")
                    # Let's generate character details using AI for each character
                    if char_name:
                        prompt = f"""Return a JSON object for a character named {char_name} in this format exactly:
                            {{
                                "description": "detailed physical appearance",
                                "personality": "personality traits",
                                "background": "background story",
                                "goals": "character goals"
                            }}"""
                            
                        response = self.api_request([
                                {"role": "system", "content": "You are a character designer that only returns valid JSON."},
                                {"role": "user", "content": prompt}
                            ], temperature=0.5)

                        if response:
                            try:
                                # Remove any non-JSON text before and after the JSON object
                                json_start = response.find("{")
                                json_end = response.rfind("}") + 1
                                if json_start >= 0 and json_end > json_start:
                                    json_str = response[json_start:json_end]
                                    char_info = json.loads(json_str)
                                    characters[char_name] = char_info
                                    st.success(f"Character {char_name} details generated!")

                                # Remove any non-JSON text before and after the JSON object
                                json_start = response.find("{")
                                json_end = response.rfind("}") + 1
                                if json_start >= 0 and json_end > json_start:
                                    json_str = response[json_start:json_end]
                                    char_info = json.loads(json_str)
                                    characters[char_name] = char_info
                                    st.success(f"Character {char_name} details generated!")

                            #     # Generate character portrait
                            #     portrait_prompt = f"Visual novel style portrait of {char_name}: full face, shoulders up, anime style, high quality render"
                            #     # Generate and save character portrait
                            #     portrait_filename = f"portrait_{char_name.lower().replace(' ', '_')}.png"
                            #     if not Path(portrait_filename).exists():
                            #         portrait = self.image_api_request(portrait_prompt)
                            #         if portrait:
                            #             portrait.save(portrait_filename)
                            #             st.image(portrait_filename, caption=f"Portrait of {char_name}")
                            #     else:
                            #         st.image(portrait_filename, caption=f"Portrait of {char_name}")

                            except:
                                st.error(f"Error generating details for {char_name}")

            # Advanced story options
            with st.expander("Advanced Story Options"):
                story_length = st.selectbox(
                    "Story Length",
                    ["Short (3-5 scenes)", "Medium (6-10 scenes)", "Long (11-15 scenes)"]
                )
                complexity = st.selectbox(
                    "Story Complexity",
                    ["Simple", "Moderate", "Complex"]
                )
                themes = st.multiselect(
                    "Additional Themes",
                    ["Friendship", "Betrayal", "Love", "Loss", "Redemption", "Growth"]
                )

            submit_button = st.form_submit_button("Generate Story")
            
            if submit_button and characters:
                with st.spinner("Crafting your story..."):
                    self.generate_story(
                        genre, story_theme, setting, tone, characters,
                        pacing, narrative_style, story_length, complexity, themes
                    )

    def generate_story(self, genre, theme, setting, tone, characters, pacing, 
                      narrative_style, story_length, complexity, themes):
        """Generate the initial story state and content using the API"""
        # Create a detailed story prompt
        story_prompt = f"""
        Create an engaging visual novel opening with these parameters:
        Genre: {genre}
        Theme: {theme}
        Setting: {setting}
        Tone: {tone}
        Pacing: {pacing}
        Narrative Style: {narrative_style}
        Story Length: {story_length}
        Complexity: {complexity}
        Additional Themes: {', '.join(themes)}
        
        Characters:
        {json.dumps(characters, indent=2)}
        
        Create:
        1. A vivid opening scene (2-3 paragraphs)
        2. Initial emotional states for each character
        3. Three meaningful dialogue options that will impact the story
        4. Brief background context for the scene
        
        Format as JSON:
        {{
            "opening_scene": "scene description",
            "background_context": "context",
            "character_emotions": {{"character_name": "emotion"}},
            "dialogue_options": ["option1", "option2", "option3"]
        }}
        """
        
        try:
            # Generate story content through API
            response = self.api_request([
                {"role": "system", "content": "You are a professional visual novel writer specializing in creating engaging, character-driven narratives."},
                {"role": "user", "content": story_prompt}
            ])
            
            if not response:
                return

            try:
                # Clean the response string to ensure valid JSON
                json_str = response.strip()
                # Find the first '{' and last '}' to extract just the JSON object
                start = json_str.find('{')
                end = json_str.rfind('}') + 1
                if start == -1 or end == 0:
                    raise ValueError("No valid JSON object found in response")
                clean_json = json_str[start:end]
                
                story_data = json.loads(clean_json)
            except json.JSONDecodeError as e:
                st.error(f"Invalid JSON response from API: {str(e)}")
                return
            except ValueError as e:
                st.error(f"Could not find valid JSON in response: {str(e)}")
                return
            
            # Initialize enhanced character objects
            char_objects = {}
            for char_name, char_info in characters.items():
                char_objects[char_name] = Character(
                    name=char_name,
                    description=char_info["description"],
                    personality=char_info["personality"],
                    relationships={},
                    current_mood=story_data["character_emotions"].get(char_name, "neutral"),
                    background_story=char_info["background"]
                )
            
            # Create detailed story state
            st.session_state.story_state = StoryState(
                current_scene=f"{story_data['background_context']}\n\n{story_data['opening_scene']}",
                characters=char_objects,
                player_choices=story_data["dialogue_options"],
                story_branch="main",
                mood=tone,
                setting=setting,
                genre=genre
            )
            
            # Initialize story history
            st.session_state.story_history.append({
                "scene": st.session_state.story_state.current_scene,
                "choices": story_data["dialogue_options"],
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })

            # # Generate images for the opening scene and characters
            # try:
            #     # Generate scene image
            #     scene_prompt = f"Visual novel style, {setting} setting: {story_data['opening_scene']}"
            #     scene_image = self.image_api_request(scene_prompt)
            #     if scene_image:
            #         scene_image.save("temp_scene.png")
            #         st.image("temp_scene.png", caption="Opening Scene")

            # except Exception as e:
            #     st.warning(f"Image generation failed: {str(e)}")
            
        except Exception as e:
            st.error(f"Error generating story: {str(e)}")
            st.error("Please try again with different inputs.")

    def display_story_interface(self):
        """Display the main story interface with enhanced features"""
        if not st.session_state.story_state:
            return
        
        st.title("📖 Your Visual Novel")
        
        # Main story display
        col1, col2 = st.columns([2, 1])
        
        with col1:
            # Scene information
            st.markdown(f"### Chapter {st.session_state.story_state.chapter}, Scene {st.session_state.story_state.scene_number}")
            st.write(st.session_state.story_state.current_scene)
            
            # Enhanced choice system
            st.markdown("### What will you do?")
            for option in st.session_state.story_state.player_choices:
                if st.button(option, key=f"choice_{option}"):
                    self.process_player_choice(option)
        
        with col2:
            # Character status and relationships
            st.markdown("### Characters")
            for char_name, char in st.session_state.story_state.characters.items():
                with st.expander(char_name):
                    st.write(f"Current Mood: {char.current_mood}")
                    st.markdown("#### Personality")
                    st.write(char.personality)
                    st.markdown("#### Background")
                    st.write(char.background_story)
                    if char.relationships:
                        st.markdown("#### Relationships")
                        for other_char, value in char.relationships.items():
                            st.write(f"- {other_char}: {value}")

        # Story management features
        col3, col4 = st.columns([1, 1])
        with col3:
            if st.button("Save Game"):
                save_name = f"Save_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                st.session_state.save_states[save_name] = {
                    "story_state": st.session_state.story_state,
                    "history": st.session_state.story_history
                }
                st.success(f"Game saved as: {save_name}")

        with col4:
            if st.session_state.save_states:
                save_files = list(st.session_state.save_states.keys())
                selected_save = st.selectbox("Load Save File:", save_files)
                if st.button("Load Game"):
                    save_data = st.session_state.save_states[selected_save]
                    st.session_state.story_state = save_data["story_state"]
                    st.session_state.story_history = save_data["history"]
                    st.success("Game loaded successfully!")
                    st.experimental_rerun()

        # Story history with timestamps
        with st.expander("Story History"):
            for i, history_item in enumerate(st.session_state.story_history):
                st.markdown(f"### Scene {i+1} ({history_item['timestamp']})")
                st.write(history_item["scene"])
                st.markdown("Choices:")
                for choice in history_item["choices"]:
                    st.write(f"- {choice}")

    def process_player_choice(self, choice):
        """Process the player's choice and generate the next scene using the API"""
        try:
            # Create a detailed context for the next scene
            context = {
                "current_scene": st.session_state.story_state.current_scene,
                "choice": choice,
                "genre": st.session_state.story_state.genre,
                "mood": st.session_state.story_state.mood,
                "characters": {name: {"mood": char.current_mood, "personality": char.personality}
                             for name, char in st.session_state.story_state.characters.items()},
                "chapter": st.session_state.story_state.chapter,
                "scene_number": st.session_state.story_state.scene_number
            }
            
            prompt = f"""
            Story Context: {json.dumps(context, indent=2)}
            
            Generate the next scene based on the player's choice. Include:
            1. Scene description
            2. Updated character emotions
            3. Three new meaningful choices
            4. Any relationship changes between characters.
            5. You should not add any additonal comments on the json file.
            
            Format as JSON:
            {{
                "scene": "scene description",
                "character_emotions": {{"character_name": "emotion"}},
                "dialogue_options": ["option1", "option2", "option3"],
                "relationship_changes": {{"character1": {{"character2": "change_value"}}}}
            }}
            """
            
            # Generate next scene through API
            response = self.api_request([
                {"role": "system", "content": "You are a professional visual novel writer continuing an ongoing story."},
                {"role": "user", "content": prompt}
            ])
            
            if not response:
                return

            try:
                # Clean the response string to ensure valid JSON
                json_str = response.strip()
                # Find the first '{' and last '}' to extract just the JSON object
                start = json_str.find('{')
                end = json_str.rfind('}') + 1
                if start == -1 or end == 0:
                    raise ValueError("No valid JSON object found in response")
                clean_json = json_str[start:end]
                scene_data = json.loads(clean_json)
            except json.JSONDecodeError as e:
                st.error(f"Invalid JSON response from API: {str(e)}")
                st.write(response)
                return
            except ValueError as e:
                st.error(f"Could not find valid JSON in response: {str(e)}")
                return
            
            # Update character states
            for char_name, emotion in scene_data["character_emotions"].items():
                if char_name in st.session_state.story_state.characters:
                    st.session_state.story_state.characters[char_name].current_mood = emotion
            
            # Update relationships
            for char1, relations in scene_data.get("relationship_changes", {}).items():
                for char2, change in relations.items():
                    if char1 in st.session_state.story_state.characters and char2 in st.session_state.story_state.characters:
                        current_value = st.session_state.story_state.characters[char1].relationships.get(char2, 0)
                        # Extract numeric value from the change string using regex
                        numeric_change = int(''.join(filter(str.isdigit, str(change))))
                        # Apply negative if "decrease" or "-" is in the change string
                        if any(term in str(change).lower() for term in ['decrease', '-']):
                            numeric_change = -numeric_change
                        st.session_state.story_state.characters[char1].relationships[char2] = current_value + numeric_change
            
            # Update story state
            st.session_state.story_state.current_scene = scene_data["scene"]
            st.session_state.story_state.player_choices = scene_data["dialogue_options"]
            st.session_state.story_state.scene_number += 1
            
            # Check if we should advance to next chapter
            if st.session_state.story_state.scene_number > 3:  # Advance chapter every 3 scenes
                st.session_state.story_state.chapter += 1
                st.session_state.story_state.scene_number = 1
            
            # Add to history
            st.session_state.story_history.append({
                "scene": scene_data["scene"],
                "choices": scene_data["dialogue_options"],
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            })
            
            # Force a experimental_rerun to update the interface
            st.experimental_rerun()
            
        except Exception as e:
            st.error(f"Error processing choice: {str(e)}")
            return

    def save_story_to_file(self):
        """Save the current story state to a JSON file"""
        if not st.session_state.story_state:
            return False
            
        try:
            save_data = {
                "story_state": {
                    "current_scene": st.session_state.story_state.current_scene,
                    "chapter": st.session_state.story_state.chapter,
                    "scene_number": st.session_state.story_state.scene_number,
                    "genre": st.session_state.story_state.genre,
                    "mood": st.session_state.story_state.mood,
                    "setting": st.session_state.story_state.setting
                },
                "characters": {
                    name: {
                        "name": char.name,
                        "description": char.description,
                        "personality": char.personality,
                        "current_mood": char.current_mood,
                        "background_story": char.background_story,
                        "relationships": char.relationships
                    }
                    for name, char in st.session_state.story_state.characters.items()
                },
                "history": st.session_state.story_history
            }
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"visual_novel_save_{timestamp}.json"
            
            with open(filename, 'w') as f:
                json.dump(save_data, f, indent=2)
                
            return filename
            
        except Exception as e:
            st.error(f"Error saving story: {str(e)}")
            return False

    def load_story_from_file(self, filename):
        """Load a story state from a JSON file"""
        try:
            with open(filename, 'r') as f:
                save_data = json.load(f)
                
            # Reconstruct characters
            characters = {}
            for name, char_data in save_data["characters"].items():
                characters[name] = Character(
                    name=char_data["name"],
                    description=char_data["description"],
                    personality=char_data["personality"],
                    current_mood=char_data["current_mood"],
                    background_story=char_data["background_story"],
                    relationships=char_data["relationships"]
                )
                
            # Reconstruct story state
            st.session_state.story_state = StoryState(
                current_scene=save_data["story_state"]["current_scene"],
                characters=characters,
                player_choices=save_data["story_state"].get("player_choices", []),
                story_branch=save_data["story_state"].get("story_branch", "main"),
                mood=save_data["story_state"]["mood"],
                setting=save_data["story_state"]["setting"],
                genre=save_data["story_state"]["genre"],
                chapter=save_data["story_state"]["chapter"],
                scene_number=save_data["story_state"]["scene_number"]
            )
            
            # Load history
            st.session_state.story_history = save_data["history"]
            
            return True
            
        except Exception as e:
            st.error(f"Error loading story: {str(e)}")
            return False

    def run(self):
        """Main application loop with enhanced save/load functionality"""
        # Sidebar options
        with st.sidebar:
            st.markdown("### Story Management")
            if st.button("Start New Story"):
                st.session_state.story_state = None
                st.session_state.story_history = []
                st.experimental_rerun()
                
            if st.session_state.story_state:
                if st.button("Save Story to File"):
                    filename = self.save_story_to_file()
                    if filename:
                        st.success(f"Story saved to {filename}")
                
                uploaded_file = st.file_uploader("Load Story from File", type="json")
                if uploaded_file:
                    if self.load_story_from_file(uploaded_file):
                        st.success("Story loaded successfully!")
                        st.experimental_rerun()

        # Main interface
        if not st.session_state.story_state:
            self.create_story_setup_interface()
        else:
            self.display_story_interface()

# Run the application
if __name__ == "__main__":
    app = VisualNovelGenerator()
    app.run()