app.py

"""
PDF → Manim Animation Pipeline
Hugging Face Spaces — Gradio 6.x
"""

import asyncio
import atexit
import queue
import threading
import uuid

import gradio as gr

from queue_manager import JobQueue, State
from pipeline import run_pipeline

def _cleanup_event_loop():
    try:
        loop = asyncio.get_event_loop()
        if not loop.is_closed():
            loop.close()
    except Exception:
        pass

atexit.register(_cleanup_event_loop)

job_queue = JobQueue(max_workers=8, max_jobs=100)


def submit_and_stream(pdf_file, api_key: str):
    """
    Generator yielding exactly 4 values per tick:
        status_md  (str)
        code_box   (gr.update)
        video_player (gr.update)
        zip_download (gr.update)
    """

    def _err(msg):
        return (
            msg,
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False),
        )

    if pdf_file is None:
        yield _err("❌ Please upload a PDF file.")
        return
    if not api_key or len(api_key) < 10:
        yield _err("❌ Please enter a valid Gemini API key.")
        return
    if job_queue.is_full():
        yield _err("⚠️ Queue is full (max 100 jobs). Please try again shortly.")
        return

    job_id   = uuid.uuid4().hex
    pdf_path = pdf_file.name
    job_queue.register(job_id)

    update_q: queue.Queue = queue.Queue()

    def status_cb(state: State, message: str = "", code: str | None = None):
        update_q.put((state, message, code))

    result_holder: dict = {}

    def _run():
        try:
            result = run_pipeline(
                job_id=job_id,
                pdf_path=pdf_path,
                gemini_api_key=api_key,
                status_cb=status_cb,
            )
            result_holder.update(result)
            update_q.put((State.DONE, "✅ Render complete!", result.get("code")))
        except Exception as exc:
            update_q.put((State.FAILED, f"❌ {exc}", None))
        finally:
            update_q.put(None)

    threading.Thread(target=_run, daemon=True).start()

    icons = {State.QUEUED: "⏳", State.RUNNING: "⚙️", State.DONE: "✅", State.FAILED: "❌"}

    code_so_far = ""

    # Initial tick
    yield (
        f"⏳ **Queued** — Starting…\n\n*Job `{job_id}`*",
        gr.update(visible=False),
        gr.update(visible=False),
        gr.update(visible=False),
    )

    while True:
        item = update_q.get()
        if item is None:
            break

        state, message, code = item
        if code:
            code_so_far = code

        status_text = f"{icons.get(state,'❓')} **{state.value.title()}** — {message}\n\n*Job `{job_id}`*"
        is_done = state == State.DONE

        yield (
            status_text,
            gr.update(value=code_so_far, visible=bool(code_so_far)),
            gr.update(value=result_holder.get("video_path") if is_done else None, visible=is_done),
            gr.update(value=result_holder.get("zip_path")   if is_done else None, visible=is_done),
        )


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="PDF → Manim Video") as demo:

    gr.Markdown("# 🎬 PDF → Manim Animation Pipeline\nUpload a PDF and get a downloadable Manim animation.")

    saved_api_key = gr.BrowserState("")

    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="📄 Upload PDF", file_types=[".pdf"])
            api_key_input = gr.Textbox(
                label="🔑 Gemini API Key",
                placeholder="AIza…",
                type="password",
                info="Saved in your browser — enter once.",
            )
            submit_btn = gr.Button("🚀 Generate Video", variant="primary")

        with gr.Column(scale=1):
            status_md    = gr.Markdown("*Submit a job to see live status here.*")
            code_box     = gr.Code(label="📝 Generated Manim Code", language="python", visible=False, interactive=False)
            video_player = gr.Video(label="🎬 Rendered Animation", visible=False, interactive=False)
            zip_download = gr.File(label="⬇️ Download Artifacts (.py + .mp4)", visible=False, interactive=False)

    gr.Markdown("---\n**Notes:** Processing takes 2–5 minutes. The artifacts ZIP contains the `.py` source and rendered `.mp4`. Your API key is never stored server-side.")

    demo.load(fn=lambda k: k, inputs=[saved_api_key], outputs=[api_key_input])
    api_key_input.change(fn=lambda v: v, inputs=[api_key_input], outputs=[saved_api_key])

    submit_btn.click(
        fn=submit_and_stream,
        inputs=[pdf_input, api_key_input],
        outputs=[status_md, code_box, video_player, zip_download],
    )


if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(),
        ssr_mode=False,
        server_name="0.0.0.0",
        server_port=7860,
    )