from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import subprocess
import tempfile
import shutil
from datetime import datetime
import traceback
import json

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# Configuration
BASE_DIR = "/app"
MEDIA_DIR = os.path.join(BASE_DIR, "media")
TEMP_DIR = os.path.join(BASE_DIR, "temp")
os.makedirs(MEDIA_DIR, exist_ok=True)
os.makedirs(TEMP_DIR, exist_ok=True)

# API Key for security (optional)
API_KEY = "rkmentormindzofficaltokenkey12345"

import textwrap
from manim import *

def make_wrapped_paragraph(content, max_width, color, font, font_size, line_spacing, align_left=True):
    """
    Build a vertically stacked group of Text lines that together form a paragraph.
    It splits content into lines that fit within max_width by measuring rendered width.
    Each line is a separate Text object joined into a VGroup and arranged downward.
    """
    words = content.split()
    lines = []
    current = ""

    # Create a temporary Text to measure width; use the same font/size as final lines
    temp = Text("", color=color, font=font, font_size=font_size)

    for w in words:
        test = w if not current else current + " " + w
        test_obj = Text(test, color=color, font=font, font_size=font_size)
        if test_obj.width <= max_width:
            current = test
        else:
            # flush the current line
            line = Text(current, color=color, font=font, font_size=font_size)
            lines.append(line)
            current = w
    if current:
        lines.append(Text(current, color=color, font=font, font_size=font_size))

    if not lines:
        return VGroup()

    para = VGroup(*lines)
    # Space lines vertically; arrange them as a column
    para.arrange(DOWN, buff=line_spacing)
    if align_left:
        para = para.align_to(LEFT)
    return para.strip()

def create_manim_script(problem_data, script_path):
    """Generate Manim script from problem data with robust wrapping for title, text, and equations."""

    # Defaults
    settings = problem_data.get("video_settings", {
        "background_color": "#0f0f23",
        "text_color": "WHITE",
        "highlight_color": "YELLOW",
        "font": "CMU Serif",
        "text_size": 36,
        "equation_size": 42,
        "title_size": 48,
        "wrap_width": 12.0  # in scene width units; adjust to taste
    })

    slides = problem_data.get("slides", [])
    if not slides:
        raise ValueError("No slides provided in input data")

    slides_repr = repr(slides)

    # Use a dedicated wrap width in scene units; you can adapt how max_width is computed
    wrap_width = float(settings.get("wrap_width", 12.0))

    manim_code = f'''
from manim import *
import textwrap
class GeneratedMathScene(Scene):
    def construct(self):
        # Scene settings
        self.camera.background_color = "{settings.get('background_color', '#0f0f23')}"
        default_color = {settings.get('text_color', 'WHITE')}
        highlight_color = {settings.get('highlight_color', 'YELLOW')}
        default_font = "{settings.get('font', 'CMU Serif')}"
        text_size = {settings.get('text_size', 36)}
        equation_size = {settings.get('equation_size', 42)}
        title_size = {settings.get('title_size', 48)}
        wrap_width = {wrap_width}
        
        # Helper to wrap text into lines that fit within max width
        def make_wrapped_paragraph(content, color, font, font_size, line_spacing=0.2):
            lines = []
            words = content.split()
            current = ""
            for w in words:
                test = w if not current else current + " " + w
                test_obj = Text(test, color=color, font=font, font_size=font_size)
                if test_obj.width <= wrap_width * 0.9:  # a bit of padding
                    current = test
                else:
                    lines.append(Text(current, color=color, font=font, font_size=font_size))
                    current = w
            if current:
                lines.append(Text(current, color=color, font=font, font_size=font_size))
            if not lines:
                return VGroup()
            para = VGroup(*lines).arrange(DOWN, buff=line_spacing)
            return para
        class GeneratedMathSceneInner(Scene):
            pass
        content_group = VGroup()
        current_y = 3.0
        line_spacing = 0.8
        slides = {slides_repr}
        
        # Build each slide
        for idx, slide in enumerate(slides):
            obj = None
            content = slide.get("content", "")
            animation = slide.get("animation", "write_left")
            duration = slide.get("duration", 1.0)
            slide_type = slide.get("type", "text")
            
            if slide_type == "title":
                # Wrap title text
                title_text = content
                # Use paragraph wrapping to keep multi-line titles readable
                lines = []
                if title_text:
                    lines = []
                    # Reuse make_wrapped_paragraph by simulating a single paragraph
                    lines_group = make_wrapped_paragraph(title_text, highlight_color, default_font, title_size, line_spacing=0.2)
                    obj = lines_group if len(lines_group) > 0 else Text(title_text, color=highlight_color, font=default_font, font_size=title_size)
                else:
                    obj = Text("", color=highlight_color, font=default_font, font_size=title_size)
                if obj.width > wrap_width:
                    obj.scale_to_fit_width(wrap_width)
                
                obj.move_to(ORIGIN)
                self.play(FadeIn(obj), run_time=duration * 0.8)
                self.wait(duration * 0.3)
                self.play(FadeOut(obj), run_time=duration * 0.3)
                continue
            
            elif slide_type == "text":
                # Use wrapping for normal text
                obj = make_wrapped_paragraph(content, default_color, default_font, text_size, line_spacing=0.25)
            
            elif slide_type == "equation":
                # Wrap long equations by splitting content into lines if needed
                # Heuristic: if content is too wide, create a multi-line TeX using \\ line breaks
                eq_content = content
                # Optional: insert line breaks at common math breakpoints if needed
                test = MathTex(eq_content, color=default_color, font_size=equation_size)
                if test.width > wrap_width:
                    # naive wrap: insert line breaks at spaces near the middle
                    parts = eq_content.split(" ")
                    mid = len(parts)//2
                    line1 = " ".join(parts[:mid])
                    line2 = " ".join(parts[mid:])
                    wrapped_eq = f"{{line1}} \\\\\\\\ {{line2}}"
                    obj = MathTex(wrapped_eq, color=default_color, font_size=equation_size)
                else:
                    obj = MathTex(eq_content, color=default_color, font_size=equation_size)
                
                if obj.width > wrap_width:
                    obj.scale_to_fit_width(wrap_width)
            
            if obj:
                # Position and animate
                obj.to_edge(LEFT, buff=0.3)
                obj.shift(UP * (current_y - obj.height/2))
                
                obj_bottom = obj.get_bottom()[1]
                if obj_bottom < -3.5:
                    scroll_amount = abs(obj_bottom - (-3.5)) + 0.3
                    self.play(content_group.animate.shift(UP * scroll_amount), run_time=0.5)
                    current_y += scroll_amount
                    obj.shift(UP * scroll_amount)
                    obj.to_edge(LEFT, buff=0.3)
                
                if animation == "write_left":
                    self.play(Write(obj), run_time=duration)
                elif animation == "fade_in":
                    self.play(FadeIn(obj), run_time=duration)
                elif animation == "highlight_left":
                    self.play(Write(obj), run_time=duration * 0.6)
                    self.play(obj.animate.set_color(highlight_color), run_time=duration * 0.4)
                else:
                    self.play(Write(obj), run_time=duration)
                
                content_group.add(obj)
                # Decrease y for next item
                current_y -= (getattr(obj, "height", 0) + line_spacing)
                self.wait(0.3)
        
        if len(content_group) > 0:
            final_box = SurroundingRectangle(content_group[-1], color=highlight_color, buff=0.2)
            self.play(Create(final_box), run_time=0.8)
            self.wait(1.5)
    '''

    with open(script_path, 'w', encoding='utf-8') as f:
        f.write(manim_code)
    
    print(f"Generated script preview (first 500 chars):{manim_code[:500]}...")

@app.route("/")
def home():
    return "Flask Manim Video Generator is Running"

@app.route("/generate", methods=["POST", "OPTIONS"])
def generate_video():
    # Handle preflight
    if request.method == "OPTIONS":
        return '', 204
    
    try:
        # Optional: Check API key
        api_key = request.headers.get('X-API-KEY')
        if api_key and api_key != API_KEY:
            return jsonify({"error": "Invalid API key"}), 401
        
        # Get JSON data


        # Try reading raw body text
        raw_body = request.data.decode("utf-8").strip()
        data = None
        
        if not raw_body:
            return jsonify({"error": "No input data provided"}), 400
        
        # Try to detect if input is JSON or plain string
        if raw_body.startswith("{") or raw_body.startswith("["):
            # Likely JSON, try parsing
            try:
                data = json.loads(raw_body)
                print("✅ Detected valid JSON input.")
            except json.JSONDecodeError:
                # Not valid JSON, fallback to manual parse
                data = None
        
        if data is None:
            print("⚙️ Detected raw string input (non-JSON). Parsing manually...")
        
            # Handle format like: [ [...], [...]] &&& Tamil explanation
            parts = raw_body.split("&&&")
            slides_part = parts[0].strip()
            extra_info = parts[1].strip() if len(parts) > 1 else ""
        
            try:
                slides = json.loads(slides_part)
            except json.JSONDecodeError as e:
                return jsonify({
                    "error": "Failed to parse slides JSON from string input",
                    "details": str(e),
                    "raw_snippet": slides_part[:200]
                }), 400
        
            # Convert list of lists → list of dicts
            slides_json = []
            for s in slides:
                if len(s) >= 4:
                    slide_type, content, animation, duration = s
                    slides_json.append({
                        "type": slide_type,
                        "content": content,
                        "animation": animation,
                        "duration": duration
                    })
        
            data = {
                "slides": slides_json,
                "language": "Tamil" if "Tamil" in extra_info else "English",
                "explanation": extra_info,
                "video_settings": {
                    "background_color": "#0f0f23",
                    "text_color": "WHITE",
                    "highlight_color": "YELLOW",
                    "font": "CMU Serif"
                }
            }
        
        # ✅ Final validation
        if "slides" not in data or not data["slides"]:
            return jsonify({"error": "No slides provided in request"}), 400
        
        print(f"✅ Parsed {len(data['slides'])} slides successfully.")
        
        # Validate input
        if "slides" not in data or not data["slides"]:
            return jsonify({"error": "No slides provided in request"}), 400
        
        print(f"Received request with {len(data['slides'])} slides")
        
        # Create unique temporary directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        temp_work_dir = os.path.join(TEMP_DIR, f"manim_{timestamp}")
        os.makedirs(temp_work_dir, exist_ok=True)
        
        # Generate Manim script
        script_path = os.path.join(temp_work_dir, "scene.py")
        create_manim_script(data, script_path)
        print(f"Created Manim script at {script_path}")
        
        # Render video using subprocess
        quality = request.args.get('quality', 'l')  # l=low, m=medium, h=high
        render_command = [
            "manim",
            f"-q{quality}",
            "--disable_caching",
            "--media_dir", temp_work_dir,
            script_path,
            "GeneratedMathScene"
        ]
        
        print(f"Running command: {' '.join(render_command)}")
        
        result = subprocess.run(
            render_command,
            capture_output=True,
            text=True,
            cwd=temp_work_dir,
            timeout=120
        )
        
        if result.returncode != 0:
            error_msg = result.stderr or result.stdout
            print(f"Manim rendering failed: {error_msg}")
            return jsonify({
                "error": "Manim rendering failed",
                "details": error_msg
            }), 500
        
        print("Manim rendering completed successfully")
        
        # Find generated video
        quality_map = {'l': '480p15', 'm': '720p30', 'h': '1080p60'}
        video_quality = quality_map.get(quality, '480p15')
        
        video_path = os.path.join(
            temp_work_dir,
            "videos",
            "scene",
            video_quality,
            "GeneratedMathScene.mp4"
        )
        
        if not os.path.exists(video_path):
            print(f"Video not found at expected path: {video_path}")
            return jsonify({
                "error": "Video file not found after rendering",
                "expected_path": video_path
            }), 500
        
        print(f"Video found at: {video_path}")
        
        # Copy to media directory
        output_filename = f"math_video_{timestamp}.mp4"
        output_path = os.path.join(MEDIA_DIR, output_filename)
        shutil.copy(video_path, output_path)
        print(f"Video copied to: {output_path}")
        
        # Clean up temp directory
        try:
            shutil.rmtree(temp_work_dir)
            print("Cleaned up temp directory")
        except Exception as e:
            print(f"Failed to clean temp dir: {e}")
        
        # Return video file as blob
        return send_file(
            output_path,
            mimetype='video/mp4',
            as_attachment=False,
            download_name=output_filename
        )
        
    except subprocess.TimeoutExpired:
        print("Video rendering timeout")
        return jsonify({"error": "Video rendering timeout (120s)"}), 504
    except Exception as e:
        print(f"Error: {str(e)}")
        traceback.print_exc()
        return jsonify({
            "error": str(e),
            "traceback": traceback.format_exc()
        }), 500

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False)