Update app.py

app.py

@@ -9,11 +9,6 @@ import numpy as np
 import sys
 import subprocess
 import json
-from pygments import highlight
-from pygments.lexers import PythonLexer
-from pygments.formatters import HtmlFormatter
-import base64
-import torch
 import re
 import shutil
 import time
@@ -25,67 +20,62 @@ import plotly.express as px
 import zipfile
 import traceback
 
-# Set up enhanced logging
+# Set up logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    format="%(asctime)s %(levelname)s %(name)s – %(message)s",
     handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
-# Model configuration mapping
+# Available quality presets and formats
+QUALITY_PRESETS = ["480p", "720p", "1080p", "4K", "8K"]
+FPS_OPTIONS = [24, 30, 60, 120]
+OUTPUT_FORMATS = {
+    "MP4 Video": "mp4",
+    "GIF Animation": "gif",
+    "WebM Video": "webm",
+    "PNG Sequence (ZIP)": "png_sequence",
+    "SVG Image": "svg"
+}
+
+# Model configurations
 MODEL_CONFIGS = {
-    "DeepSeek-V3-0324": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek"},
-    "DeepSeek-R1": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek"},
-    "Llama-4-Scout-17B-16E-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
-    "Llama-4-Maverick-17B-128E-Instruct-FP8": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
-    "gpt-4o-mini": {"max_tokens": 15000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "gpt-4o": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "gpt-4.1": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "gpt-4.1-mini": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "gpt-4.1-nano": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "o3-mini": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI"},
-    "o1": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI"},
-    "o1-mini": {"max_completion_tokens": 66000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI"},
-    "o1-preview": {"max_tokens": 33000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
-    "Phi-4-multimodal-instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Microsoft"},
-    "Mistral-large-2407": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral"},
-    "Codestral-2501": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral"},
-    "default": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Other"}
+    "gpt-4o": {},
+    "gpt-4o-mini": {},
+    "gpt-4.1": {},
+    "gpt-4.1-mini": {},
+    "o1": {},
+    "o1-mini": {},
+    "default": {}
 }
 
-# Try to import Streamlit Ace
+# Try to import st_ace
 try:
     from streamlit_ace import st_ace
     ACE_EDITOR_AVAILABLE = True
 except ImportError:
     ACE_EDITOR_AVAILABLE = False
-    logger.warning("streamlit-ace not available, falling back to text area")
-
-def prepare_api_params(messages, model_name):
-    config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
-    params = {"messages": messages, "model": model_name}
-    params[config["param_name"]] = config.get(config["param_name"])
-    return params, config
+    logger.warning("streamlit-ace not available, using text_area")
 
 def get_secret(env_var):
     val = os.environ.get(env_var)
     if not val:
-        logger.warning(f"Secret '{env_var}' not found")
+        logger.warning(f"Secret '{env_var}' not set")
     return val
 
 def check_password():
-    correct = get_secret("password")
-    if not correct:
+    pwd = get_secret("password")
+    if not pwd:
         st.error("Admin password not configured")
         return False
-    if "password_entered" not in st.session_state:
-        st.session_state.password_entered = False
-    if not st.session_state.password_entered:
-        pwd = st.text_input("Enter password to access AI features", type="password")
-        if pwd:
-            if pwd == correct:
-                st.session_state.password_entered = True
+    if "pwd_ok" not in st.session_state:
+        st.session_state.pwd_ok = False
+    if not st.session_state.pwd_ok:
+        entry = st.text_input("Enter password", type="password")
+        if entry:
+            if entry == pwd:
+                st.session_state.pwd_ok = True
                 return True
             else:
                 st.error("Incorrect password")
@@ -94,18 +84,23 @@ def check_password():
     return True
 
 def ensure_packages():
-    required = {
-        'manim': '0.17.3', 'Pillow': '9.0.0', 'numpy': '1.22.0',
-        'transformers': '4.30.0', 'torch': '2.0.0', 'pygments': '2.15.1',
-        'streamlit-ace': '0.1.1', 'pydub': '0.25.1', 'plotly': '5.14.0',
-        'pandas': '2.0.0', 'python-pptx': '0.6.21', 'fpdf': '1.7.2',
-        'matplotlib': '3.5.0', 'seaborn': '0.11.2', 'scipy': '1.7.3',
-        'huggingface_hub': '0.16.0'
+    reqs = {
+        "manim": "0.17.3",
+        "Pillow": "9.0.0",
+        "numpy": "1.22.0",
+        "plotly": "5.14.0",
+        "pandas": "2.0.0",
+        "python-pptx": "0.6.21",
+        "fpdf": "1.7.2",
+        "matplotlib": "3.5.0",
+        "seaborn": "0.11.2",
+        "scipy": "1.7.3",
+        "streamlit-ace": "0.1.1"
     }
     missing = {}
-    for pkg, ver in required.items():
+    for pkg, ver in reqs.items():
         try:
-            __import__(pkg if pkg != 'Pillow' else 'PIL')
+            __import__(pkg if pkg != "Pillow" else "PIL")
         except ImportError:
             missing[pkg] = ver
     if not missing:
@@ -113,88 +108,36 @@ def ensure_packages():
     bar = st.progress(0)
     txt = st.empty()
     for i, (pkg, ver) in enumerate(missing.items()):
-        bar.progress(i / len(missing))
         txt.text(f"Installing {pkg}...")
         res = subprocess.run([sys.executable, "-m", "pip", "install", f"{pkg}>={ver}"], capture_output=True, text=True)
         if res.returncode != 0:
             st.error(f"Failed to install {pkg}")
             return False
-    bar.progress(1.0)
+        bar.progress((i + 1) / len(missing))
     txt.empty()
     return True
 
-def install_custom_packages(pkgs):
-    if not pkgs.strip():
-        return True, "No packages specified"
-    parts = [p.strip() for p in pkgs.split(",") if p.strip()]
-    if not parts:
-        return True, "No valid packages"
-    sidebar_txt = st.sidebar.empty()
-    bar = st.sidebar.progress(0)
-    results, success = [], True
-    for i, p in enumerate(parts):
-        bar.progress(i / len(parts))
-        sidebar_txt.text(f"Installing {p}...")
-        res = subprocess.run([sys.executable, "-m", "pip", "install", p], capture_output=True, text=True)
-        if res.returncode != 0:
-            results.append(f"Failed {p}: {res.stderr}")
-            success = False
-        else:
-            results.append(f"Installed {p}")
-    bar.progress(1.0)
-    sidebar_txt.empty()
-    return success, "\n".join(results)
-
-@st.cache_resource(ttl=3600)
-def init_ai_models_direct():
-    token = get_secret("github_token_api")
-    if not token:
-        st.error("API token not configured")
-        return None
-    try:
-        from azure.ai.inference import ChatCompletionsClient
-        from azure.ai.inference.models import UserMessage
-        from azure.core.credentials import AzureKeyCredential
-        client = ChatCompletionsClient(
-            endpoint="https://models.inference.ai.azure.com",
-            credential=AzureKeyCredential(token)
-        )
-        return {"client": client, "model_name": "gpt-4o", "last_loaded": datetime.now().isoformat()}
-    except ImportError as e:
-        st.error("Azure AI SDK not installed")
-        logger.error(str(e))
-        return None
+def extract_scene_class_name(code):
+    m = re.findall(r"class\s+(\w+)\s*\([^)]*Scene", code)
+    return m[0] if m else "MyScene"
 
 def generate_manim_preview(code):
-    objects = []
-    if "Circle" in code: objects.append("⭕")
-    if "Square" in code: objects.append("🔲")
-    if "MathTex" in code or "Tex" in code: objects.append("📊")
-    if "Text" in code: objects.append("📝")
-    if "Axes" in code: objects.append("📈")
-    icons = "".join(objects) or "🎬"
+    icons = []
+    if "Circle" in code: icons.append("⭕")
+    if "Square" in code: icons.append("🔲")
+    if "MathTex" in code or "Tex" in code: icons.append("📊")
+    if "Text" in code: icons.append("📝")
+    if "Axes" in code: icons.append("📈")
+    preview_icons = "".join(icons) or "🎬"
     return f"""
-    <div style="background:#000;color:#fff;padding:1rem;border-radius:10px;text-align:center;">
-      <h3>Animation Preview</h3>
-      <div style="font-size:2rem;">{icons}</div>
-      <p>Full rendering required for accurate preview</p>
+    <div style="background:#000;color:#fff;padding:1rem;border-radius:8px;text-align:center;">
+      <h3>Preview</h3>
+      <div style="font-size:2rem;">{preview_icons}</div>
+      <p>Full render for accurate output</p>
     </div>
     """
 
-def extract_scene_class_name(code):
-    m = re.findall(r'class\s+(\w+)\s*\([^)]*Scene', code)
-    return m[0] if m else "MyScene"
-
-def mp4_to_gif(mp4_path, gif_path, fps=15):
-    cmd = [
-        "ffmpeg", "-i", mp4_path,
-        "-vf", f"fps={fps},scale=640:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse",
-        "-loop", "0", gif_path
-    ]
-    res = subprocess.run(cmd, capture_output=True, text=True)
-    return gif_path if res.returncode == 0 else None
-
-def generate_manim_video(code, fmt, quality, speed=1.0, audio_path=None):
+def generate_manim_video(code, fmt, quality, fps, audio_path=None):
     temp_dir = tempfile.mkdtemp(prefix="manim_")
     scene = extract_scene_class_name(code)
     scene_file = os.path.join(temp_dir, "scene.py")
@@ -202,345 +145,240 @@ def generate_manim_video(code, fmt, quality, speed=1.0, audio_path=None):
         f.write(code)
     qflags = {"480p":"-ql","720p":"-qm","1080p":"-qh","4K":"-qk","8K":"-qp"}
     qf = qflags.get(quality, "-qm")
-    cmd = ["manim", scene_file, scene, qf, f"--format={fmt}"]
-    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
-    output, out_path, mp4_path = [], None, None
+    cmd = ["manim", scene_file, scene, qf, f"--format={fmt}", f"--fps={fps}"]
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    progress = st.progress(0)
     log = st.empty()
-    for line in proc.stdout:
-        output.append(line)
-        log.code("".join(output[-10:]))
-        if "File ready at" in line:
-            m = re.search(r'["\'](.+?\.(?:mp4|gif|webm|svg))["\']', line)
-            if m:
-                out_path = m.group(1)
-                if out_path.endswith(".mp4"):
-                    mp4_path = out_path
-    proc.wait()
-    time.sleep(1)
-    if fmt=="gif" and (not out_path or not os.path.exists(out_path)) and mp4_path:
-        gif = os.path.join(temp_dir, "converted.gif")
-        conv = mp4_to_gif(mp4_path, gif)
-        if conv and os.path.exists(conv):
-            out_path = conv
+    total = None
+    for line in process.stdout:
+        log.code(line)
+        m = re.search(r"(\d+)\s*/\s*(\d+)", line)
+        if m:
+            cur, tot = map(int, m.groups())
+            total = tot
+            progress.progress(min(0.99, cur / tot))
+        else:
+            p = re.search(r"(\d+)%", line)
+            if p:
+                progress.progress(min(0.99, int(p.group(1)) / 100))
+    process.wait()
+    progress.progress(1.0)
+    # locate output
+    out_path = None
+    for root, _, files in os.walk(temp_dir):
+        for f in files:
+            if f.endswith(f".{fmt}") or (fmt=="png_sequence" and f.endswith(".zip")):
+                out_path = os.path.join(root, f)
+                break
     data = None
     if out_path and os.path.exists(out_path):
         with open(out_path, "rb") as f:
             data = f.read()
     shutil.rmtree(temp_dir)
     if data:
-        size_mb = len(data)/(1024*1024)
-        return data, f"✅ Generated ({size_mb:.1f} MB)"
+        size_mb = len(data) / (1024 * 1024)
+        return data, f"✅ Done ({size_mb:.1f} MB)"
     else:
-        return None, "❌ No output generated. See logs."
-
-def detect_input_calls(code):
-    calls=[]
-    for i,line in enumerate(code.split("\n"),1):
-        if "input(" in line and not line.strip().startswith("#"):
-            m=re.search(r'input\(["\'](.+?)["\']\)', line)
-            prompt=m.group(1) if m else f"Input at line {i}"
-            calls.append({"line":i,"prompt":prompt})
-    return calls
-
-def run_python_script(code, inputs=None, timeout=60):
-    res={"stdout":"","stderr":"","exception":None,"plots":[],"dataframes":[],"execution_time":0}
-    mod=""
-    if inputs:
-        mod=f"""
-__INPUTS={inputs}
-__IDX=0
-def input(prompt=''):
-    global __IDX
-    print(prompt,end='')
-    if __IDX<len(__INPUTS):
-        val=__INPUTS[__IDX]; __IDX+=1
-        print(val)
-        return val
-    print()
-    return ''
-"""
-    full_code=mod+code
-    with tempfile.TemporaryDirectory() as td:
-        path=os.path.join(td,"script.py")
-        with open(path,"w") as f: f.write(full_code)
-        outf, errf = os.path.join(td,"out.txt"), os.path.join(td,"err.txt")
-        start=time.time()
-        try:
-            with open(outf,"w") as o, open(errf,"w") as e:
-                proc=subprocess.Popen([sys.executable, path], stdout=o, stderr=e, cwd=td)
-                proc.wait(timeout=timeout)
-        except subprocess.TimeoutExpired:
-            proc.kill()
-            res["stderr"]+="\nTimed out"
-            res["exception"]="Timeout"
-        res["execution_time"]=time.time()-start
-        res["stdout"]=open(outf).read()
-        res["stderr"]+=open(errf).read()
-    return res
-
-def display_python_script_results(r):
-    st.info(f"Completed in {r['execution_time']:.2f}s")
-    if r["exception"]:
-        st.error(f"Exception: {r['exception']}")
-    if r["stderr"]:
-        st.error("Errors:")
-        st.code(r["stderr"], language="bash")
-    if r["plots"]:
-        st.markdown("### Plots")
-        cols=st.columns(min(3,len(r["plots"])))
-        for i,p in enumerate(r["plots"]):
-            cols[i%len(cols)].image(p,use_column_width=True)
-    if r["dataframes"]:
-        st.markdown("### DataFrames")
-        for df in r["dataframes"]:
-            with st.expander(f"{df['name']} {df['shape']}"):
-                st.markdown(df["preview_html"], unsafe_allow_html=True)
-    if r["stdout"]:
-        st.markdown("### Output")
-        st.code(r["stdout"], language="bash")
+        return None, "❌ No output generated"
 
 def main():
-    if 'init' not in st.session_state:
+    st.set_page_config("Manim Studio", "🎬", layout="wide")
+    if "inited" not in st.session_state:
         st.session_state.update({
-            'init':True, 'video_data':None, 'status':None, 'ai_models':None,
-            'generated_code':"", 'code':"", 'temp_code':"", 'editor_key':str(uuid.uuid4()),
-            'packages_checked':False, 'audio_path':None, 'image_paths':[],
-            'custom_library_result':"", 'python_script':"", 'python_result':None,
-            'active_tab':0, 'settings':{"quality":"720p","format_type":"mp4","animation_speed":"Normal"},
-            'password_entered':False, 'custom_model':"gpt-4o", 'pending_tab_switch':None
+            "inited": True,
+            "code": "",
+            "video_data": None,
+            "status": "",
+            "settings": {"quality":"720p","fps":30,"format":"mp4"},
+            "pwd_ok": False,
+            "custom_model": "gpt-4o",
+            "generated_code": ""
         })
-    st.set_page_config(page_title="Manim Animation Studio", page_icon="🎬", layout="wide")
 
-    if not st.session_state.packages_checked:
+    if not st.session_state.get("packages_checked", False):
         if ensure_packages():
-            st.session_state.packages_checked=True
+            st.session_state.packages_checked = True
         else:
-            st.error("Package installation failed")
             return
 
-    tab_names=[
-        "✨ Editor","🤖 AI Assistant","🎨 Assets",
-        "🎞️ Timeline","🎓 Educational Export","🐍 Python Runner"
-    ]
-    tabs = st.tabs(tab_names)
+    tabs = st.tabs([
+        "✨ Editor", "🤖 AI Assistant", "🎨 Assets",
+        "🎞️ Timeline", "🎓 Export", "🐍 Python Runner"
+    ])
+
+    # Sidebar settings
+    with st.sidebar:
+        st.markdown("## ⚙️ Render Settings")
+        q = st.selectbox("Quality", QUALITY_PRESETS, index=QUALITY_PRESETS.index(st.session_state["settings"]["quality"]))
+        f = st.selectbox("FPS", FPS_OPTIONS, index=FPS_OPTIONS.index(st.session_state["settings"]["fps"]))
+        fmt_disp = st.selectbox("Format", list(OUTPUT_FORMATS.keys()),
+                                 index=list(OUTPUT_FORMATS.values()).index(st.session_state["settings"]["format"]))
+        st.session_state["settings"] = {"quality": q, "fps": f, "format": OUTPUT_FORMATS[fmt_disp]}
 
-    # Editor
+    # Editor tab
     with tabs[0]:
-        col1,col2 = st.columns([3,2])
-        with col1:
+        c1, c2 = st.columns([3,2])
+        with c1:
             st.markdown("### 📝 Animation Editor")
-            mode = st.radio("Code Input", ["Type Code","Upload File"], key="editor_mode")
-            if mode=="Upload File":
-                up = st.file_uploader("Upload .py", type=["py"])
-                if up:
-                    txt=up.getvalue().decode()
-                    if txt.strip():
-                        st.session_state.code=txt
-                        st.session_state.temp_code=txt
             if ACE_EDITOR_AVAILABLE:
-                st.session_state.temp_code = st_ace(
-                    value=st.session_state.code, language="python",
-                    theme="monokai", min_lines=20,
-                    key=f"ace_{st.session_state.editor_key}"
+                st.session_state.code = st_ace(
+                    value=st.session_state.code,
+                    language="python",
+                    theme="monokai",
+                    min_lines=20,
+                    key="ace"
                 )
             else:
-                st.session_state.temp_code = st.text_area(
-                    "Code", st.session_state.code, height=400,
-                    key=f"ta_{st.session_state.editor_key}"
+                st.session_state.code = st.text_area(
+                    "Manim Code",
+                    value=st.session_state.code,
+                    height=400
                 )
-            if st.session_state.temp_code!=st.session_state.code:
-                st.session_state.code=st.session_state.temp_code
             if st.button("🚀 Generate Animation"):
-                if not st.session_state.code:
+                if not st.session_state.code.strip():
                     st.error("Enter code first")
                 else:
                     data, msg = generate_manim_video(
                         st.session_state.code,
-                        st.session_state.settings["format_type"],
-                        st.session_state.settings["quality"],
-                        {"Slow":0.5,"Normal":1.0,"Fast":2.0,"Very Fast":3.0}[st.session_state.settings["animation_speed"]],
-                        st.session_state.audio_path
+                        st.session_state["settings"]["format"],
+                        st.session_state["settings"]["quality"],
+                        st.session_state["settings"]["fps"]
                     )
-                    st.session_state.video_data=data
-                    st.session_state.status=msg
-        with col2:
+                    st.session_state.video_data = data
+                    st.session_state.status = msg
+        with c2:
             if st.session_state.code:
-                components.html(
-                    generate_manim_preview(st.session_state.code),
-                    height=250
-                )
+                components.html(generate_manim_preview(st.session_state.code), height=250)
             if st.session_state.video_data:
-                fmt=st.session_state.settings["format_type"]
-                if fmt=="png_sequence":
-                    st.download_button(
-                        "⬇️ Download PNG ZIP", data=st.session_state.video_data,
-                        file_name=f"manim_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip",
-                        mime="application/zip"
-                    )
-                elif fmt=="svg":
+                fmt = st.session_state["settings"]["format"]
+                if fmt == "png_sequence":
+                    st.download_button("⬇️ Download ZIP", st.session_state.video_data, "animation.zip", "application/zip")
+                elif fmt == "svg":
                     try:
-                        svg=st.session_state.video_data.decode('utf-8')
+                        svg = st.session_state.video_data.decode("utf-8")
                         components.html(svg, height=400)
                     except:
                         st.error("Cannot display SVG")
-                    st.download_button(
-                        "⬇️ Download SVG", data=st.session_state.video_data,
-                        file_name="animation.svg", mime="image/svg+xml"
-                    )
+                    st.download_button("⬇️ Download SVG", st.session_state.video_data, "animation.svg", "image/svg+xml")
                 else:
                     st.video(st.session_state.video_data, format=fmt)
-                    st.download_button(
-                        f"⬇️ Download {fmt.upper()}", st.session_state.video_data,
-                        file_name=f"animation.{fmt}", mime=f"video/{fmt}" if fmt!="gif" else "image/gif"
-                    )
+                    st.download_button(f"⬇️ Download {fmt.upper()}", st.session_state.video_data,
+                                       f"animation.{fmt}", f"video/{fmt}")
             if st.session_state.status:
-                if "❌" in st.session_state.status:
+                if st.session_state.status.startswith("❌"):
                     st.error(st.session_state.status)
                 else:
                     st.success(st.session_state.status)
 
-    # AI Assistant
+    # AI Assistant tab
     with tabs[1]:
-        st.markdown("### 🤖 AI Animation Assistant")
-        if check_password():
-            client_data = init_ai_models_direct()
-            if client_data:
-                if st.button("Test API Connection"):
-                    from azure.ai.inference.models import UserMessage
-                    params,_=prepare_api_params([UserMessage("Hello")], client_data["model_name"])
-                    resp=client_data["client"].complete(**params)
-                    if resp.choices:
-                        st.success("✅ Connection successful!")
-                        st.session_state.ai_models=client_data
-                    else:
-                        st.error("❌ No response")
-                if st.session_state.ai_models:
-                    st.info(f"Using model {st.session_state.ai_models['model_name']}")
-                    prompt = st.text_area("Describe animation or paste partial code", height=150)
-                    if st.button("Generate Animation Code"):
-                        if prompt.strip():
-                            from azure.ai.inference.models import UserMessage
-                            params,_=prepare_api_params(
-                                [UserMessage(f"Write a complete Manim scene for:\n{prompt}")],
-                                st.session_state.ai_models["model_name"]
-                            )
-                            resp=st.session_state.ai_models["client"].complete(**params)
-                            if resp.choices:
-                                code = resp.choices[0].message.content
-                                if "```python" in code:
-                                    code=code.split("```python")[1].split("```")[0]
-                                st.session_state.generated_code=code
-                            else:
-                                st.error("No code generated")
-                        else:
-                            st.warning("Enter prompt first")
-                    if st.session_state.generated_code:
-                        st.code(st.session_state.generated_code, language="python")
-                        if st.button("Use This Code"):
-                            st.session_state.code=st.session_state.generated_code
-                            st.session_state.temp_code=st.session_state.generated_code
-                            st.session_state.pending_tab_switch=0
-                            st.rerun()
-        else:
-            st.info("Enter password to access AI")
+        st.markdown("### 🤖 AI Assistant")
+        if not check_password():
+            return
+        model = st.selectbox("Select AI Model", list(MODEL_CONFIGS.keys()),
+                             index=list(MODEL_CONFIGS.keys()).index(st.session_state.custom_model))
+        st.session_state.custom_model = model
+        token = get_secret("github_token_api")
+        if st.button("Test Connection"):
+            from azure.ai.inference import ChatCompletionsClient
+            from azure.core.credentials import AzureKeyCredential
+            client = ChatCompletionsClient(
+                endpoint="https://models.inference.ai.azure.com",
+                credential=AzureKeyCredential(token)
+            )
+            from azure.ai.inference.models import UserMessage
+            resp = client.complete(**{"messages":[UserMessage("Hello")], "model":model, "max_tokens":1000})
+            if resp.choices:
+                st.success("✅ Connected")
+                st.session_state.ai_client = client
+            else:
+                st.error("❌ No response")
+        if "ai_client" in st.session_state:
+            prompt = st.text_area("Describe animation or provide code")
+            if st.button("Generate Code"):
+                from azure.ai.inference.models import UserMessage
+                msgs = [UserMessage(f"Write a complete Manim scene:\n{prompt}")]
+                resp = st.session_state.ai_client.complete(**{"messages":msgs, "model":model, "max_tokens":8000})
+                if resp.choices:
+                    code = resp.choices[0].message.content
+                    if "```python" in code:
+                        code = code.split("```python")[1].split("```")[0]
+                    st.session_state.generated_code = code
+                else:
+                    st.error("No code returned")
+            if st.session_state.generated_code:
+                st.code(st.session_state.generated_code, language="python")
+                if st.button("Use This Code"):
+                    st.session_state.code = st.session_state.generated_code
+                    st.experimental_rerun()
 
-    # Assets
+    # Assets tab
     with tabs[2]:
-        st.markdown("### 🎨 Asset Management")
-        c1,c2 = st.columns(2)
-        with c1:
-            imgs = st.file_uploader(
-                "Upload Images", type=["png","jpg","jpeg","svg"],
-                accept_multiple_files=True
-            )
-            if imgs:
-                idir = os.path.join(os.getcwd(),"manim_assets","images")
-                os.makedirs(idir, exist_ok=True)
-                for up in imgs:
-                    ext=up.name.split(".")[-1]
-                    fname=f"img_{int(time.time())}_{uuid.uuid4().hex[:6]}.{ext}"
-                    path=os.path.join(idir,fname)
-                    with open(path,"wb") as f: f.write(up.getvalue())
-                    st.session_state.image_paths.append({"name":up.name,"path":path})
-            for info in st.session_state.image_paths:
-                img=Image.open(info["path"])
-                st.image(img, caption=info["name"], width=100)
-                if st.button(f"Use {info['name']}"):
-                    snippet=f"""
-# Image asset
-image = ImageMobject(r"{info['path']}")
-image.scale(2)
+        st.markdown("### 🎨 Assets")
+        imgs = st.file_uploader("Upload Images", type=["png","jpg","jpeg","svg"], accept_multiple_files=True)
+        if imgs:
+            adir = os.path.join(os.getcwd(),"manim_assets","images")
+            os.makedirs(adir, exist_ok=True)
+            for up in imgs:
+                ext = up.name.split(".")[-1]
+                fn = f"img_{int(time.time())}_{uuid.uuid4().hex[:6]}.{ext}"
+                path = os.path.join(adir, fn)
+                with open(path,"wb") as f: f.write(up.getvalue())
+                st.image(path, caption=up.name, width=150)
+                if st.button(f"Use {up.name}"):
+                    snippet = f"""image = ImageMobject(r"{path}")
 self.play(FadeIn(image))
 self.wait(1)
 """
-                    st.session_state.code+=snippet
-                    st.session_state.temp_code=st.session_state.code
-                    st.success(f"Added {info['name']}")
-                    st.session_state.pending_tab_switch=0
-                    st.rerun()
-        with c2:
-            aud = st.file_uploader("Upload Audio", type=["mp3","wav","ogg"])
-            if aud:
-                adir = os.path.join(os.getcwd(),"manim_assets","audio")
-                os.makedirs(adir, exist_ok=True)
-                ext=aud.name.split(".")[-1]
-                aname=f"audio_{int(time.time())}.{ext}"
-                ap=os.path.join(adir,aname)
-                with open(ap,"wb") as f: f.write(aud.getvalue())
-                st.session_state.audio_path=ap
-                st.audio(aud)
-                st.success("Audio uploaded")
+                    st.session_state.code += "\n" + snippet
+                    st.experimental_rerun()
+        aud = st.file_uploader("Upload Audio", type=["mp3","wav","ogg"])
+        if aud:
+            adir = os.path.join(os.getcwd(),"manim_assets","audio")
+            os.makedirs(adir, exist_ok=True)
+            fn = f"audio_{int(time.time())}.{aud.name.split('.')[-1]}"
+            path = os.path.join(adir, fn)
+            with open(path,"wb") as f: f.write(aud.getvalue())
+            st.audio(aud)
+            st.success("Audio uploaded")
 
-    # Timeline
+    # Timeline tab
     with tabs[3]:
         st.markdown("### 🎞️ Timeline Editor")
-        st.info("Use code editor to adjust timing of self.play and self.wait calls.")
+        st.info("Adjust code manually – timeline UI coming soon.")
 
-    # Educational Export
+    # Export tab
     with tabs[4]:
-        st.markdown("### 🎓 Educational Export")
-        if not st.session_state.video_data:
-            st.warning("Generate animation first")
-        else:
-            title=st.text_input("Title","Manim Animation")
-            expl=st.text_area("Explanation (use ## to separate steps)",height=150)
-            fmt=st.selectbox("Format",["PowerPoint","HTML","PDF Sequence"])
-            if st.button("Export"):
-                st.success(f"{fmt} export not implemented yet")
+        st.markdown("### 🎓 Export")
+        st.warning("Export features coming soon.")
 
-    # Python Runner
+    # Python Runner tab
     with tabs[5]:
-        st.markdown("### 🐍 Python Script Runner")
-        examples={"Select...":"","Sine Plot":"""import matplotlib.pyplot as plt
-import numpy as np
-x=np.linspace(0,10,100)
-y=np.sin(x)
-plt.plot(x,y)
-print("Done")"""}
-        sel=st.selectbox("Example",list(examples.keys()))
-        code = examples.get(sel, st.session_state.python_script)
-        if ACE_EDITOR_AVAILABLE:
-            code=st_ace(value=code, language="python", theme="monokai", min_lines=15, key="pyace")
-        else:
-            code=st.text_area("Code", code, height=300, key="pyta")
-        st.session_state.python_script=code
-        inputs=detect_input_calls(code)
-        vals=[]
-        if inputs:
-            st.info(f"{len(inputs)} input() calls detected")
-            for i,c in enumerate(inputs):
-                vals.append(st.text_input(f"{c['prompt']} (line {c['line']})", key=f"in{i}"))
-        timeout=st.slider("Timeout",5,300,30)
-        if st.button("▶️ Run"):
-            res=run_python_script(code, inputs=vals, timeout=timeout)
-            st.session_state.python_result=res
-        if st.session_state.python_result:
-            display_python_script_results(st.session_state.python_result)
-
-    # Handle pending tab switch
-    if st.session_state.pending_tab_switch is not None:
-        st.session_state.active_tab = st.session_state.pending_tab_switch
-        st.session_state.pending_tab_switch = None
+        st.markdown("### 🐍 Python Runner")
+        code = st.text_area("Script", height=300, key="py_code")
+        inputs = []
+        for i, line in enumerate(code.split("\n"), 1):
+            if "input(" in line:
+                prompt = re.search(r'input\(["\'](.+?)["\']\)', line)
+                label = prompt.group(1) if prompt else f"Input line {i}"
+                inputs.append(st.text_input(label, key=f"in_{i}"))
+        timeout = st.slider("Timeout (s)", 5, 300, 30)
+        if st.button("Run"):
+            temp = tempfile.NamedTemporaryFile(delete=False, suffix=".py")
+            temp.write(code.encode())
+            temp.flush()
+            proc = subprocess.Popen([sys.executable, temp.name],
+                                    stdin=subprocess.PIPE,
+                                    stdout=subprocess.PIPE,
+                                    stderr=subprocess.PIPE,
+                                    text=True)
+            # feed inputs
+            out, err = proc.communicate(input="\n".join(inputs), timeout=timeout)
+            if err:
+                st.error(err)
+            if out:
+                st.code(out)
+            temp.close()
 
 if __name__ == "__main__":
     main()