# app.py
import streamlit as st
import time
import random
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import textwrap
import os
import base64
from io import BytesIO

# Set up the page
st.set_page_config(
    page_title="CodeTales ✨",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@700&display=swap');
            
.stApp {
    background: linear-gradient(135deg, #6e8efb, #a777e3);
}
            
.header {
    font-family: 'Comic Neue', cursive;
    color: white;
    text-align: center;
    font-size: 3.5rem;
    text-shadow: 3px 3px 0 #000;
    padding: 1rem;
}
            
.subheader {
    font-family: 'Comic Neue', cursive;
    color: #ffeb3b;
    text-align: center;
    font-size: 1.8rem;
    margin-bottom: 2rem;
}
            
.story-box {
    background-color: rgba(255, 255, 255, 0.9);
    border-radius: 20px;
    padding: 2rem;
    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
    margin-bottom: 2rem;
}
            
.robot-speech {
    background-color: #4caf50;
    color: white;
    border-radius: 20px;
    padding: 1.5rem;
    font-size: 1.2rem;
    position: relative;
    margin-top: 2rem;
}
            
.robot-speech:after {
    content: '';
    position: absolute;
    top: -20px;
    left: 50px;
    border: 10px solid transparent;
    border-bottom-color: #4caf50;
}
            
.generate-btn {
    background: #ff5722 !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.2rem !important;
    padding: 0.7rem 2rem !important;
    border-radius: 50px !important;
    border: none !important;
    box-shadow: 0 4px 8px rgba(0,0,0,0.3) !important;
    transition: all 0.3s !important;
    margin-top: 1rem;
}
            
.generate-btn:hover {
    transform: scale(1.05) !important;
    box-shadow: 0 6px 12px rgba(0,0,0,0.4) !important;
}
            
.code-block {
    background: #2d2d2d;
    color: #f8f8f2;
    padding: 1rem;
    border-radius: 10px;
    font-family: monospace;
    font-size: 1.1rem;
    margin: 1rem 0;
    overflow-x: auto;
}

.animation-frame {
    border: 3px solid #ffeb3b;
    border-radius: 10px;
    margin: 5px;
}
</style>
""", unsafe_allow_html=True)

# Initialize AI models
@st.cache_resource
def load_models():
    """Load open-source AI models"""
    try:
        # Text generation model for code explanations
        story_to_code = pipeline("text-generation", model="gpt2", max_length=100)
        
        # Sentiment analysis for story understanding
        sentiment_analyzer = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
        
        # Named entity recognition for identifying objects
        ner_model = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english")
        
        return story_to_code, sentiment_analyzer, ner_model
    except Exception as e:
        st.error(f"Error loading models: {e}")
        return None, None, None

# Image generation functions
def create_storyboard_image(text, width=400, height=300):
    """Create a storyboard image from text"""
    # Create blank image
    img = Image.new('RGB', (width, height), color=(25, 25, 112))  # Dark blue background
    
    # Load a comic-style font (fallback to default if not available)
    try:
        font = ImageFont.truetype("comic.ttf", 16)
    except:
        font = ImageFont.load_default()
    
    draw = ImageDraw.Draw(img)
    
    # Draw title
    draw.text((10, 10), "Your Story Comes to Life!", fill=(255, 215, 0), font=font)
    
    # Draw text box
    draw.rectangle([10, 40, width-10, height-40], fill=(240, 248, 255), outline=(255, 215, 0), width=2)
    
    # Wrap text
    wrapped_text = textwrap.fill(text, width=40)
    draw.text((20, 50), wrapped_text, fill=(25, 25, 112), font=font)
    
    # Draw decorations
    draw.rectangle([width-50, height-30, width-30, height-10], fill=(220, 20, 60), outline=(255, 215, 0), width=1)
    draw.ellipse([20, height-50, 70, height], fill=(30, 144, 255), outline=(255, 215, 0), width=1)
    
    return img

def generate_sprite_animation(story, num_frames=4):
    """Generate a sprite-based animation from story"""
    frames = []
    width, height = 300, 200
    
    for i in range(num_frames):
        # Create base image
        img = Image.new('RGB', (width, height), color=(0, 0, 30))
        draw = ImageDraw.Draw(img)
        
        # Draw stars
        for _ in range(30):
            x = random.randint(0, width)
            y = random.randint(0, height)
            draw.ellipse([x, y, x+2, y+2], fill=(255, 255, 255))
        
        # Draw moving elements based on frame
        if "spaceship" in story.lower():
            ship_x = 50 + i * 60
            ship_y = 80
            draw.polygon([(ship_x, ship_y), (ship_x+30, ship_y), 
                          (ship_x+15, ship_y-20)], fill=(169, 169, 169))
            
            if "shoot" in story.lower() and i > 1:
                for j in range(3):
                    laser_x = ship_x + 15
                    laser_y = ship_y - 20 + j*5
                    draw.line([(laser_x, laser_y), (width, laser_y)], fill=(255, 0, 0), width=2)
        
        if "alien" in story.lower():
            alien_x = 200
            alien_y = 100 - i*10
            draw.ellipse([alien_x, alien_y, alien_x+20, alien_y+20], fill=(50, 205, 50))
            draw.ellipse([alien_x+5, alien_y+5, alien_x+7, alien_y+7], fill=(0, 0, 0))
            draw.ellipse([alien_x+13, alien_y+5, alien_x+15, alien_y+7], fill=(0, 0, 0))
        
        if "dragon" in story.lower():
            dragon_x = 150 + i*20
            dragon_y = 150
            draw.ellipse([dragon_x, dragon_y, dragon_x+40, dragon_y+20], fill=(178, 34, 34))
            draw.line([(dragon_x+40, dragon_y+10), (dragon_x+60, dragon_y)], fill=(178, 34, 34), width=3)
            
            if "fire" in story.lower() and i > 0:
                for j in range(5):
                    flame_x = dragon_x + 60
                    flame_y = dragon_y - j*5
                    flame_size = random.randint(5, 15)
                    draw.ellipse([flame_x, flame_y, flame_x+flame_size, flame_y+flame_size], 
                                fill=(255, random.randint(100, 200), 0))
        
        frames.append(img)
    
    return frames

def generate_code_explanation(story, story_to_code):
    """Generate code explanation using open-source model"""
    try:
        # Create a prompt for the model
        prompt = f"Explain how this story would be translated to code: '{story}'. Show a code snippet and explanation."
        
        # Generate text
        result = story_to_code(
            prompt, 
            max_length=150,
            num_return_sequences=1,
            temperature=0.7,
            top_k=50
        )
        
        return result[0]['generated_text']
    except:
        # Fallback explanation if model fails
        return f"""See how your story became real code? For example, when you wrote "{story.split()[0]}", 
we used code like: character.move(). That's how we turn words into actions!"""

# Header section
st.markdown('<div class="header">CodeTales ✨</div>', unsafe_allow_html=True)
st.markdown('<div class="subheader">Storytime + Coding Magic</div>', unsafe_allow_html=True)
st.markdown('<div style="text-align:center; color:white; font-size:1.2rem; margin-bottom:2rem;">Turn wild stories into playable games with open-source AI magic!</div>', unsafe_allow_html=True)

# How it works section
with st.expander("✨ How It Works (Like Baking a Cake) 🎂"):
    st.markdown("""
    **1. Kids Write a Story (The Recipe)**  
    Example: *"The spaceship zooms past aliens and shoots lasers!"*
    
    **2. AI is the Magic Oven 🔮**  
    We turn words into code using open-source AI models
    
    **3. Animation Pops Out (The Cake!) 🎮**  
    See your story come to life with custom animations!
    
    **4. Robot Teacher Explains 🤖**  
    Tavus shows: *"See? 'spaceship.move_right()' makes it fly! That's coding!"*
    """)

# Load models
story_to_code, sentiment_analyzer, ner_model = load_models()

# Initialize session state
if 'animation_generated' not in st.session_state:
    st.session_state.animation_generated = False
if 'story_text' not in st.session_state:
    st.session_state.story_text = ""
if 'animation_frames' not in st.session_state:
    st.session_state.animation_frames = []
if 'code_explanation' not in st.session_state:
    st.session_state.code_explanation = ""

# Main content
col1, col2 = st.columns([1, 1])

with col1:
    st.markdown('<div class="story-box">', unsafe_allow_html=True)
    st.markdown("### 📖 Write Your Story Here:")
    story_text = st.text_area(
        "Tell your adventure story...", 
        height=200,
        placeholder="Once upon a time, a brave spaceship zoomed through space, shooting lasers at alien spaceships...",
        label_visibility="collapsed",
        value=st.session_state.story_text
    )
    
    # Generate button with animation
    if st.button("✨ Generate Animation!", use_container_width=True, key="generate", type="primary"):
        if story_text.strip():
            st.session_state.story_text = story_text
            st.session_state.animation_generated = True
            
            with st.spinner("🧙‍♂️ Cooking your story in the magic oven..."):
                # Generate animation frames
                st.session_state.animation_frames = generate_sprite_animation(story_text)
                
                # Generate code explanation
                st.session_state.code_explanation = generate_code_explanation(story_text, story_to_code)
                
                # Analyze story sentiment (just for fun)
                if sentiment_analyzer:
                    sentiment = sentiment_analyzer(story_text[:512])[0]
                    st.session_state.sentiment = f"Your story feels {sentiment['label'].lower()}! (Confidence: {sentiment['score']:.2f})"
                else:
                    st.session_state.sentiment = "Your story is full of adventure!"
                
                # Identify story characters
                if ner_model and len(story_text) > 10:
                    entities = ner_model(story_text[:256])
                    characters = list(set([ent['word'] for ent in entities if ent['entity'] in ['B-PER', 'I-PER']]))
                    if characters:
                        st.session_state.characters = f"Main characters: {', '.join(characters)}"
                    else:
                        st.session_state.characters = "Exciting characters are ready for action!"
                else:
                    st.session_state.characters = "Your heroes are preparing for adventure!"
        else:
            st.warning("Please enter a story first!")
    st.markdown('</div>', unsafe_allow_html=True)

with col2:
    st.markdown('<div class="story-box">', unsafe_allow_html=True)
    st.markdown("### 🎮 Your Game Animation")
    
    if st.session_state.animation_generated and st.session_state.animation_frames:
        # Display animation frames
        st.markdown("**Your Story Comes to Life!**")
        cols = st.columns(len(st.session_state.animation_frames))
        for i, frame in enumerate(st.session_state.animation_frames):
            with cols[i]:
                st.image(frame, caption=f"Frame {i+1}", use_container_width=True)
        
        # Create an animated GIF
        gif_buffer = BytesIO()
        st.session_state.animation_frames[0].save(
            gif_buffer,
            format='GIF',
            save_all=True,
            append_images=st.session_state.animation_frames[1:],
            duration=500,
            loop=0
        )
        gif_data = gif_buffer.getvalue()
        st.download_button(
            label="⬇️ Download Animation",
            data=gif_data,
            file_name="your_story.gif",
            mime="image/gif",
            use_container_width=True
        )
        
        # Display story analysis
        st.success("✨ Animation created!")
        st.info(f"🎭 {st.session_state.characters}")
        st.info(f"😊 {st.session_state.sentiment}")
        
    elif st.session_state.animation_generated:
        st.warning("Couldn't generate animation. Try a different story!")
        st.image("https://img.freepik.com/free-vector/hand-drawn-colorful-space-background_23-2148821798.jpg", 
                 use_container_width=True)
    elif story_text:
        # Show storyboard preview
        preview_img = create_storyboard_image(story_text)
        st.image(preview_img, caption="Your Story Preview", use_container_width=True)
        st.info("👆 Click Generate to bring your story to life!")
    else:
        st.image("https://img.freepik.com/free-vector/hand-drawn-colorful-space-background_23-2148821798.jpg", 
                 use_container_width=True)
        st.info("👆 Write your story and click Generate to see the magic!")
    
    st.markdown('</div>', unsafe_allow_html=True)

# Tavus explanation section
if st.session_state.animation_generated and st.session_state.story_text:
    st.markdown('<div class="robot-speech">', unsafe_allow_html=True)
    st.markdown("### 🤖 Tavus the Robot Teacher says:")
    
    # Extract action words from story
    action_words = ["zoom", "shoot", "fly", "move", "jump", "run", "attack", "laser", "alien", "spaceship", 
                   "dragon", "fire", "hero", "sword", "castle", "run", "escape", "fight", "win"]
    found_actions = [word for word in action_words if word in st.session_state.story_text.lower()]
    
    if found_actions:
        # Create explanation based on found words
        explanations = []
        code_snippets = []
        
        for action in found_actions[:3]:  # Limit to 3 explanations
            if action == "zoom":
                explanations.append("makes your spaceship move super fast!")
                code_snippets.append("spaceship.accelerate(speed=10)")
            elif action == "shoot":
                explanations.append("creates laser beams to defeat enemies!")
                code_snippets.append("laser = spaceship.fire_weapon()")
            elif action == "fly":
                explanations.append("keeps your character moving through the air!")
                code_snippets.append("character.apply_gravity(False)")
            elif action == "move":
                explanations.append("changes position on the screen!")
                code_snippets.append("player.move(direction='right')")
            elif action == "jump":
                explanations.append("makes your character leap upwards!")
                code_snippets.append("hero.jump(height=100)")
            elif action == "run":
                explanations.append("makes your character move faster!")
                code_snippets.append("player.speed = player.speed * 2")
            elif action == "attack":
                explanations.append("lets your hero fight the bad guys!")
                code_snippets.append("sword.swing(damage=15)")
            elif action == "laser":
                explanations.append("creates powerful energy beams!")
                code_snippets.append("weapon = Laser(color='red', damage=20)")
            elif action == "alien":
                explanations.append("adds enemies to your game!")
                code_snippets.append("enemy = Alien(position=(100, 200))")
            elif action == "spaceship":
                explanations.append("creates your main character!")
                code_snippets.append("player = Spaceship(image='spaceship.png')")
            elif action == "dragon":
                explanations.append("adds a fire-breathing challenge!")
                code_snippets.append("boss = Dragon(health=100, damage=25)")
            elif action == "fire":
                explanations.append("creates dangerous fire effects!")
                code_snippets.append("flame = FireEffect(position, size=30)")
            elif action == "hero":
                explanations.append("creates the main player character!")
                code_snippets.append("player = Hero(name='Your Hero')")
            elif action == "castle":
                explanations.append("adds a majestic building to your world!")
                code_snippets.append("castle = Building(type='castle', position=(300, 150))")
            elif action == "escape":
                explanations.append("creates exciting chase sequences!")
                code_snippets.append("player.start_escape_sequence(speed=15)")
            elif action == "fight":
                explanations.append("initiates battle mechanics!")
                code_snippets.append("initiate_combat(enemy)")
            elif action == "win":
                explanations.append("creates victory conditions!")
                code_snippets.append("if player.score > 100: show_victory_screen()")
        
        st.markdown("See how your story became real code? Here's what happened:")
        
        for i, (action, explanation, snippet) in enumerate(zip(found_actions, explanations, code_snippets)):
            st.markdown(f"**{i+1}. When you said '{action}'**:")
            st.markdown(f'<div class="code-block">{snippet}</div>', unsafe_allow_html=True)
            st.markdown(f"This code {explanation}")
        
        st.markdown("That's the magic of coding - turning words into actions!")
    else:
        st.markdown("I see you created an awesome story! Every word you write can become game code. Try adding action words like 'jump', 'run', or 'shoot' next time!")
    
    # Show AI-generated explanation
    st.markdown("### 🧠 AI-Powered Explanation:")
    st.write(st.session_state.code_explanation)
    
    st.markdown("</div>", unsafe_allow_html=True)

# Benefits section
st.markdown("""
## ❤ Why Everyone Will Love CodeTales

| For Kids 👧👦 | For Parents & Teachers 👪👩‍🏫 |
|--------------|----------------------------|
| ✨ Feels like playing, not learning | 🧠 Secretly teaches programming concepts |
| 🚀 Brings imagination to life | 🔍 Develops logical thinking skills |
| 🎮 Creates personal video games | ➗ Reinforces math fundamentals |
| 😄 Makes learning fun and exciting | 🧩 Encourages problem-solving abilities |
| 🌟 100% private with no API keys | 💰 Completely free and open-source |
""")

# Footer
st.markdown("---")
st.markdown("""
<center>
    <p style="color:white; font-size:1.1rem;">
        ✨ Made with open-source magic by CodeTales Team ✨<br>
        Transforming stories into games since 2023 | No API Keys Required!
    </p>
</center>
""", unsafe_allow_html=True)