# 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

# Use compatible import for diffusers
try:
    from diffusers import StableDiffusionPipeline
except ImportError:
    from diffusers import DiffusionPipeline as StableDiffusionPipeline

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 a fast, optimized model
        models['text_to_image'] = StableDiffusionPipeline.from_pretrained(
            "OFA-Sys/small-stable-diffusion-v0",  # Optimized for speed
            torch_dtype=torch.float32,
            safety_checker=None
        )
        device = "cuda" if torch.cuda.is_available() else "cpu"
        models['text_to_image'] = 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 optimized settings
        image = models['text_to_image'](
            prompt,
            num_inference_steps=15,  # Faster generation
            guidance_scale=7.5,
            height=256,  # Smaller image for faster generation
            width=256
        ).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!")
    
    # 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"
                    
                    # Load models only when needed
                    with st.spinner("🚀 Loading AI models..."):
                        ai_models = load_models()
                    
                    with st.spinner("🎨 Generating AI story scene (this may take 5-10 seconds)..."):
                        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_container_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_container_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_container_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_container_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_container_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()