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
import re
import shutil
import time
from datetime import datetime
import streamlit.components.v1 as components
import uuid
import pandas as pd
import plotly.express as px
import zipfile
import traceback

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

# 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 = {
    "gpt-4o": {},
    "gpt-4o-mini": {},
    "gpt-4.1": {},
    "gpt-4.1-mini": {},
    "o1": {},
    "o1-mini": {},
    "default": {}
}

# 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, using text_area")

def get_secret(env_var):
    val = os.environ.get(env_var)
    if not val:
        logger.warning(f"Secret '{env_var}' not set")
    return val

def check_password():
    pwd = get_secret("password")
    if not pwd:
        st.error("Admin password not configured")
        return False
    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")
                return False
        return False
    return True

def ensure_packages():
    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 reqs.items():
        try:
            __import__(pkg if pkg != "Pillow" else "PIL")
        except ImportError:
            missing[pkg] = ver
    if not missing:
        return True
    bar = st.progress(0)
    txt = st.empty()
    for i, (pkg, ver) in enumerate(missing.items()):
        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((i + 1) / len(missing))
    txt.empty()
    return True

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):
    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:8px;text-align:center;">
      <h3>Preview</h3>
      <div style="font-size:2rem;">{preview_icons}</div>
      <p>Full render for accurate output</p>
    </div>
    """

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")
    with open(scene_file, "w") as f:
        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}", f"--fps={fps}"]
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    progress = st.progress(0)
    log = st.empty()
    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"✅ Done ({size_mb:.1f} MB)"
    else:
        return None, "❌ No output generated"

def main():
    st.set_page_config("Manim Studio", "🎬", layout="wide")
    if "inited" not in st.session_state:
        st.session_state.update({
            "inited": True,
            "code": "",
            "video_data": None,
            "status": "",
            "settings": {"quality":"720p","fps":30,"format":"mp4"},
            "pwd_ok": False,
            "custom_model": "gpt-4o",
            "generated_code": ""
        })

    if not st.session_state.get("packages_checked", False):
        if ensure_packages():
            st.session_state.packages_checked = True
        else:
            return

    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 tab
    with tabs[0]:
        c1, c2 = st.columns([3,2])
        with c1:
            st.markdown("### 📝 Animation Editor")
            if ACE_EDITOR_AVAILABLE:
                st.session_state.code = st_ace(
                    value=st.session_state.code,
                    language="python",
                    theme="monokai",
                    min_lines=20,
                    key="ace"
                )
            else:
                st.session_state.code = st.text_area(
                    "Manim Code",
                    value=st.session_state.code,
                    height=400
                )
            if st.button("🚀 Generate Animation"):
                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"],
                        st.session_state["settings"]["quality"],
                        st.session_state["settings"]["fps"]
                    )
                    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)
            if st.session_state.video_data:
                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")
                        components.html(svg, height=400)
                    except:
                        st.error("Cannot display SVG")
                    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,
                                       f"animation.{fmt}", f"video/{fmt}")
            if 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 tab
    with tabs[1]:
        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 tab
    with tabs[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 += "\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 tab
    with tabs[3]:
        st.markdown("### 🎞️ Timeline Editor")
        st.info("Adjust code manually – timeline UI coming soon.")

    # Export tab
    with tabs[4]:
        st.markdown("### 🎓 Export")
        st.warning("Export features coming soon.")

    # Python Runner tab
    with tabs[5]:
        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()