app.py

import streamlit as st
import tempfile
import os
import logging
from pathlib import Path
from PIL import Image
import io
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
from transformers import pipeline
import torch
import re
import shutil
import time
from datetime import datetime, timedelta
import streamlit.components.v1 as components
import uuid
import platform
import pandas as pd
import plotly.express as px
import markdown
import zipfile
import contextlib
import threading
import traceback
from io import StringIO, BytesIO

# Set up enhanced logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Model configuration mapping for different API requirements and limits
MODEL_CONFIGS = {
    "DeepSeek-V3-0324": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek", "warning": None},
    "DeepSeek-R1": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek", "warning": None},
    "Llama-4-Scout-17B-16E-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta", "warning": None},
    "Llama-4-Maverick-17B-128E-Instruct-FP8": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta", "warning": None},
    "gpt-4o-mini": {"max_tokens": 15000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "gpt-4o": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "gpt-4.1": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "gpt-4.1-mini": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "gpt-4.1-nano": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "o3-mini": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
    "o1": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
    "o1-mini": {"max_completion_tokens": 66000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
    "o1-preview": {"max_tokens": 33000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
    "Phi-4-multimodal-instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Microsoft", "warning": None},
    "Mistral-large-2407": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral", "warning": None},
    "Codestral-2501": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral", "warning": None},
    # Default configuration for other models
    "default": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Other", "warning": None}
}

# Try to import Streamlit 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 standard text editor")

def prepare_api_params(messages, model_name):
    """Create appropriate API parameters based on model configuration"""
    config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
    api_params = {"messages": messages, "model": model_name}
    token_param = config["param_name"]
    token_value = config[token_param]
    api_params[token_param] = token_value
    return api_params, config

def get_secret(key):
    """Retrieve a secret from environment or Streamlit secrets."""
    if hasattr(st, "secrets") and key in st.secrets:
        return st.secrets[key]
    return os.environ.get(key)

def check_password():
    correct_password = get_secret("password")
    if not correct_password:
        st.error("Admin password not configured in secrets or env var 'password'")
        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_password:
                st.session_state.password_entered = True
                return True
            else:
                st.error("Incorrect password")
                return False
        return False
    return True

def ensure_packages():
    required_packages = {
        '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',
        'markdown': '3.4.3',
        'fpdf': '1.7.2',
        'matplotlib': '3.5.0',
        'seaborn': '0.11.2',
        'scipy': '1.7.3',
        'huggingface_hub': '0.16.0',
    }
    missing = {}
    for pkg, ver in required_packages.items():
        try:
            __import__(pkg if pkg != 'Pillow' else 'PIL')
        except ImportError:
            missing[pkg] = ver
    if not missing:
        return True
    progress = st.progress(0)
    status = st.empty()
    for i, (pkg, ver) in enumerate(missing.items()):
        status.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}: {res.stderr}")
            return False
        progress.progress((i + 1) / len(missing))
    return True

@st.cache_resource(ttl=3600)
def init_ai_models_direct():
    try:
        token = get_secret("github_token_api")
        if not token:
            st.error("GitHub token not found in secrets or env var 'github_token_api'")
            return None
        from azure.ai.inference import ChatCompletionsClient
        from azure.ai.inference.models import SystemMessage, UserMessage
        from azure.core.credentials import AzureKeyCredential
        endpoint = "https://models.inference.ai.azure.com"
        model_name = "gpt-4o"
        client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(token))
        return {
            "client": client,
            "model_name": model_name,
            "endpoint": endpoint,
            "last_loaded": datetime.now().isoformat(),
            "category": MODEL_CONFIGS[model_name]["category"],
            "api_version": MODEL_CONFIGS[model_name].get("api_version")
        }
    except Exception as e:
        st.error(f"Error initializing AI model: {e}")
        logger.error(str(e))
        return None

def suggest_code_completion(code_snippet, models):
    if not models:
        st.error("AI models not initialized")
        return None
    try:
        prompt = f"""Write a complete Manim animation scene based on this code or idea:
{code_snippet}

The code should be a complete, working Manim animation that includes:
- Proper Scene class definition
- Constructor with animations
- Proper use of self.play() for animations
- Proper wait times between animations

Here's the complete Manim code:
"""
        from openai import OpenAI
        token = get_secret("github_token_api")
        client = OpenAI(base_url="https://models.github.ai/inference", api_key=token)
        messages = [{"role": "system", "content": "You are an expert in Manim animations."},
                    {"role": "user", "content": prompt}]
        config = MODEL_CONFIGS.get(models["model_name"], MODEL_CONFIGS["default"])
        params = {"messages": messages, "model": models["model_name"], config["param_name"]: config[config["param_name"]]}
        response = client.chat.completions.create(**params)
        content = response.choices[0].message.content
        if "```python" in content:
            content = content.split("```python")[1].split("```")[0]
        elif "```" in content:
            content = content.split("```")[1].split("```")[0]
        if "Scene" not in content:
            content = f"from manim import *\n\nclass MyScene(Scene):\n    def construct(self):\n        {content}"
        return content
    except Exception as e:
        st.error(f"Error generating code: {e}")
        logger.error(traceback.format_exc())
        return None

QUALITY_PRESETS = {
    "480p": {"resolution": "480p", "fps": "30"},
    "720p": {"resolution": "720p", "fps": "30"},
    "1080p": {"resolution": "1080p", "fps": "60"},
    "4K": {"resolution": "2160p", "fps": "60"},
    "8K": {"resolution": "4320p", "fps": "60"}
}

ANIMATION_SPEEDS = {
    "Slow": 0.5,
    "Normal": 1.0,
    "Fast": 2.0,
    "Very Fast": 3.0
}

EXPORT_FORMATS = {
    "MP4 Video": "mp4",
    "GIF Animation": "gif",
    "WebM Video": "webm",
    "PNG Image Sequence": "png_sequence",
    "SVG Image": "svg"
}

def highlight_code(code):
    formatter = HtmlFormatter(style='monokai')
    highlighted = highlight(code, PythonLexer(), formatter)
    return highlighted, formatter.get_style_defs()

def generate_manim_preview(python_code):
    scene_objects = []
    if "Circle" in python_code: scene_objects.append("circle")
    if "Square" in python_code: scene_objects.append("square")
    if "MathTex" in python_code or "Tex" in python_code: scene_objects.append("equation")
    if "Text" in python_code: scene_objects.append("text")
    if "Axes" in python_code: scene_objects.append("graph")
    if "ThreeDScene" in python_code or "ThreeDAxes" in python_code: scene_objects.append("3D scene")
    if "Sphere" in python_code: scene_objects.append("sphere")
    if "Cube" in python_code: scene_objects.append("cube")
    icons = {"circle":"⭕","square":"🔲","equation":"📊","text":"📝","graph":"📈","3D scene":"🧊","sphere":"🌐","cube":"🧊"}
    icon_html = "".join(f'<span style="font-size:2rem; margin:0.3rem;">{icons[o]}</span>' for o in scene_objects)
    preview_html = f"""
    <div style="background-color:#000; width:100%; height:220px; border-radius:10px; display:flex; flex-direction:column; align-items:center; justify-content:center; color:white; text-align:center;">
        <h3>Animation Preview</h3>
        <div>{icon_html if icon_html else '<span style="font-size:2rem;">🎬</span>'}</div>
        <p>Scene contains: {', '.join(scene_objects) if scene_objects else 'No detected objects'}</p>
        <p style="font-size:0.8rem; opacity:0.7;">Full rendering required for accurate preview</p>
    </div>
    """
    return preview_html

def render_latex_preview(latex):
    if not latex:
        return """
        <div style="background:#f8f9fa; width:100%; height:100px; border-radius:5px; display:flex; align-items:center; justify-content:center; color:#6c757d;">
            Enter LaTeX formula to see preview
        </div>
        """
    return f"""
    <div style="background:#202124; width:100%; padding:20px; border-radius:5px; color:white; text-align:center;">
        <script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
        <script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
        <div><h3>LaTeX Preview</h3><div id="math-preview">$$ {latex} $$</div><p style="font-size:0.8rem; opacity:0.7;">Use MathTex(r"{latex}") in Manim</p></div>
    </div>
    """

def extract_scene_class_name(python_code):
    match = re.search(r'class\s+(\w+)\s*\([^)]*Scene[^)]*\)', python_code)
    return match.group(1) if match else "MyScene"

def prepare_audio_for_manim(audio_file, target_dir):
    audio_dir = os.path.join(target_dir, "audio")
    os.makedirs(audio_dir, exist_ok=True)
    filename = f"audio_{int(time.time())}.mp3"
    path = os.path.join(audio_dir, filename)
    with open(path, "wb") as f: f.write(audio_file.getvalue())
    return path

def mp4_to_gif(mp4, out, fps=15):
    cmd = ["ffmpeg","-i",mp4,"-vf",f"fps={fps},scale=640:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse","-loop","0",out]
    res = subprocess.run(cmd,capture_output=True,text=True)
    return out if res.returncode==0 else None

def generate_manim_video(code, fmt, quality, speed, audio_path=None):
    temp_dir = tempfile.mkdtemp(prefix="manim_render_")
    try:
        scene = extract_scene_class_name(code)
        if audio_path and "with_sound" not in code:
            code = "from manim.scene.scene_file_writer import SceneFileWriter\n" + code
            pat = re.search(f"class {scene}\\(.*?\\):", code)
            if pat:
                decor = f"@with_sound(\"{audio_path}\")\n"
                code = code[:pat.start()] + decor + code[pat.start():]
        path_py = os.path.join(temp_dir, "scene.py")
        with open(path_py, "w", encoding="utf-8") as f: f.write(code)
        qmap = {"480p":"-ql","720p":"-qm","1080p":"-qh","4K":"-qk","8K":"-qp"}
        qflag = qmap.get(quality,"-qm")
        if fmt=="png_sequence":
            farg="--format=png"; extra=["--save_pngs"]
        elif fmt=="svg":
            farg="--format=svg"; extra=[]
        else:
            farg=f"--format={fmt}"; extra=[]
        cmd = ["manim", path_py, scene, qflag, farg] + extra
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        output=[]
        out_path=None; mp4_path=None
        while True:
            line = proc.stdout.readline()
            if not line and proc.poll() is not None: break
            output.append(line)
            if "%" in line:
                try:
                    p=float(line.split("%")[0].strip().split()[-1]); 
                except: pass
            if "File ready at" in line:
                chunk = line.split("File ready at")[-1].strip()
                m=re.search(r'([\'"]?)(.*?\.(mp4|gif|webm|svg))\1',chunk)
                if m: 
                    out_path=m.group(2)
                    if out_path.endswith(".mp4"): mp4_path=out_path
        proc.wait()
        time.sleep(2)
        data=None
        if fmt=="gif" and (not out_path or not os.path.exists(out_path)) and mp4_path:
            gif=os.path.join(temp_dir,f"{scene}_converted.gif")
            if mp4_to_gif(mp4_path,gif): out_path=gif
        if fmt=="png_sequence":
            dirs=[os.path.join(temp_dir,"media","images",scene,"Animations")]
            pngs=[]
            for d in dirs:
                if os.path.isdir(d):
                    pngs+= [os.path.join(d,f) for f in os.listdir(d) if f.endswith(".png")]
            if pngs:
                zipf=os.path.join(temp_dir,f"{scene}_pngs.zip")
                with zipfile.ZipFile(zipf,"w") as z:
                    for p in pngs: z.write(p,os.path.basename(p))
                data=open(zipf,"rb").read()
        elif out_path and os.path.exists(out_path):
            data=open(out_path,"rb").read()
        else:
            # fallback search
            files=[]
            for root,_,fs in os.walk(temp_dir):
                for f in fs:
                    if f.endswith(f".{fmt}") and "partial" not in f:
                        files.append(os.path.join(root,f))
            if files:
                latest=max(files,key=os.path.getctime)
                data=open(latest,"rb").read()
                if fmt=="gif" and latest.endswith(".mp4"):
                    gif=os.path.join(temp_dir,f"{scene}_converted.gif")
                    if mp4_to_gif(latest,gif): data=open(gif,"rb").read()
        if data:
            size=len(data)/(1024*1024)
            return data, f"✅ Animation generated successfully! ({size:.1f} MB)"
        else:
            return None, "❌ Error: No output files generated.\n" + "".join(output)[:500]
    except Exception as e:
        logger.error(traceback.format_exc())
        return None, f"❌ Error: {e}"
    finally:
        try: shutil.rmtree(temp_dir)
        except: pass

def detect_input_calls(code):
    calls=[]
    for i,line in enumerate(code.splitlines(),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 for line {i}"
            calls.append({"line":i,"prompt":prompt})
    return calls

def run_python_script(code, inputs=None, timeout=60):
    result={"stdout":"","stderr":"","exception":None,"plots":[],"dataframes":[],"execution_time":0}
    if inputs:
        inject = f"""
__INPUT_VALUES={inputs}
__INPUT_INDEX=0
def input(prompt=''):
    global __INPUT_INDEX
    print(prompt,end='')
    if __INPUT_INDEX<len(__INPUT_VALUES):
        v=__INPUT_VALUES[__INPUT_INDEX]; __INPUT_INDEX+=1
        print(v); return v
    print(); return ''
"""
        code = inject + code
    with tempfile.TemporaryDirectory() as td:
        plot_dir=os.path.join(td,'plots'); os.makedirs(plot_dir,exist_ok=True)
        stdout_f=os.path.join(td,'stdout.txt')
        stderr_f=os.path.join(td,'stderr.txt')
        if 'plt' in code or 'matplotlib' in code:
            if 'import matplotlib.pyplot as plt' not in code:
                code="import matplotlib.pyplot as plt\n"+code
            save_plots=f"""
import matplotlib.pyplot as plt,os
for i,num in enumerate(plt.get_fignums()):
    plt.figure(num).savefig(os.path.join(r'{plot_dir}','plot_{{i}}.png'))
"""
            code+=save_plots
        if 'pd.' in code or 'import pandas' in code:
            if 'import pandas as pd' not in code:
                code="import pandas as pd\n"+code
            dfcap=f"""
import pandas as pd, json,os
for name,val in globals().items():
    if isinstance(val,pd.DataFrame):
        info={{"name":name,"shape":val.shape,"columns":list(val.columns),"preview":val.head().to_html()}}
        open(os.path.join(r'{td}',f'df_{{name}}.json'),'w').write(json.dumps(info))
"""
            code+=dfcap
        script=os.path.join(td,'script.py')
        open(script,'w').write(code)
        start=time.time()
        try:
            with open(stdout_f,'w') as so, open(stderr_f,'w') as se:
                p=subprocess.Popen([sys.executable,script],stdout=so,stderr=se,cwd=td)
                p.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            p.kill()
            result["stderr"]+="\nTimeout"
            result["exception"]="Timeout"
            return result
        result["execution_time"]=time.time()-start
        result["stdout"]=open(stdout_f).read()
        result["stderr"]=open(stderr_f).read()
        for f in sorted(os.listdir(plot_dir)):
            if f.endswith('.png'):
                result["plots"].append(open(os.path.join(plot_dir,f),'rb').read())
        for f in os.listdir(td):
            if f.startswith('df_') and f.endswith('.json'):
                result["dataframes"].append(json.load(open(os.path.join(td,f))))
    return result

def display_python_script_results(result):
    if not result: return
    st.info(f"Execution completed in {result['execution_time']:.2f}s")
    if result["exception"]:
        st.error(f"Exception: {result['exception']}")
    if result["stderr"]:
        st.error("Errors:")
        st.code(result["stderr"], language="bash")
    if result["plots"]:
        st.markdown("### Plots")
        cols=st.columns(min(3,len(result["plots"])))
        for i,p in enumerate(result["plots"]):
            cols[i%len(cols)].image(p,use_column_width=True)
    if result["dataframes"]:
        st.markdown("### DataFrames")
        for df in result["dataframes"]:
            with st.expander(f"{df['name']} {df['shape']}"):
                st.write(pd.read_html(df["preview"])[0])
    if result["stdout"]:
        st.markdown("### Stdout")
        st.code(result["stdout"], language="bash")

def parse_animation_steps(code):
    steps=[]
    plays=re.findall(r'self\.play\((.*?)\)',code,re.DOTALL)
    waits=re.findall(r'self\.wait\((.*?)\)',code,re.DOTALL)
    cum=0
    for i,pc in enumerate(plays):
        anims=[a.strip() for a in pc.split(',')]
        dur=1.0
        if i<len(waits):
            m=re.search(r'(\d+\.?\d*)',waits[i])
            if m: dur=float(m.group(1))
        steps.append({"id":i+1,"type":"play","animations":anims,"duration":dur,"start_time":cum,"code":f"self.play({pc})"})
        cum+=dur
    return steps

def generate_code_from_timeline(steps,orig):
    m=re.search(r'(class\s+\w+\s*\([^)]*\)\s*:.*?def\s+construct\s*\(\s*self\s*\)\s*:)',orig,re.DOTALL)
    if not m: return orig
    header=m.group(1)
    new=[header]
    indent="        "
    for s in sorted(steps,key=lambda x:x["id"]):
        new.append(f"{indent}{s['code']}")
        if s["duration"]>0:
            new.append(f"{indent}self.wait({s['duration']})")
    return "\n".join(new)

def create_timeline_editor(code):
    st.markdown("### 🎞️ Animation Timeline Editor")
    if not code:
        st.warning("Add animation code first")
        return code
    steps=parse_animation_steps(code)
    if not steps:
        st.warning("No steps detected")
        return code
    df=pd.DataFrame(steps)
    st.markdown("#### Animation Timeline")
    fig=px.timeline(df,x_start="start_time",x_end=df["start_time"]+df["duration"],y="id",color="type",hover_name="animations",labels={"id":"Step","start_time":"Time(s)"})
    fig.update_layout(height=300,xaxis=dict(title="Time(s)",rangeslider_visible=True))
    st.plotly_chart(fig,use_container_width=True)
    sel=st.selectbox("Select Step:",options=df["id"],format_func=lambda x:f"Step {x}")
    new_dur=st.number_input("Duration(s):",min_value=0.1,max_value=10.0,value=float(df[df["id"]==sel]["duration"].iloc[0]),step=0.1)
    action=st.selectbox("Action:",["Update Duration","Move Up","Move Down","Delete"])
    if st.button("Apply"):
        idx=df[df["id"]==sel].index[0]
        if action=="Update Duration":
            df.at[idx,"duration"]=new_dur
        elif action=="Move Up" and sel>1:
            j=df[df["id"]==sel-1].index[0]
            df.at[idx,"id"],df.at[j,"id"]=sel-1,sel
        elif action=="Move Down" and sel<len(df):
            j=df[df["id"]==sel+1].index[0]
            df.at[idx,"id"],df.at[j,"id"]=sel+1,sel
        elif action=="Delete":
            df=df[df["id"]!=sel]
            df["id"]=range(1,len(df)+1)
        cum=0
        for i in df.sort_values("id").index:
            df.at[i,"start_time"]=cum; cum+=df.at[i,"duration"]
        new_code=generate_code_from_timeline(df.to_dict('records'),code)
        st.success("Timeline updated, code regenerated.")
        return new_code
    return code

def export_to_educational_format(video_data,fmt,title,explanation,temp_dir):
    try:
        if fmt=="powerpoint":
            import pptx
            from pptx.util import Inches
            prs=pptx.Presentation()
            s0=prs.slides.add_slide(prs.slide_layouts[0]); s0.shapes.title.text=title; s0.placeholders[1].text="Created with Manim"
            s1=prs.slides.add_slide(prs.slide_layouts[5]); s1.shapes.title.text="Animation"
            vid_path=os.path.join(temp_dir,"anim.mp4"); open(vid_path,"wb").write(video_data)
            try:
                s1.shapes.add_movie(vid_path,Inches(1),Inches(1.5),Inches(8),Inches(4.5))
            except:
                thumb=os.path.join(temp_dir,"thumb.png")
                subprocess.run(["ffmpeg","-i",vid_path,"-ss","00:00:01","-vframes","1",thumb],check=True)
                s1.shapes.add_picture(thumb,Inches(1),Inches(1.5),Inches(8),Inches(4.5))
            if explanation:
                s2=prs.slides.add_slide(prs.slide_layouts[1]); s2.shapes.title.text="Explanation"; s2.placeholders[1].text=explanation
            out=os.path.join(temp_dir,f"{title.replace(' ','_')}.pptx"); prs.save(out)
            return open(out,"rb").read(),"pptx"
        if fmt=="html":
            html=f"""<!DOCTYPE html><html><head><title>{title}</title>
<style>body{{font-family:Arial;max-width:800px;margin:auto;padding:20px}}
.controls button{{margin-right:10px;padding:5px 10px}}</style>
<script>window.onload=function(){{const v=document.getElementById('anim');
document.getElementById('play').onclick=()=>v.play();
document.getElementById('pause').onclick=()=>v.pause();
document.getElementById('restart').onclick=()=>{{v.currentTime=0;v.play()}};
}};</script>
</head><body><h1>{title}</h1>
<video id="anim" width="100%" controls><source src="data:video/mp4;base64,{base64.b64encode(video_data).decode()}" type="video/mp4"></video>
<div class="controls"><button id="play">Play</button><button id="pause">Pause</button><button id="restart">Restart</button></div>
<div class="explanation">{markdown.markdown(explanation)}</div>
</body></html>"""
            out=os.path.join(temp_dir,f"{title.replace(' ','_')}.html"); open(out,"w").write(html)
            return open(out,"rb").read(),"html"
        if fmt=="sequence":
            from fpdf import FPDF
            vid=os.path.join(temp_dir,"anim.mp4"); open(vid,"wb").write(video_data)
            fr_dir=os.path.join(temp_dir,"frames"); os.makedirs(fr_dir,exist_ok=True)
            subprocess.run(["ffmpeg","-i",vid,"-r","1",os.path.join(fr_dir,"frame_%03d.png")],check=True)
            pdf=FPDF(); pdf.set_auto_page_break(True,15)
            pdf.add_page(); pdf.set_font("Arial","B",20); pdf.cell(190,10,title,0,1,"C")
            segs=explanation.split("##") if explanation else ["No explanation"]
            imgs=sorted([f for f in os.listdir(fr_dir) if f.endswith(".png")])
            for i,img in enumerate(imgs):
                pdf.add_page(); pdf.image(os.path.join(fr_dir,img),10,10,190)
                pdf.ln(100); pdf.set_font("Arial","B",12); pdf.cell(190,10,f"Step {i+1}",0,1)
                pdf.set_font("Arial","",10); pdf.multi_cell(190,5,segs[min(i,len(segs)-1)].strip())
            out=os.path.join(temp_dir,f"{title.replace(' ','_')}_seq.pdf"); pdf.output(out)
            return open(out,"rb").read(),"pdf"
    except Exception as e:
        logger.error(traceback.format_exc())
    return None,None

def main():
    if 'init' not in st.session_state:
        st.session_state.init=True
        st.session_state.video_data=None
        st.session_state.status=None
        st.session_state.ai_models=None
        st.session_state.generated_code=""
        st.session_state.code=""
        st.session_state.temp_code=""
        st.session_state.editor_key=str(uuid.uuid4())
        st.session_state.packages_checked=False
        st.session_state.latex_formula=""
        st.session_state.audio_path=None
        st.session_state.image_paths=[]
        st.session_state.custom_library_result=""
        st.session_state.python_script="""import matplotlib.pyplot as plt
import numpy as np

# Example: Create a simple plot
x = np.linspace(0, 10, 100)
y = np.sin(x)

plt.figure(figsize=(10, 6))
plt.plot(x, y, 'b-', label='sin(x)')
plt.title('Sine Wave')
plt.xlabel('x')
plt.ylabel('sin(x)')
plt.grid(True)
plt.legend()
"""
        st.session_state.python_result=None
        st.session_state.settings={"quality":"720p","format_type":"mp4","animation_speed":"Normal"}
        st.session_state.password_entered=False
    st.set_page_config(page_title="Manim Animation Studio", page_icon="🎬", layout="wide")
    st.markdown("""
    <style>
    /* custom CSS */
    </style>
    """, unsafe_allow_html=True)
    st.markdown("<h1 style='text-align:center;'>🎬 Manim Animation Studio</h1>", unsafe_allow_html=True)
    if not st.session_state.packages_checked:
        if ensure_packages():
            st.session_state.packages_checked=True
        else:
            st.error("Failed to install packages"); st.stop()
    if not ACE_EDITOR_AVAILABLE:
        try:
            from streamlit_ace import st_ace
            ACE_EDITOR_AVAILABLE=True
        except ImportError:
            pass
    tabs = st.tabs(["✨ Editor","🤖 AI Assistant","📚 LaTeX Formulas","🎨 Assets","🎞️ Timeline","🎓 Educational Export","🐍 Python Runner"])
    # --- Editor Tab ---
    with tabs[0]:
        col1,col2=st.columns([3,2])
        with col1:
            st.markdown("### 📝 Animation Editor")
            mode=st.radio("Input code:",["Type Code","Upload File"],key="editor_mode")
            if mode=="Upload File":
                up=st.file_uploader("Upload .py",type=["py"],key="file_up")
                if up:
                    txt=up.getvalue().decode("utf-8")
                    st.session_state.code=txt; st.session_state.temp_code=txt
            if ACE_EDITOR_AVAILABLE:
                code_in=st_ace(value=st.session_state.code,language="python",theme="monokai",min_lines=20,key=f"ace_{st.session_state.editor_key}")
            else:
                code_in=st.text_area("Code",value=st.session_state.code,height=400,key=f"ta_{st.session_state.editor_key}")
            if code_in!=st.session_state.code:
                st.session_state.code=code_in; st.session_state.temp_code=code_in
            if st.button("🚀 Generate Animation",key="gen"):
                if not st.session_state.code.strip():
                    st.error("Enter code first")
                else:
                    sc=extract_scene_class_name(st.session_state.code)
                    if sc=="MyScene" and "class MyScene" not in st.session_state.code:
                        df="""\nclass MyScene(Scene):\n    def construct(self):\n        text=Text("Default Scene"); self.play(Write(text)); self.wait(2)\n"""
                        st.session_state.code+=df; st.warning("No scene class; added default")
                    with st.spinner("Rendering..."):
                        d,s=generate_manim_video(st.session_state.code,st.session_state.settings["format_type"],st.session_state.settings["quality"],ANIMATION_SPEEDS[st.session_state.settings["animation_speed"]],st.session_state.audio_path)
                        st.session_state.video_data=d; st.session_state.status=s
        with col2:
            st.markdown("### 🖥️ Preview & Output")
            if st.session_state.code:
                st.markdown("<div style='border:1px solid #ccc;padding:10px;'>",unsafe_allow_html=True)
                st.components.v1.html(generate_manim_preview(st.session_state.code),height=250)
                st.markdown("</div>",unsafe_allow_html=True)
            if st.session_state.video_data:
                fmt=st.session_state.settings["format_type"]
                if fmt=="png_sequence":
                    st.download_button("⬇️ Download PNG Zip",st.session_state.video_data,file_name=f"frames_{int(time.time())}.zip")
                elif fmt=="svg":
                    try: st.components.v1.html(st.session_state.video_data.decode('utf-8'),height=400)
                    except: pass
                    st.download_button("⬇️ Download SVG",st.session_state.video_data,file_name=f"anim.svg")
                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"anim.{fmt}")
            if st.session_state.status:
                if "❌" in st.session_state.status: st.error(st.session_state.status)
                else: st.success(st.session_state.status)
    # --- AI Assistant Tab ---
    with tabs[1]:
        st.markdown("### 🤖 AI Animation Assistant")
        if check_password():
            if not st.session_state.ai_models:
                st.session_state.ai_models=init_ai_models_direct()
            # Debug & selection & generation (as in original)
            with st.expander("🔧 Debug Connection"):
                if st.button("Test API Connection"):
                    with st.spinner("Testing..."):
                        try:
                            token=get_secret("github_token_api")
                            if not token: st.error("Token missing"); st.stop()
                            model=st.session_state.ai_models["model_name"]
                            from openai import OpenAI
                            client=OpenAI(base_url="https://models.github.ai/inference",api_key=token)
                            params={"messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Hi"}],"model":model}
                            params[MODEL_CONFIGS[model]["param_name"]]=MODEL_CONFIGS[model][MODEL_CONFIGS[model]["param_name"]]
                            resp=client.chat.completions.create(**params)
                            if resp and resp.choices:
                                st.success("✅ Connected")
                            else: st.error("No response")
                        except Exception as e:
                            st.error(f"Error: {e}")
            st.markdown("### 🤖 Model Selection")
            cats={}
            for m,cfg in MODEL_CONFIGS.items():
                if m!="default":
                    cats.setdefault(cfg["category"],[]).append(m)
            cat_tabs=st.tabs(sorted(cats.keys()))
            for i,cat in enumerate(sorted(cats.keys())):
                with cat_tabs[i]:
                    for m in sorted(cats[cat]):
                        cfg=MODEL_CONFIGS[m]
                        sel=(m==st.session_state.ai_models["model_name"])
                        st.markdown(f"<div style='background:#f8f9fa;padding:10px;border-left:4px solid {'#0d6efd' if sel else '#4F46E5'};margin-bottom:8px;'>"
                                    f"<h4>{m}</h4><p>Max Tokens: {cfg.get(cfg['param_name'],'?')}</p><p>API Ver: {cfg['api_version'] or 'default'}</p></div>",
                                    unsafe_allow_html=True)
                        if st.button("Select" if not sel else "Selected ✓",key=f"sel_{m}",disabled=sel):
                            st.session_state.ai_models["model_name"]=m
                            st.experimental_rerun()
            if st.session_state.ai_models:
                st.info(f"Using model: {st.session_state.ai_models['model_name']}")

            if st.session_state.ai_models and "client" in st.session_state.ai_models:
                st.markdown("#### Generate Animation from Description")
                ideas=["...","3D sphere to torus","Pythagorean proof","Fourier transform","Neural network propagation","Integration area"]
                sel=st.selectbox("Try idea",ideas)
                prompt=sel if sel!="..." else ""
                inp=st.text_area("Your prompt or code",value=prompt,height=150)
                if st.button("Generate Animation Code"):
                    if inp:
                        with st.spinner("Generating..."):
                            code=suggest_code_completion(inp,st.session_state.ai_models)
                            if code:
                                st.session_state.generated_code=code
                            else: st.error("Failed")
                    else: st.warning("Enter prompt")
                if st.session_state.generated_code:
                    st.code(st.session_state.generated_code,language="python")
                    c1,c2=st.columns(2)
                    if c1.button("Use This Code"):
                        st.session_state.code=st.session_state.generated_code
                        st.experimental_rerun()
                    if c2.button("Render Preview"):
                        vd,stt=generate_manim_video(st.session_state.generated_code,"mp4","480p",1.0)
                        if vd: st.video(vd); st.download_button("Download Preview",vd,file_name="preview.mp4")
                        else: st.error(f"Error: {stt}")
        else:
            st.info("Enter password to access AI")

    # --- LaTeX Formulas Tab ---
    with tabs[2]:
        st.markdown("### 📚 LaTeX Formula Builder")
        c1,c2=st.columns([3,2])
        with c1:
            lt=st.text_area("LaTeX Formula",value=st.session_state.latex_formula,placeholder=r"e^{i\pi}+1=0",height=100)
            st.session_state.latex_formula=lt
            categories={
                "Basic Math":[{"name":"Fraction","latex":r"\frac{a}{b}"},...],
                # fill in as original categories...
            }
            tab_cats=st.tabs(list(categories.keys()))
            for i,(cat,forms) in enumerate(categories.items()):
                with tab_cats[i]:
                    for f in forms:
                        if st.button(f["name"],key=f"lt_{f['name']}"):
                            st.session_state.latex_formula=f["latex"]; st.experimental_rerun()
            if lt:
                snippet=f"""
formula=MathTex(r"{lt}")
self.play(Write(formula))
self.wait(2)
"""
                st.code(snippet,language="python")
                if st.button("Insert into Editor"):
                    if "def construct" in st.session_state.code:
                        lines=st.session_state.code.split("\n")
                        idx=[i for i,l in enumerate(lines) if "def construct" in l][0]
                        indent=re.match(r"(\s*)",lines[idx+1]).group(1) if idx+1<len(lines) else "        "
                        insert="\n".join(indent+line for line in snippet.strip().split("\n"))
                        lines.insert(idx+2,insert)
                        st.session_state.code="\n".join(lines)
                        st.experimental_rerun()
                    else:
                        base=f"""from manim import *\n\nclass LatexScene(Scene):\n    def construct(self):\n        {snippet.strip().replace('\n','\n        ')}\n"""
                        st.session_state.code=base; st.experimental_rerun()
        with c2:
            st.components.v1.html(render_latex_preview(st.session_state.latex_formula),height=300)

    # --- Assets Tab ---
    with tabs[3]:
        st.markdown("### 🎨 Asset Management")
        a1,a2=st.columns(2)
        with a1:
            imgs=st.file_uploader("Upload Images",type=["png","jpg","jpeg","svg"],accept_multiple_files=True)
            if imgs:
                d="manim_assets/images";os.makedirs(d,exist_ok=True)
                for up in imgs:
                    ext=up.name.split(".")[-1]
                    fn=f"img_{int(time.time())}_{uuid.uuid4().hex[:8]}.{ext}"
                    p=os.path.join(d,fn)
                    open(p,"wb").write(up.getvalue())
                    st.session_state.image_paths.append({"name":up.name,"path":p})
                st.success("Images uploaded")
            if st.session_state.image_paths:
                for ip in st.session_state.image_paths:
                    st.image(Image.open(ip["path"]),caption=ip["name"],width=100)
                    if st.button(f"Use {ip['name']}",key=f"use_img_{ip['name']}"):
                        code=f"""
image=ImageMobject(r"{ip['path']}")
self.play(FadeIn(image))
self.wait(1)
"""
                        st.session_state.code+=code; st.experimental_rerun()
        with a2:
            au=st.file_uploader("Upload Audio",type=["mp3","wav","ogg"])
            if au:
                d="manim_assets/audio";os.makedirs(d,exist_ok=True)
                fn=f"audio_{int(time.time())}.{au.name.split('.')[-1]}"
                p=os.path.join(d,fn)
                open(p,"wb").write(au.getvalue())
                st.session_state.audio_path=p
                st.audio(au)
                st.success("Audio uploaded")

    # --- Timeline Tab ---
    with tabs[4]:
        updated=create_timeline_editor(st.session_state.code)
        if updated!=st.session_state.code:
            st.session_state.code=updated; st.experimental_rerun()

    # --- Educational Export Tab ---
    with tabs[5]:
        st.markdown("### 🎓 Educational Export")
        if not st.session_state.video_data:
            st.warning("Generate animation first")
        else:
            title=st.text_input("Animation Title","Manim Animation")
            expl=st.text_area("Explanation",height=150)
            fmt=st.selectbox("Format",["PowerPoint Presentation","Interactive HTML","Explanation Sequence PDF"])
            if st.button("Export"):
                mp={"PowerPoint Presentation":"powerpoint","Interactive HTML":"html","Explanation Sequence PDF":"sequence"}
                data,typ=export_to_educational_format(st.session_state.video_data,mp[fmt],title,expl,tempfile.mkdtemp())
                if data:
                    ext={"powerpoint":"pptx","html":"html","sequence":"pdf"}[typ]
                    st.download_button("Download",data,file_name=f"{title.replace(' ','_')}.{ext}")
                else: st.error("Export failed")

    # --- Python Runner Tab ---
    with tabs[6]:
        st.markdown("### 🐍 Python Script Runner")
        examples={
            "Basic Plot":st.session_state.python_script,
            "Input Example":"""# input demo...""",
            "DataFrame":"""import pandas as pd...""",
        }
        choice=st.selectbox("Examples",list(examples.keys()))
        code=examples[choice] if choice in examples else st.session_state.python_script
        if ACE_EDITOR_AVAILABLE:
            code_in=st_ace(value=code,language="python",theme="monokai",min_lines=15,key=f"pyace_{st.session_state.editor_key}")
        else:
            code_in=st.text_area("Code",value=code,height=400,key=f"pyta_{st.session_state.editor_key}")
        st.session_state.python_script=code_in
        calls=detect_input_calls(code_in)
        vals=[]
        if calls:
            st.markdown("Provide inputs:")
            for i,c in enumerate(calls):
                v=st.text_input(c["prompt"],key=f"inp_{i}")
                vals.append(v)
        timeout=st.slider("Timeout",5,300,30)
        if st.button("▶️ Run"):
            res=run_python_script(code_in,vals,timeout)
            st.session_state.python_result=res
        if st.session_state.python_result:
            display_python_script_results(st.session_state.python_result)
            if st.session_state.python_result["plots"]:
                st.markdown("Add plot to animation:")
                for i,p in enumerate(st.session_state.python_result["plots"]):
                    st.image(p); 
                    if st.button(f"Use Plot {i+1}",key=f"use_plot_{i}"):
                        path=tempfile.NamedTemporaryFile(delete=False,suffix=".png").name
                        open(path,"wb").write(p)
                        code=f"""
plot_img=ImageMobject(r"{path}")
self.play(FadeIn(plot_img))
self.wait(1)
"""
                        st.session_state.code+=code; st.experimental_rerun()

if __name__ == "__main__":
    main()