app.py

from fastapi import FastAPI, Form
from fastapi.middleware.cors import CORSMiddleware
from time import sleep
from pydantic import BaseModel, Field
from langchain_google_genai import ChatGoogleGenerativeAI
from SystemPrompts.execute import system_prompt, debug_prompt
from SystemPrompts.planning import system_prompt_planner
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
import os
import subprocess
from dotenv import load_dotenv
from fastapi.responses import StreamingResponse
from fastapi.responses import FileResponse, JSONResponse
import uvicorn


app = FastAPI()

GENERATE_PATH = "/tmp/generate.py"
MEDIA_DIR = "/tmp"

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def read_root():
    return {"Hello": "World"}

# from waitress import serve
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# Define the desired structure
class ManimScript(BaseModel):
    """Manim script."""
    code: str = Field(..., description="manim script code with imports")
    classname: str = Field(..., description="classname for manim script")
    instructions: str = Field(..., description="summery of manim script")

# --- Constants ---
MAX_ATTEMPTS = 10

def generateSequences(query: str):
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-preview-04-17", temperature=0)
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt_planner.replace("{", "{{").replace("}", "}}")),
        ("human", "Generate animation sequences for Manim that give full context and avoid errors for this input:- {input}"),
    ])
    parser = JsonOutputParser(pydantic_object=ManimScript)
    chain = prompt | llm | parser
    print("[💾] Generated The Sequence, How its going to work!")
    return chain.invoke({"input": query})

def generateScript(sequence) -> ManimScript:
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-preview-04-17", temperature=0)
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt.replace("{", "{{").replace("}", "}}")),
        ("human", "Generate error-free, production ready manim animation script based on this sequences: {input}"),
    ])
    parser = JsonOutputParser(pydantic_object=ManimScript)
    chain = prompt | llm | parser
    print("[✅] Generated The Initial Script")
    return chain.invoke({"input": sequence})
# result = generateScript("Create Circle Animation")


def debugCode(code) -> ManimScript:
    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    prompt = ChatPromptTemplate.from_messages([
        ("system", debug_prompt.replace("{", "{{").replace("}", "}}")),
        ("human", "find all the errors and fix the code, make sure to check all the code. this is the code -> {input}"),
    ])
    parser = JsonOutputParser(pydantic_object=ManimScript)
    chain = prompt | llm | parser
    print("[💾] Debugging Done")
    return chain.invoke({"input": code['code']})
# result = generateScript("Create Circle Animation")

# Keep only this
def writeInFile(parsed):
    print(type(parsed))  # Should be <class 'dict'>
    try:
        with open(GENERATE_PATH, "w", encoding="utf-8") as file:
            file.write(parsed['code'])
        print(f"[💾] Code written to {GENERATE_PATH}")
    except Exception as e:
        print(f"[❌] Writing to file failed: {e}")



def fixScript(error_message: str, current_code: str, depth: int = 1) -> ManimScript:
    if depth > 3:
        print("[🔁] Fixing recursion limit reached.")
        return ManimScript(code=current_code, classname="Unknown", instructions="Fix failed.")

    try:
        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
        parser = JsonOutputParser(pydantic_object=ManimScript)

        sleep(5)
        print("Taking 10 sec Break")
        sleep(10)
        print("Ready>")
        sleep(5)
        # Define prompt and structured chain
        prompt = ChatPromptTemplate.from_messages([
            ("system", """
    You are a Manim animation expert and code debugger.

    Your task is to analyze Manim code, identify issues, and return corrected code that keeps the same animation intent. You must also explain the errors briefly unless instructed otherwise.

    ## Rules:
    - Target Manim Community Edition (latest version).
    - Check for common issues like wrong syntax, deprecated functions, missing imports, object misuse, etc.
    - Keep the animation logic, style, and structure as close to the original as possible.
    - If a function is incorrect, replace it with a valid one (`FadeIn`, `Write`, `Create`, etc.).
    - If an object is misused (e.g., treating `Text` like a list), fix it.
    - Handle missing `self.wait()` or `self.play()` issues if relevant.
    - Output both the **corrected code** and a **summary of the fixes**.

    ## Input Format:
    The user will provide faulty or non-working Manim code. You will respond with:

    1. A brief explanation of the issues
    2. The corrected Python code in a single block


"""),
            ("human", """Fix this ManimCE code with the given error:
                        ```python
                        {code}
                        ```

                        Error message:
                        ```
                        {error}
                        ```

                        Return JSON with:
                "classname": the class name
                "code": fixed code
                "instructions": what the animation does
             """)])

        chain = prompt | llm | parser
        # Invoke chain with input
        return chain.invoke({"code": current_code, "error": error_message})
    except Exception as e:
        print(f"[⚠️] Fix failed at depth {depth}, retrying...")
        return fixScript(f"{error_message}\n\n{e}", current_code, depth + 1)

def executeWithAutoFix(parsed):
    if not parsed:
        print("[❌] No parsed code to execute.")
        return

    for attempt in range(1, MAX_ATTEMPTS + 1):
        writeInFile(parsed)
        command = f"manim -ql --media_dir {MEDIA_DIR} {GENERATE_PATH} {parsed['classname']}"

        try:
            result = subprocess.run(
                command,
                shell=True,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            print("[✅] Execution successful!")
            print(result.stdout)
            path = f"/tmp/videos/generate/480p15/{parsed['classname']}.mp4"
            return path

        except subprocess.CalledProcessError as e:
            error_msg = e.stderr or e.stdout or str(e)
            print(f"[❌] Attempt {attempt}/{MAX_ATTEMPTS} failed:\n{error_msg}\nTrying auto-fix...")
            parsed = fixScript(error_msg, parsed.get('code', ''), depth=1)

    print("[🚫] Max attempts reached. Could not execute successfully.")

@app.get('/')
def hello_world():
   return 'Hello World'


@app.post('/generate')
def generate(query: str = Form(...)):
    def generate_stream():
        yield f"data: Starting Animation: {query}\n\n"
        result = generateSequences(query)

        yield f"data: Preparing: {query}!\n\n"
        code = generateScript(result)

        yield f"data: Generating Code: {code['classname']}.\n\n"
        yield f"data: {code['instructions']}\n\n"

        debugged_code = debugCode(code)
        yield "data: Fixing Errors!\n\n"

        path = executeWithAutoFix(debugged_code)

        if path:
            yield f"data: <path>{path}</path> : contains the video\n\n"
        else:
            yield "data: Failed to generate video after multiple attempts.\n\n"

    return StreamingResponse(generate_stream(), media_type="text/event-stream")


@app.get('/stream')
def stream_data():
    def generate():
        for i in range(10):
            sleep(1)
            yield f"Data chunk {i}\n"
    return StreamingResponse(generate(), status=200, content_type="text/event-stream")

@app.post("/video")
def get_video(path: str = Form(...)):
    # Convert to an absolute path
    video_path = path  # Full path already, no need to join with os.getcwd()

    print(f"Video path: {video_path}")  # Log the path for debugging

    # Check if file exists
    if os.path.exists(video_path):
        return FileResponse(video_path, media_type="video/mp4", filename=os.path.basename(video_path))
    else:
        return JSONResponse(content={"error": "File not found"}, status_code=404)


if __name__ == "__main__":
    try:
        uvicorn.run("main:app", host="0.0.0.0", port=5000)
    except KeyboardInterrupt:
        print("Server stopped by user")