app.py

import math
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Literal

import gradio as gr
import spaces
import uvicorn
from fastapi import BackgroundTasks, FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field, field_validator


api = FastAPI(title="MathX Manim Renderer", version="1.0.0")
api.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://mathx-ifal.vercel.app",
        "https://mathx-nine.vercel.app",
        "http://localhost:3000",
        "http://localhost:8000",
    ],
    allow_methods=["GET", "POST"],
    allow_headers=["Content-Type"],
)


class GraphSpec(BaseModel):
    type: Literal["quadratic", "linear", "sine"]
    formula: str = Field(default="Visualização matemática", max_length=160)
    a: float = 1
    b: float = 0
    c: float = 0
    d: float = 0
    m: float = 1

    @field_validator("a", "b", "c", "d", "m")
    @classmethod
    def finite_coefficient(cls, value: float) -> float:
        if not math.isfinite(value) or abs(value) > 1000:
            raise ValueError("Coeficiente inválido")
        return value


def scene_source(spec: GraphSpec) -> str:
    if spec.type == "quadratic":
        expression = f"({spec.a})*x**2 + ({spec.b})*x + ({spec.c})"
    elif spec.type == "linear":
        expression = f"({spec.m})*x + ({spec.b})"
    else:
        expression = f"({spec.a})*np.sin(({spec.b})*x + ({spec.c})) + ({spec.d})"

    safe_formula = spec.formula.replace("\\", "\\\\").replace('"', '\\"')
    return f'''from manim import *
import numpy as np

class MathXScene(Scene):
    def construct(self):
        self.camera.background_color = "#090A0D"
        axes = Axes(
            x_range=[-6, 6, 1],
            y_range=[-5, 7, 1],
            x_length=10,
            y_length=6,
            axis_config={{"color": "#60636D", "stroke_width": 1.5}},
            tips=True,
        )
        labels = axes.get_axis_labels(Text("x", font_size=24), Text("y", font_size=24))
        title = Text("{safe_formula}", font_size=28, color=WHITE).to_edge(UP)
        graph = axes.plot(lambda x: {expression}, x_range=[-6, 6], color="#D8FF3E", stroke_width=5)
        self.play(Create(axes), FadeIn(labels), Write(title), run_time=1.1)
        self.play(Create(graph), run_time=2.8, rate_func=smooth)
        self.wait(0.7)
'''


def remove_job(path: str) -> None:
    shutil.rmtree(path, ignore_errors=True)


@api.get("/health")
def health():
    return {"service": "MathX Manim Renderer", "status": "ready"}


@api.post("/render")
@spaces.GPU(duration=60)
def render(spec: GraphSpec, background_tasks: BackgroundTasks):
    job_dir = Path(tempfile.mkdtemp(prefix="mathx-"))
    scene_file = job_dir / "scene.py"
    scene_file.write_text(scene_source(spec), encoding="utf-8")
    command = [
        "manim", "-ql", "--disable_caching", "--format=mp4",
        "--media_dir", str(job_dir / "media"), str(scene_file), "MathXScene",
    ]
    try:
        subprocess.run(
            command,
            cwd=job_dir,
            check=True,
            timeout=45,
            capture_output=True,
            text=True,
        )
        videos = list((job_dir / "media").rglob("MathXScene.mp4"))
        if not videos:
            raise RuntimeError("O Manim não produziu o vídeo esperado.")
    except subprocess.TimeoutExpired:
        remove_job(str(job_dir))
        raise HTTPException(status_code=504, detail="A renderização excedeu 45 segundos.")
    except subprocess.CalledProcessError as exc:
        remove_job(str(job_dir))
        detail = (exc.stderr or exc.stdout or "Falha ao executar o Manim.")[-1200:]
        raise HTTPException(status_code=500, detail=detail)
    except Exception as exc:
        remove_job(str(job_dir))
        raise HTTPException(status_code=500, detail=str(exc))

    background_tasks.add_task(remove_job, str(job_dir))
    return FileResponse(
        videos[0],
        media_type="video/mp4",
        filename="mathx-manim.mp4",
        background=background_tasks,
    )


with gr.Blocks(title="MathX Manim Renderer") as demo:
    gr.Markdown(
        "# MathX Manim Renderer\n"
        "Serviço de renderização de animações matemáticas utilizado pelo MathX. "
        "A API está disponível em `/render`."
    )

app = gr.mount_gradio_app(api, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)