import os
import streamlit as st
from dotenv import load_dotenv
from novita_client import NovitaClient
import ast
import subprocess
import tempfile
import time

load_dotenv()

# Initialize Novita API client
class StoryToCodeConverter:
    def __init__(self):
        api_key = os.getenv("NOVITA_API_KEY")
        if not api_key:
            raise ValueError("NOVITA_API_KEY not found in environment")
        self.client = NovitaClient(api_key)

    def generate_code(self, story):
        prompt = f"""
Convert this children's story into Python animation code using turtle graphics.
Add code comments that explain the basic concepts to kids, like loops, movement, and colors.

Story: {story}

Only return the Python code wrapped in triple backticks like ```python ... ```
"""
        response = self.client.chat_completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}]
        )
        content = response.choices[0].message.content
        return self._extract_code(content)

    def _extract_code(self, content):
        if "```python" in content:
            return content.split("```python")[1].split("```")[0]
        return content


# Safe execution
class CodeExecutor:
    def execute_safely(self, code):
        try:
            ast.parse(code)  # Safety check
            with tempfile.NamedTemporaryFile(delete=False, suffix=".py") as tmp:
                tmp.write(code.encode("utf-8"))
                tmp_path = tmp.name
            subprocess.Popen(["python", tmp_path])
            return True
        except Exception as e:
            st.error(f"Robot Chef burned the cake: {e}")
            return False


# ---------- Streamlit UI ----------
st.set_page_config(page_title="CodeTales", page_icon="✨")

with st.sidebar:
    st.header("🍰 How It Works")
    st.markdown("""
1. 📝 Kids write a story (like a recipe!)  
2. 🔮 AI turns it into **Python code**  
3. 🎮 Animation appears like magic!  
4. 🤖 Robot Teacher explains the code!

Made for ages 6–12. Fun, visual, and interactive!
""")

st.title("✨ CodeTales: Storytime + Coding Magic")

story = st.text_area("Write Your Story Recipe 📝", height=200,
                     placeholder="Once upon a time, a rocket zoomed past rainbow stars and met an alien...")

if st.button("✨ Bake My Story Cake!"):
    if not story.strip():
        st.warning("Please write a story first!")
    else:
        with st.spinner("Baking your story cake... 🍰"):
            converter = StoryToCodeConverter()
            try:
                code = converter.generate_code(story)
            except Exception as e:
                st.error(f"AI oven broke down! {e}")
                st.stop()

            st.subheader("📜 Your Code Recipe")
            st.code(code, language='python')

            executor = CodeExecutor()
            if executor.execute_safely(code):
                st.balloons()
                with st.expander("🤖 Robot Teacher Says:"):
                    st.markdown("""
- `turtle.forward(100)` – Moves the rocket forward 🚀  
- `turtle.color("red")` – Paints with colors 🎨  
- `for loop` – Repeats an action like magic 🔁  
                    """)
                st.success("🎉 Your story came to life!")
            else:
                st.error("Oops! Couldn't bake the animation.")
else:
    st.caption("Turn your words into animated code with magic 🪄")

st.markdown("---")
st.caption("Made with ❤️ for young coders | Powered by Novita + Python Turtle")