from flask import Flask, render_template, request, url_for, send_from_directory
from google import genai
import re
import subprocess
import os
import shutil
import time
from threading import Timer

app = Flask(__name__)

# Load API key from environment variable
API_KEY = os.environ.get("GOOGLE_API_KEY")
if not API_KEY:
    raise ValueError("Missing GOOGLE_API_KEY environment variable.")
client = genai.Client(api_key=API_KEY)

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        prompt = request.form.get("prompt")
        if not prompt:
            return render_template("index.html")
        
        while True:
            try:
                ai_response = client.models.generate_content(
                    model="gemini-2.0-flash-lite-preview-02-05",
                    contents=f"""You are 'Manimator', an expert Manim animator and coder. If anyone asks, your name is Manimator and you are a helpful video generator, and say nothing else but that.
The user wants you to code this: {prompt}. 
Plan out in chain of thought what you are going to do first, then give the final code output in ```python``` codeblock.
Make sure to not use external images or resources other than default Manim.
Keep the scene uncluttered and aesthetically pleasing.
You got this!! <3
"""
                )
            except Exception as e:
                print("Error calling GenAI API:", e)
                time.sleep(2)
                continue

            pattern = r"```python\s*(.*?)\s*```"
            match = re.search(pattern, ai_response.text, re.DOTALL)
            if match:
                code = match.group(1)
            else:
                print("No python code block found in the AI response, retrying...")
                time.sleep(2)
                continue

            scene_match = re.search(r"class\s+(\w+)\(.*Scene.*\):", code)
            if scene_match:
                scene_name = scene_match.group(1)
            else:
                scene_name = "MyScene"
                
            code_filename = "generated_manim_scene.py"
            try:
                with open(code_filename, "w") as f:
                    f.write(code)
            except Exception as e:
                print("Error writing code file:", e)
                time.sleep(2)
                continue 

            video_filename = "output_video.mp4"
            cmd = [
                "manim",
                "-qm",
                "-o", video_filename,
                code_filename,
                scene_name
            ]
            try:
                subprocess.run(cmd, check=True)
            except subprocess.CalledProcessError as e:
                print("Error during Manim rendering:", e)
                time.sleep(2)
                continue

            # Construct the expected output path from Manim
            video_path_in_media = os.path.join("media", "videos", code_filename.replace(".py", ""), "720p30", video_filename)
            # Instead of moving to a read-only directory, move it to /tmp (which is writable)
            tmp_video_path = os.path.join("/tmp", video_filename)

            try:
                shutil.move(video_path_in_media, tmp_video_path)
            except FileNotFoundError:
                print("Manim output file not found. Check your manim code and output location.")
                time.sleep(2)
                continue
            except Exception as e:
                print("Error moving video file:", e)
                time.sleep(2)
                continue

            # Generate a URL that points to the get_video route
            video_url = url_for('get_video', filename=video_filename)
            return render_template("result.html", video_url=video_url)
            
    return render_template("index.html")

@app.route("/video/<filename>")
def get_video(filename):
    video_path = os.path.join("/tmp", filename)
    response = send_from_directory("/tmp", filename)

    # Schedule deletion of the video file after 5 seconds
    def remove_file():
        try:
            os.remove(video_path)
            print("Removed video file:", video_path)
        except Exception as e:
            app.logger.error("Error removing video file: %s", e)
    Timer(5, remove_file).start()

    return response

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