# app.py - Final Version with AI-Powered Animations
# Fix for huggingface_hub compatibility
import huggingface_hub
if not hasattr(huggingface_hub, 'cached_download'):
    from huggingface_hub import hf_hub_download
    huggingface_hub.cached_download = hf_hub_download

import streamlit as st
import os
import time
import random
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from gtts import gTTS
import base64
from PIL import Image
import io
import matplotlib.pyplot as plt
import requests
from io import BytesIO
import json
import torch
from diffusers import DiffusionPipeline, StableDiffusionPipeline
import torch
from transformers import pipeline

# Configure Streamlit page
st.set_page_config(
    page_title="StoryCoder - Learn Python Through Stories",
    page_icon="🧙‍♂️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for colorful UI
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@700&display=swap');
    
    :root {
        --primary: #FF6B6B;
        --secondary: #4ECDC4;
        --accent: #FFD166;
        --dark: #1A535C;
        --light: #F7FFF7;
    }
    
    body {
        background: linear-gradient(135deg, var(--light) 0%, #E8F4F8 100%);
        font-family: 'Comic Neue', cursive;
    }
    
    .stApp {
        background: url('https://www.transparenttextures.com/patterns/cartographer.png');
    }
    
    .story-box {
        background-color: white;
        border-radius: 20px;
        padding: 25px;
        box-shadow: 0 8px 16px rgba(26, 83, 92, 0.15);
        border: 3px solid var(--accent);
        margin-bottom: 25px;
    }
    
    .header {
        color: var(--dark);
        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    }
    
    .concept-card {
        background: linear-gradient(145deg, #ffffff, #f0f0f0);
        border-radius: 15px;
        padding: 15px;
        margin: 10px 0;
        border-left: 5px solid var(--secondary);
        box-shadow: 0 4px 8px rgba(0,0,0,0.05);
    }
    
    .stButton>button {
        background: linear-gradient(45deg, var(--primary), var(--secondary));
        color: white;
        border-radius: 12px;
        padding: 10px 24px;
        font-weight: bold;
        font-size: 18px;
        border: none;
        transition: all 0.3s;
    }
    
    .stButton>button:hover {
        transform: scale(1.05);
        box-shadow: 0 6px 12px rgba(0,0,0,0.15);
    }
    
    .stTextInput>div>div>input {
        border-radius: 12px;
        padding: 12px;
        border: 2px solid var(--accent);
    }
    
    .tabs {
        display: flex;
        gap: 10px;
        margin-bottom: 20px;
        overflow-x: auto;
    }
    
    .tab {
        padding: 10px 20px;
        background-color: var(--accent);
        border-radius: 10px;
        cursor: pointer;
        font-weight: bold;
        white-space: nowrap;
    }
    
    .tab.active {
        background-color: var(--secondary);
        color: white;
    }
    
    @media (max-width: 768px) {
        .tabs {
            flex-wrap: wrap;
        }
    }
    
    .animation-container {
        background-color: #1a1a2e;
        border-radius: 15px;
        padding: 20px;
        box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
        margin-bottom: 25px;
        position: relative;
        overflow: hidden;
    }
    
    .animation-canvas {
        border-radius: 10px;
        overflow: hidden;
        margin: 0 auto;
        display: block;
        max-width: 100%;
    }
    
    .character {
        font-size: 48px;
        text-align: center;
        margin: 10px 0;
    }
    
    .audio-player {
        width: 100%;
        margin: 20px 0;
        border-radius: 20px;
        background: #f0f8ff;
        padding: 15px;
    }
    
    .animation-gif {
        max-width: 100%;
        border-radius: 10px;
        box-shadow: 0 8px 16px rgba(0,0,0,0.2);
    }
    
    .ai-animation {
        border: 3px solid var(--accent);
        border-radius: 15px;
        padding: 15px;
        background: white;
        margin: 20px 0;
    }
</style>
""", unsafe_allow_html=True)

# Concept database
CONCEPTS = {
    "loop": {
        "name": "Loop",
        "emoji": "🔄",
        "description": "Loops repeat actions multiple times",
        "example": "for i in range(5):\n    print('Hello!')",
        "color": "#FF9E6D"
    },
    "conditional": {
        "name": "Conditional",
        "emoji": "❓",
        "description": "Conditionals make decisions in code",
        "example": "if sunny:\n    go_outside()\nelse:\n    stay_inside()",
        "color": "#4ECDC4"
    },
    "function": {
        "name": "Function",
        "emoji": "✨",
        "description": "Functions are reusable blocks of code",
        "example": "def greet(name):\n    print(f'Hello {name}!')",
        "color": "#FFD166"
    },
    "variable": {
        "name": "Variable",
        "emoji": "📦",
        "description": "Variables store information",
        "example": "score = 10\nplayer = 'Alex'",
        "color": "#FF6B6B"
    },
    "list": {
        "name": "List",
        "emoji": "📝",
        "description": "Lists store collections of items",
        "example": "fruits = ['apple', 'banana', 'orange']",
        "color": "#1A535C"
    }
}

# Pre-generated animation examples
ANIMATION_EXAMPLES = {
    "loop": "https://i.imgur.com/7zQY1eE.gif",
    "conditional": "https://i.imgur.com/5X8jYAy.gif",
    "function": "https://i.imgur.com/9zJkQ7P.gif"
}

# Initialize AI models
@st.cache_resource
def load_models():
    """Load AI models for animation generation"""
    models = {}
    
    try:
        # Use CPU-friendly model and float32 precision
        models['text_to_image'] = DiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", 
            torch_dtype=torch.float32,  # Use float32 instead of float16
            safety_checker=None  # Disable safety checker for simplicity
        )
        # Move to GPU if available, otherwise keep on CPU
        device = "cuda" if torch.cuda.is_available() else "cpu"
        models['text_to_image'].to(device)
    except Exception as e:
        st.error(f"Could not load text-to-image model: {str(e)}")
        models['text_to_image'] = None
    
    try:
        # Text-to-speech model
        models['text_to_speech'] = pipeline("text-to-speech", model="suno/bark-small")
    except:
        models['text_to_speech'] = None
    
    return models

def analyze_story(story):
    """Analyze story and identify programming concepts"""
    story_lower = story.lower()
    detected_concepts = []
    
    # Check for loops
    if any(word in story_lower for word in ["times", "repeat", "again", "multiple"]):
        detected_concepts.append("loop")
    
    # Check for conditionals
    if any(word in story_lower for word in ["if", "when", "unless", "whether"]):
        detected_concepts.append("conditional")
    
    # Check for functions
    if any(word in story_lower for word in ["make", "create", "do", "perform", "cast"]):
        detected_concepts.append("function")
    
    # Check for variables
    if any(word in story_lower for word in ["is", "has", "set to", "value"]):
        detected_concepts.append("variable")
    
    # Check for lists
    if any(word in story_lower for word in ["and", "many", "several", "collection", "items"]):
        detected_concepts.append("list")
    
    return list(set(detected_concepts))

def extract_count_from_story(story):
    """Extract a number from the story to use in animations"""
    for word in story.split():
        if word.isdigit():
            return min(int(word), 10)
    return 3  # Default value

def create_story_scene(story, concept, models):
    """Create a story scene using AI text-to-image model"""
    try:
        # Create a prompt based on the story and concept
        style = "cartoon style, bright colors, children's book illustration"
        prompt = f"{story}. {CONCEPTS[concept]['name']} concept. {style}"
        
        # Generate image with fixed settings
        image = models['text_to_image'](
            prompt,
            num_inference_steps=25,  # Fewer steps for faster generation
            guidance_scale=7.5,
            height=512,
            width=512
        ).images[0]
        
        # Convert to bytes
        buf = io.BytesIO()
        image.save(buf, format='PNG')
        buf.seek(0)
        
        return buf
    
    except Exception as e:
        st.error(f"AI scene generation error: {str(e)}")
        return None

def create_animation_preview(story, concept):
    """Create an animation preview for the concept"""
    try:
        # Return a sample GIF for the concept
        return ANIMATION_EXAMPLES.get(concept, "https://i.imgur.com/7zQY1eE.gif")
    except:
        return "https://i.imgur.com/7zQY1eE.gif"

def text_to_speech_custom(text, models, filename="story_audio.wav"):
    """Convert text to speech using AI model"""
    try:
        if models['text_to_speech']:
            # Generate audio
            audio = models['text_to_speech'](text)
            
            # Save to file
            with open(filename, "wb") as f:
                f.write(audio["audio"])
            
            return filename
        
        # Fallback to gTTS
        tts = gTTS(text=text, lang='en')
        tts.save(filename)
        return filename
    except Exception as e:
        st.error(f"Audio creation error: {str(e)}")
        return None

def autoplay_audio(file_path):
    """Autoplay audio in Streamlit"""
    with open(file_path, "rb") as f:
        data = f.read()
        b64 = base64.b64encode(data).decode()
        md = f"""
            <audio autoplay>
            <source src="data:audio/wav;base64,{b64}" type="audio/wav">
            </audio>
            """
        st.markdown(md, unsafe_allow_html=True)

def generate_animation_code(story, concept):
    """Generate Python animation code based on story"""
    try:
        # Sample code based on concept
        if concept == "loop":
            code = f"""
# Loop Animation for: {story[:30]}...
import time

def animate_story():
    actions = {extract_count_from_story(story)}
    
    for i in range(actions):
        print(f"Action {{i+1}}: {story[:20]}...")
        # Animation code would go here
        time.sleep(0.5)
        
animate_story()
            """
            description = "Looping through actions multiple times"
            visual_cue = "Repeating actions in a loop"
        
        elif concept == "conditional":
            code = f"""
# Conditional Animation for: {story[:30]}...
import random

condition = random.choice([True, False])

if condition:
    print("Condition is True: {story[:20]}...")
    # True branch animation
else:
    print("Condition is False: {story[:20]}...")
    # False branch animation
            """
            description = "Making decisions based on conditions"
            visual_cue = "Branching paths based on conditions"
        
        else:  # function
            code = f"""
# Function Animation for: {story[:30]}...
def perform_action():
    print("Performing action: {story[:20]}...")
    # Action animation code

# Call the function multiple times
for _ in range({extract_count_from_story(story)}):
    perform_action()
            """
            description = "Reusing code with functions"
            visual_cue = "Calling a reusable function"
        
        return code, description, visual_cue
    
    except Exception as e:
        return f"# Error generating code\nprint('Could not generate code: {str(e)}')", "", ""

def create_interactive_animation(story, concept):
    """Create an interactive animation using Plotly"""
    try:
        # Create animation data
        frames = []
        count = extract_count_from_story(story)
        
        for i in range(count):
            frames.append({
                "frame": i,
                "x": np.random.uniform(1, 9),
                "y": np.random.uniform(1, 9),
                "size": np.random.uniform(20, 40),
                "label": f"Action {i+1}",
                "color": f"hsl({i*40}, 100%, 50%)"
            })
        
        df = pd.DataFrame(frames)
        
        # Create animated scatter plot
        fig = px.scatter(
            df, 
            x="x", 
            y="y", 
            animation_frame="frame",
            text="label",
            size="size",
            size_max=45,
            color="color",
            color_discrete_sequence=px.colors.qualitative.Alphabet,
            range_x=[0, 10],
            range_y=[0, 10],
            title=f"Interactive Animation: {story[:40]}{'...' if len(story) > 40 else ''}"
        )
        
        # Customize layout
        fig.update_layout(
            showlegend=False,
            plot_bgcolor='rgba(240,248,255,1)',
            paper_bgcolor='rgba(240,248,255,1)',
            width=800,
            height=600,
            xaxis=dict(showgrid=False, zeroline=False, visible=False),
            yaxis=dict(showgrid=False, zeroline=False, visible=False),
        )
        
        fig.update_traces(
            textfont_size=20,
            textposition='middle center',
            marker=dict(sizemode='diameter')
        )
        
        return fig
    
    except Exception as e:
        st.error(f"Interactive animation error: {str(e)}")
        return None

def main():
    """Main application function"""
    st.title("🧙‍♂️ StoryCoder - Learn Python Through Stories!")
    st.subheader("Turn your story into an animation and discover coding secrets!")
    
    # Load AI models
    ai_models = load_models()
    
    # Initialize session state
    if 'story' not in st.session_state:
        st.session_state.story = ""
    if 'concepts' not in st.session_state:
        st.session_state.concepts = []
    if 'story_scene' not in st.session_state:
        st.session_state.story_scene = None
    if 'interactive_animation' not in st.session_state:
        st.session_state.interactive_animation = None
    if 'animation_code' not in st.session_state:
        st.session_state.animation_code = ""
    if 'code_description' not in st.session_state:
        st.session_state.code_description = ""
    if 'visual_cue' not in st.session_state:
        st.session_state.visual_cue = ""
    if 'audio_file' not in st.session_state:
        st.session_state.audio_file = None
    if 'active_tab' not in st.session_state:
        st.session_state.active_tab = "story"
    
    # Create tabs
    tabs = st.empty()
    tab_cols = st.columns(5)
    with tab_cols[0]:
        if st.button("📖 Create Story"):
            st.session_state.active_tab = "story"
    with tab_cols[1]:
        if st.button("🎬 Animation"):
            st.session_state.active_tab = "animation"
    with tab_cols[2]:
        if st.button("🔍 Concepts"):
            st.session_state.active_tab = "concepts"
    with tab_cols[3]:
        if st.button("💻 Code"):
            st.session_state.active_tab = "code"
    with tab_cols[4]:
        if st.button("🔄 Reset"):
            st.session_state.story = ""
            st.session_state.concepts = []
            st.session_state.story_scene = None
            st.session_state.interactive_animation = None
            st.session_state.animation_code = ""
            st.session_state.code_description = ""
            st.session_state.visual_cue = ""
            st.session_state.audio_file = None
            st.session_state.active_tab = "story"
    
    # Story creation tab
    if st.session_state.active_tab == "story":
        with st.container():
            st.header("📖 Create Your Story")
            st.write("Write a short story (2-5 sentences) and I'll turn it into an animation!")
            
            story = st.text_area(
                "Your story:",
                height=200,
                placeholder="Once upon a time, a rabbit hopped 3 times to reach a carrot...",
                value=st.session_state.story,
                key="story_input"
            )
            
            if st.button("Create Animation!", use_container_width=True):
                if len(story) < 10:
                    st.error("Your story needs to be at least 10 characters long!")
                else:
                    st.session_state.story = story
                    with st.spinner("🧠 Analyzing your story for coding concepts..."):
                        st.session_state.concepts = analyze_story(story)
                    
                    # Get the main concept
                    main_concept = st.session_state.concepts[0] if st.session_state.concepts else "variable"
                    
                    with st.spinner("🎨 Generating AI story scene..."):
                        st.session_state.story_scene = create_story_scene(
                            story, main_concept, ai_models
                        )
                    
                    with st.spinner("🚀 Creating interactive animation..."):
                        st.session_state.interactive_animation = create_interactive_animation(
                            story, main_concept
                        )
                    
                    with st.spinner("💻 Generating animation code..."):
                        st.session_state.animation_code, st.session_state.code_description, st.session_state.visual_cue = generate_animation_code(
                            story, main_concept
                        )
                    
                    with st.spinner("🔊 Creating audio narration..."):
                        audio_text = f"Your story: {story}. This demonstrates the programming concept: {st.session_state.code_description}"
                        st.session_state.audio_file = text_to_speech_custom(
                            audio_text, ai_models
                        )
                    
                    st.session_state.active_tab = "animation"
                    st.rerun()
            
            # Show examples
            st.subheader("✨ Story Examples")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.caption("Loop Example")
                st.code('"A dragon breathes fire 5 times at the castle"', language="text")
                st.image(create_animation_preview("", "loop"), 
                         use_column_width=True, 
                         caption="Loop Animation Preview")
            with col2:
                st.caption("Conditional Example")
                st.code('"If it rains, the cat stays inside, else it goes out"', language="text")
                st.image(create_animation_preview("", "conditional"), 
                         use_column_width=True, 
                         caption="Conditional Animation Preview")
            with col3:
                st.caption("Function Example")
                st.code('"A wizard casts a spell to make flowers grow"', language="text")
                st.image(create_animation_preview("", "function"), 
                         use_column_width=True, 
                         caption="Function Animation Preview")
    
    # Animation tab
    elif st.session_state.active_tab == "animation":
        st.header("🎬 Your Story Animation")
        
        if not st.session_state.story:
            st.warning("Please create a story first!")
            st.session_state.active_tab = "story"
            st.rerun()
        
        # Display AI-generated scene
        st.subheader("🖼️ AI-Generated Story Scene")
        if st.session_state.story_scene:
            st.image(st.session_state.story_scene, use_container_width=True)
        else:
            st.warning("Scene couldn't be generated. Showing example instead.")
            concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
            st.image(create_animation_preview("", concept), 
                     use_container_width=True)
        
        # Display interactive animation
        st.subheader("🎮 Interactive Animation")
        if st.session_state.interactive_animation:
            st.plotly_chart(st.session_state.interactive_animation, use_container_width=True)
        else:
            st.info("Interactive animation preview. The real animation would run on your computer!")
            concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
            st.image(create_animation_preview("", concept), 
                    use_container_width=True)
        
        # Animation concept preview
        st.subheader("📽️ Animation Concept Preview")
        if st.session_state.concepts:
            concept = st.session_state.concepts[0]
            st.image(create_animation_preview("", concept), 
                     caption=f"{CONCEPTS[concept]['name']} Animation Example", 
                     use_container_width=True)
        else:
            st.info("Animation preview would appear here")
        
        # Play audio narration
        st.subheader("🔊 Story Narration")
        if st.session_state.audio_file:
            audio_bytes = open(st.session_state.audio_file, 'rb').read()
            st.audio(audio_bytes, format='audio/wav')
            
            if st.button("▶️ Play Automatically"):
                autoplay_audio(st.session_state.audio_file)
        else:
            st.info("Audio narration would play automatically here")
        
        st.success("✨ Animation created successfully!")
        st.caption("This animation was generated with Python code based on your story!")
        
        if st.button("Reveal Coding Secrets!", use_container_width=True):
            st.session_state.active_tab = "concepts"
            st.rerun()
    
    # Concepts tab
    elif st.session_state.active_tab == "concepts":
        st.header("🔍 Coding Concepts in Your Story")
        st.subheader("We secretly used these programming concepts:")
        
        if not st.session_state.concepts:
            st.warning("No concepts detected in your story! Try adding words like '3 times', 'if', or 'make'.")
        else:
            for concept in st.session_state.concepts:
                if concept in CONCEPTS:
                    details = CONCEPTS[concept]
                    st.markdown(f"""
                    <div class="concept-card" style="border-left: 5px solid {details['color']};">
                        <div style="display:flex; align-items:center; gap:15px;">
                            <span style="font-size:36px;">{details['emoji']}</span>
                            <h3 style="color:{details['color']};">{details['name']}</h3>
                        </div>
                        <p>{details['description']}</p>
                        <pre style="background:#f0f0f0; padding:10px; border-radius:8px;">{details['example']}</pre>
                    </div>
                    """, unsafe_allow_html=True)
                    
                    # Show animation preview for the concept
                    st.image(create_animation_preview("", concept), 
                             caption=f"{details['name']} Animation Example", 
                             use_column_width=True)
        
        if st.button("See the Magic Code!", use_container_width=True):
            st.session_state.active_tab = "code"
            st.rerun()
    
    # Code tab
    elif st.session_state.active_tab == "code":
        st.header("💻 The Magic Code Behind Your Animation")
        st.write("Here's the Python code that would create your animation:")
        
        if st.session_state.animation_code:
            st.subheader(st.session_state.code_description)
            st.markdown(f"**Visual Cue:** {st.session_state.visual_cue}")
            st.code(st.session_state.animation_code, language="python")
            
            # Download button
            st.download_button(
                label="Download Animation Code",
                data=st.session_state.animation_code,
                file_name="story_animation.py",
                mime="text/python",
                use_container_width=True
            )
            
            st.info("💡 To run this animation on your computer:")
            st.markdown("""
            1. Install Python from [python.org](https://python.org)
            2. Install required libraries: `pip install matplotlib numpy`
            3. Copy the code above into a file named `animation.py`
            4. Run it with: `python animation.py`
            """)
            
            st.success("🎉 When you run this code, you'll see your story come to life!")
            
            # Show what the animation would look like
            st.subheader("🎬 What Your Animation Would Look Like")
            concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
            st.image(create_animation_preview("", concept), 
                    caption="This is similar to what your animation would look like", 
                    use_column_width=True)
        else:
            st.warning("No code generated yet!")
        
        if st.button("Create Another Story!", use_container_width=True):
            st.session_state.active_tab = "story"
            st.session_state.story = ""
            st.rerun()

if __name__ == "__main__":
    main()