"""
Scene Manager for ManimGL MCP Server

Handles writing scene files, rendering, and listing scenes.
"""

import ast
import asyncio
import base64
import os
import re
import tempfile
from typing import Any


class SceneManager:
    """Manages ManimGL scene files and rendering."""

    def __init__(self, default_output_dir: str | None = None):
        self.default_output_dir = default_output_dir or os.path.join(os.getcwd(), "manim_output")
        self.temp_dir = tempfile.mkdtemp(prefix="manim_mcp_")

    async def write_scene(
        self, code: str, filename: str | None = None, output_dir: str | None = None
    ) -> dict[str, Any]:
        """
        Write ManimGL scene code to a Python file.

        Args:
            code: Python code defining ManimGL scene(s)
            filename: Optional filename (without .py extension)
            output_dir: Optional directory to save the file

        Returns:
            Dict with filepath and any detected scene names
        """
        # Ensure code has the import
        if "from manimlib import" not in code and "import manimlib" not in code:
            code = "from manimlib import *\n\n" + code

        # Determine output path
        directory = output_dir or self.temp_dir
        os.makedirs(directory, exist_ok=True)

        if filename:
            # Sanitize filename
            filename = re.sub(r"[^\w\-_]", "_", filename)
            if not filename.endswith(".py"):
                filename += ".py"
        else:
            # Generate a temp filename
            import uuid

            filename = f"scene_{uuid.uuid4().hex[:8]}.py"

        filepath = os.path.join(directory, filename)

        # Write the file
        with open(filepath, "w") as f:
            f.write(code)

        # Detect scene names
        scenes = self._extract_scene_names(code)

        return {
            "success": True,
            "filepath": filepath,
            "scenes": scenes,
            "message": f"Wrote scene file with {len(scenes)} scene(s)",
        }

    async def render_scene(
        self,
        scene_name: str,
        filepath: str | None = None,
        code: str | None = None,
        quality: str = "medium",
        format: str = "mp4",
        output_dir: str | None = None,
    ) -> dict[str, Any]:
        """
        Render a ManimGL scene to video or image.

        Args:
            scene_name: Name of the Scene class to render
            filepath: Path to Python file containing the scene
            code: Python code defining the scene (alternative to filepath)
            quality: Render quality (low, medium, hd, 4k)
            format: Output format (mp4, gif, png)
            output_dir: Directory for output files

        Returns:
            Dict with output file path and render information
        """
        if not filepath and not code:
            return {"error": "Either filepath or code must be provided"}

        # If code provided, write it to a file first
        if code:
            result = await self.write_scene(code)
            filepath = result["filepath"]

        if not os.path.exists(filepath):
            return {"error": f"File not found: {filepath}"}

        # Set up output directory
        output_directory = output_dir or self.default_output_dir
        os.makedirs(output_directory, exist_ok=True)

        # Build manimgl command
        cmd = ["manimgl", filepath, scene_name]

        # Add quality flag
        quality_flags = {"low": "-l", "medium": "-m", "hd": "--hd", "4k": "--uhd"}
        if quality in quality_flags:
            cmd.append(quality_flags[quality])

        # Add format-specific flags
        if format == "png":
            cmd.extend(["-s"])  # Skip to end, save last frame
        elif format == "gif":
            cmd.extend(["-i"])  # Save as gif

        # Always write to file (not window) and specify output dir
        cmd.extend(["-w", "--video_dir", output_directory])

        # Run manimgl
        try:
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=os.environ.copy(),
            )
            try:
                stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=300)
            except asyncio.TimeoutError:
                process.kill()
                return {"error": "Render timed out after 300 seconds"}

            if process.returncode != 0:
                return {
                    "error": "Render failed",
                    "stderr": stderr.decode(),
                    "stdout": stdout.decode(),
                    "command": " ".join(cmd),
                }

            # Find the output file
            output_file = self._find_output_file(output_directory, scene_name, format)

            return {
                "success": True,
                "output_file": output_file,
                "format": format,
                "quality": quality,
                "stdout": stdout.decode(),
                "message": f"Rendered {scene_name} to {output_file}",
            }

        except FileNotFoundError:
            return {
                "error": "manimgl not found. Is ManimGL installed?",
                "hint": "Install with: pip install manimgl",
            }
        except Exception as e:
            return {"error": str(e), "type": type(e).__name__}

    async def list_scenes(
        self, filepath: str | None = None, code: str | None = None
    ) -> dict[str, Any]:
        """
        List all Scene classes in a file or code.

        Args:
            filepath: Path to Python file to analyze
            code: Python code to analyze

        Returns:
            Dict with list of scene names
        """
        if not filepath and not code:
            return {"error": "Either filepath or code must be provided"}

        if filepath:
            if not os.path.exists(filepath):
                return {"error": f"File not found: {filepath}"}
            with open(filepath) as f:
                code = f.read()

        scenes = self._extract_scene_names(code)

        return {"success": True, "scenes": scenes, "count": len(scenes)}

    async def get_scene_preview(
        self,
        scene_name: str,
        filepath: str | None = None,
        code: str | None = None,
        return_base64: bool = False,
    ) -> dict[str, Any]:
        """
        Render the final frame of a scene as a preview.

        Args:
            scene_name: Name of the Scene class to preview
            filepath: Path to Python file containing the scene
            code: Python code defining the scene
            return_base64: If True, return image as base64 string

        Returns:
            Dict with image path or base64 data
        """
        # Render as PNG (last frame only)
        result = await self.render_scene(
            scene_name=scene_name, filepath=filepath, code=code, quality="medium", format="png"
        )

        if "error" in result:
            return result

        output_file = result.get("output_file")
        if not output_file or not os.path.exists(output_file):
            return {"error": "Preview file not generated"}

        if return_base64:
            with open(output_file, "rb") as f:
                image_data = base64.b64encode(f.read()).decode("utf-8")
            return {"success": True, "base64": image_data, "filepath": output_file}

        return {"success": True, "filepath": output_file}

    def _extract_scene_names(self, code: str) -> list[str]:
        """Extract Scene class names from Python code using AST."""
        try:
            tree = ast.parse(code)
            scenes = []

            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    # Check if it inherits from Scene or a known Scene subclass
                    for base in node.bases:
                        base_name = ""
                        if isinstance(base, ast.Name):
                            base_name = base.id
                        elif isinstance(base, ast.Attribute):
                            base_name = base.attr

                        if "Scene" in base_name:
                            scenes.append(node.name)
                            break

            return scenes
        except SyntaxError as e:
            return [f"SyntaxError: {e}"]

    def _find_output_file(self, output_dir: str, scene_name: str, format: str) -> str | None:
        """Find the rendered output file."""
        ext_map = {"mp4": ".mp4", "gif": ".gif", "png": ".png"}
        ext = ext_map.get(format, ".mp4")

        # ManimGL typically puts files in output_dir/scene_name.ext
        # or output_dir/filename/scene_name.ext
        possible_paths = [
            os.path.join(output_dir, f"{scene_name}{ext}"),
        ]

        # Also check subdirectories
        if os.path.isdir(output_dir):
            for item in os.listdir(output_dir):
                item_path = os.path.join(output_dir, item)
                if os.path.isdir(item_path):
                    possible_paths.append(os.path.join(item_path, f"{scene_name}{ext}"))

        # Find the most recently modified matching file
        existing_files = [p for p in possible_paths if os.path.exists(p)]
        if existing_files:
            return max(existing_files, key=os.path.getmtime)

        # Fallback: search recursively
        for root, _dirs, files in os.walk(output_dir):
            for f in files:
                if scene_name in f and f.endswith(ext):
                    return os.path.join(root, f)

        return None

    def cleanup(self):
        """Clean up temporary files."""
        import shutil

        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)
