import streamlit as st
import os
import time
import base64
import random
import tempfile
import subprocess
import numpy as np
import pygame
import matplotlib.pyplot as plt
from PIL import Image
import io
from manim import *

# 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;
    }
</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"
    }
}

# Character database
CHARACTERS = {
    "rabbit": {"color": (255, 150, 150), "speed": 5, "size": 30},
    "dragon": {"color": (255, 100, 100), "speed": 3, "size": 60},
    "cat": {"color": (200, 150, 255), "speed": 6, "size": 25},
    "dog": {"color": (255, 200, 150), "speed": 7, "size": 35},
    "knight": {"color": (150, 200, 255), "speed": 4, "size": 40},
    "wizard": {"color": (200, 255, 150), "speed": 4, "size": 35},
    "scientist": {"color": (150, 255, 200), "speed": 5, "size": 30},
    "pirate": {"color": (255, 255, 150), "speed": 5, "size": 40}
}

# Animation templates
ANIMATION_TEMPLATES = {
    "loop": {
        "pygame": """
import pygame
import sys
import math

# Initialize pygame
pygame.init()

# Screen setup
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Your Story Animation")

# Colors
BACKGROUND = (25, 25, 50)
GROUND = (40, 100, 40)
SUN = (255, 255, 200)

# Character setup
character = "{character}"
char_color = {char_color}
char_size = {char_size}
char_x = 100
char_y = HEIGHT - 100
char_speed = {char_speed}

# Target setup
target_x = 700
target_y = HEIGHT - 100
target_reached = False
hops = 0

# Animation clock
clock = pygame.time.Clock()
frames = []

# Main animation loop
for frame in range(120):
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
    
    # Clear screen
    screen.fill(BACKGROUND)
    
    # Draw ground
    pygame.draw.rect(screen, GROUND, (0, HEIGHT - 50, WIDTH, 50))
    
    # Draw sun
    pygame.draw.circle(screen, SUN, (700, 100), 50)
    
    # Draw target
    pygame.draw.circle(screen, (255, 200, 100), (target_x, target_y), 30)
    pygame.draw.circle(screen, (255, 150, 50), (target_x, target_y), 20)
    pygame.draw.circle(screen, (255, 100, 0), (target_x, target_y), 10)
    
    # Move character
    if not target_reached:
        char_x += char_speed
        # Hop animation
        hop_height = 50 * math.sin(frame * 0.2)
        char_y = HEIGHT - 100 - abs(hop_height)
        
        # Check if target reached
        if char_x >= target_x - 40:
            hops += 1
            char_x = 100
            if hops >= {count}:
                target_reached = True
    
    # Draw character
    pygame.draw.circle(screen, char_color, (int(char_x), int(char_y)), char_size)
    
    # Draw eyes
    pygame.draw.circle(screen, (255, 255, 255), (int(char_x + 10), int(char_y - 5)), 8)
    pygame.draw.circle(screen, (0, 0, 0), (int(char_x + 10), int(char_y - 5)), 4)
    
    # Draw hop counter
    font = pygame.font.SysFont(None, 36)
    text = font.render(f"Hops: {{hops}}/{{count}}", True, (255, 255, 255))
    screen.blit(text, (20, 20))
    
    # Capture frame
    frames.append(pygame.surfarray.array3d(screen).swapaxes(0, 1))
    
    pygame.display.flip()
    clock.tick(30)

# Save animation as GIF
import imageio
imageio.mimsave('animation.gif', frames, fps=30)
        """
    },
    "conditional": {
        "pygame": """
import pygame
import sys

# Initialize pygame
pygame.init()

# Screen setup
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Your Story Animation")

# Colors
BACKGROUND = (25, 25, 50)
GROUND = (40, 100, 40)
CHEST_COLOR = (139, 69, 19)
LOCK_COLOR = (192, 192, 192)
TREASURE_COLOR = (255, 215, 0)

# Knight setup
knight_x = 100
knight_y = HEIGHT - 100
knight_speed = 4
knight_size = 40
knight_color = (100, 150, 255)

# Chest setup
chest_x = 600
chest_y = HEIGHT - 100
chest_width = 80
chest_height = 50
is_unlocked = {condition}

# Animation clock
clock = pygame.time.Clock()
frames = []

# Main animation loop
for frame in range(180):
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
    
    # Clear screen
    screen.fill(BACKGROUND)
    
    # Draw ground
    pygame.draw.rect(screen, GROUND, (0, HEIGHT - 50, WIDTH, 50))
    
    # Draw chest
    pygame.draw.rect(screen, CHEST_COLOR, (chest_x, chest_y, chest_width, chest_height))
    pygame.draw.rect(screen, (101, 67, 33), (chest_x, chest_y, chest_width, 20))
    
    # Draw lock or treasure
    if not is_unlocked:
        pygame.draw.rect(screen, LOCK_COLOR, (chest_x + chest_width//2 - 10, chest_y - 15, 20, 15))
        pygame.draw.circle(screen, LOCK_COLOR, (chest_x + chest_width//2, chest_y - 15), 8)
    else:
        pygame.draw.rect(screen, (255, 223, 0), (chest_x + 10, chest_y + 10, 20, 30))
        pygame.draw.rect(screen, (218, 165, 32), (chest_x + 50, chest_y + 15, 20, 20))
    
    # Move knight
    if knight_x < chest_x - 100:
        knight_x += knight_speed
    
    # Draw knight
    pygame.draw.circle(screen, knight_color, (knight_x, knight_y), knight_size)
    pygame.draw.rect(screen, (50, 50, 150), (knight_x - 15, knight_y - 50, 30, 50))  # Body
    pygame.draw.circle(screen, (200, 200, 200), (knight_x, knight_y - 60), 15)  # Head
    
    # Draw text
    font = pygame.font.SysFont(None, 36)
    if is_unlocked:
        text = font.render("Chest is unlocked! Treasure acquired!", True, (255, 255, 100))
    else:
        text = font.render("Chest is locked! Find the key!", True, (255, 100, 100))
    screen.blit(text, (WIDTH//2 - text.get_width()//2, 50))
    
    # Capture frame
    frames.append(pygame.surfarray.array3d(screen).swapaxes(0, 1))
    
    pygame.display.flip()
    clock.tick(30)

# Save animation as GIF
import imageio
imageio.mimsave('animation.gif', frames, fps=30)
        """
    },
    "function": {
        "manim": """
from manim import *

class FunctionMagic(Scene):
    def construct(self):
        # Create wizard and flower
        wizard = Circle(radius=0.5, color=BLUE, fill_opacity=1).shift(LEFT*3)
        wizard_label = Text("Wizard").next_to(wizard, DOWN)
        
        flower = VGroup(
            Circle(radius=0.3, color=GREEN).shift(UP*0.3),
            Circle(radius=0.2, color=YELLOW).shift(UP*0.3),
            Line(UP*0.3, DOWN*1, color=GREEN)
        ).shift(RIGHT*3)
        flower_label = Text("Flower").next_to(flower, DOWN)
        
        # Create spell function
        func_label = Text("cast_grow_spell()", color=YELLOW).shift(UP*2)
        
        # Animation
        self.play(Create(wizard), Write(wizard_label))
        self.play(Create(flower), Write(flower_label))
        self.play(Write(func_label))
        self.wait(1)
        
        # Cast spell multiple times
        for i in range({count}):
            # Create spell effect
            spell = VGroup(
                Star(color=PURPLE, fill_opacity=1).scale(0.3),
                MathTex("\\star", color=PURPLE).scale(1.5)
            )
            
            # Animate spell moving
            self.play(
                Create(spell.move_to(wizard.get_center())),
                run_time=0.5
            )
            self.play(
                spell.animate.move_to(flower.get_center()),
                run_time=1
            )
            self.play(FadeOut(spell))
            
            # Grow flower
            self.play(
                flower.animate.scale(1.2),
                run_time=0.5
            )
            self.wait(0.5)
            
            # Show function call
            call_text = Text(f"cast_grow_spell()  #{i+1}", color=GREEN).shift(DOWN*2)
            self.play(Write(call_text))
            self.play(FadeOut(call_text))
        
        self.wait(2)
        """
    }
}

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 generate_animation_code(story, concepts):
    """Generate Python animation code based on story"""
    try:
        # Choose a random character
        character = random.choice(list(CHARACTERS.keys()))
        char_details = CHARACTERS[character]
        
        # Choose the most relevant concept
        concept = concepts[0] if concepts else "loop"
        
        # Get template for the concept
        template = ANIMATION_TEMPLATES.get(concept, ANIMATION_TEMPLATES["loop"])
        
        # Fill in template parameters
        if concept == "loop":
            # Extract count from story
            count = 3
            for word in story.split():
                if word.isdigit():
                    count = min(int(word), 10)
                    break
            
            code = template["pygame"].format(
                character=character,
                char_color=char_details["color"],
                char_size=char_details["size"],
                char_speed=char_details["speed"],
                count=count
            )
        elif concept == "conditional":
            condition = "True" if random.random() > 0.5 else "False"
            code = template["pygame"].format(condition=condition)
        elif concept == "function":
            count = 3
            for word in story.split():
                if word.isdigit():
                    count = min(int(word), 5)
                    break
            code = template["manim"].format(count=count)
        else:
            # Default to loop animation
            code = ANIMATION_TEMPLATES["loop"]["pygame"].format(
                character=character,
                char_color=char_details["color"],
                char_size=char_details["size"],
                char_speed=char_details["speed"],
                count=3
            )
        
        return code
    except Exception as e:
        return f"# Error generating code\nprint('Could not generate code: {str(e)}')"

def run_pygame_animation(code):
    """Run PyGame animation code and generate GIF"""
    try:
        # Create a temporary Python file
        with tempfile.NamedTemporaryFile(suffix=".py", delete=False, mode='w') as tmpfile:
            tmpfile.write(code)
            tmpfile_path = tmpfile.name
        
        # Run the animation script
        subprocess.run(["python", tmpfile_path], check=True)
        
        # Check for generated GIF
        if os.path.exists("animation.gif"):
            return "animation.gif"
        
        return None
    except Exception as e:
        st.error(f"Animation error: {str(e)}")
        return None

def run_manim_animation(code):
    """Run Manim animation code and generate video"""
    try:
        # Create a temporary Python file
        with tempfile.NamedTemporaryFile(suffix=".py", delete=False, mode='w') as tmpfile:
            tmpfile.write(code)
            tmpfile_path = tmpfile.name
        
        # Run Manim render
        result = subprocess.run(
            ["manim", "-ql", "-o", "animation", tmpfile_path],
            capture_output=True,
            text=True
        )
        
        # Find the output file
        output_dir = "media/videos"
        if os.path.exists(output_dir):
            for file in os.listdir(output_dir):
                if file.startswith("animation") and file.endswith(".mp4"):
                    return os.path.join(output_dir, file)
        
        return None
    except Exception as e:
        st.error(f"Manim error: {str(e)}")
        return None

def create_story_image(story):
    """Create a story image with Matplotlib"""
    try:
        # Create figure
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.set_facecolor('#f0f8ff')
        ax.set_xlim(0, 10)
        ax.set_ylim(0, 10)
        ax.axis('off')
        
        # Add title
        ax.text(5, 8, '✨ Your Story ✨', fontsize=24, 
                ha='center', color='purple', fontweight='bold')
        
        # Add story text (wrapped)
        words = story.split()
        lines = []
        current_line = []
        
        for word in words:
            if len(' '.join(current_line + [word])) < 60:
                current_line.append(word)
            else:
                lines.append(' '.join(current_line))
                current_line = [word]
        
        if current_line:
            lines.append(' '.join(current_line))
        
        for i, line in enumerate(lines):
            ax.text(5, 6.5 - i*0.7, line, fontsize=14, ha='center')
        
        # Add decoration
        ax.text(2, 2, '🐰', fontsize=40, ha='center')
        ax.text(8, 2, '🐉', fontsize=40, ha='center')
        ax.text(5, 1, 'Created with StoryCoder', fontsize=12, 
                ha='center', style='italic', color='gray')
        
        # Save to buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.5)
        buf.seek(0)
        image = Image.open(buf)
        plt.close()
        
        return image
    except Exception as e:
        st.error(f"Image generation 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 'animation_code' not in st.session_state:
        st.session_state.animation_code = ""
    if 'active_tab' not in st.session_state:
        st.session_state.active_tab = "story"
    if 'animation_path' not in st.session_state:
        st.session_state.animation_path = None
    if 'animation_type' not in st.session_state:
        st.session_state.animation_type = None
    
    # 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.animation_code = ""
            st.session_state.animation_path = None
            st.session_state.animation_type = 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)
                    
                    with st.spinner("✨ Generating animation code..."):
                        st.session_state.animation_code = generate_animation_code(
                            story, st.session_state.concepts
                        )
                    
                    # Determine animation type
                    if "manim" in st.session_state.animation_code:
                        st.session_state.animation_type = "manim"
                        with st.spinner("🎬 Rendering Manim animation (this may take a minute)..."):
                            st.session_state.animation_path = run_manim_animation(
                                st.session_state.animation_code
                            )
                    else:
                        st.session_state.animation_type = "pygame"
                        with st.spinner("🎬 Creating PyGame animation..."):
                            st.session_state.animation_path = run_pygame_animation(
                                st.session_state.animation_code
                            )
                    
                    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")
            with col2:
                st.caption("Conditional Example")
                st.code('"If it rains, the cat stays inside, else it goes out"', language="text")
            with col3:
                st.caption("Function Example")
                st.code('"A wizard casts a spell to make flowers grow 3 times"', language="text")
    
    # Animation tab
    elif st.session_state.active_tab == "animation":
        st.header("🎬 Your Story Animation")
        
        if not st.session_state.animation_path:
            st.warning("Please create a story first!")
            st.session_state.active_tab = "story"
            st.rerun()
        
        # Display animation
        st.markdown(f"""
        <div class="animation-container">
            <h3 style="color: white; text-align: center;">"{st.session_state.story[:60]}{'...' if len(st.session_state.story) > 60 else ''}"</h3>
        </div>
        """, unsafe_allow_html=True)
        
        try:
            if st.session_state.animation_type == "pygame":
                st.image(st.session_state.animation_path, use_container_width=True)
            elif st.session_state.animation_type == "manim":
                st.video(st.session_state.animation_path)
        except Exception as e:
            st.error(f"Couldn't display animation: {str(e)}")
            story_image = create_story_image(st.session_state.story)
            if story_image:
                st.image(story_image, use_container_width=True)
        
        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)
        
        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 created your animation:")
        
        if st.session_state.animation_code:
            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.write("You can run this code on your computer to see the animation!")
        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()