"""
ManimGL MCP Server

This server provides tools for AI models to generate and render
mathematical animations using ManimGL (3Blue1Brown's animation engine).

Features:
- Write and render ManimGL scenes
- Interactive sessions with live coding
- Screenshot capture for visual feedback
- Comprehensive help system
"""

import asyncio
import json
import logging
from typing import Any

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    ImageContent,
    TextContent,
    Tool,
)

from .docs import (
    get_example_template,
    get_full_documentation,
    get_topic_help,
)
from .interactive import InteractiveSessionManager
from .scene_manager import SceneManager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize server
server = Server("manimgl-mcp")

# Initialize managers
scene_manager = SceneManager()
interactive_manager = InteractiveSessionManager()


@server.list_tools()
async def list_tools() -> list[Tool]:
    """List all available ManimGL tools."""
    return [
        # === File-based Scene Tools ===
        Tool(
            name="write_scene",
            description="""Write ManimGL scene code to a Python file.

The code should define Scene subclasses with construct() methods.
Always start with 'from manimlib import *'.

Example:
```python
from manimlib import *

class MyScene(Scene):
    def construct(self):
        circle = Circle(color=BLUE)
        self.play(ShowCreation(circle))
        self.wait()
```""",
            inputSchema={
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "Python code defining ManimGL scene(s). Should include 'from manimlib import *' and Scene class definitions.",
                    },
                    "filename": {
                        "type": "string",
                        "description": "Optional filename (without .py extension). If not provided, a temp file will be created.",
                    },
                    "output_dir": {
                        "type": "string",
                        "description": "Optional directory to save the file. Defaults to a temp directory.",
                    },
                },
                "required": ["code"],
            },
        ),
        Tool(
            name="render_scene",
            description="""Render a ManimGL scene to a video or image file.

Can render from a file path or directly from code. Returns the path to the rendered file.

Quality options:
- low: 480p (fast)
- medium: 720p (default)
- hd: 1080p
- 4k: 2160p (slow)

Format options:
- mp4: Video file (default)
- gif: Animated GIF
- png: Final frame only""",
            inputSchema={
                "type": "object",
                "properties": {
                    "scene_name": {
                        "type": "string",
                        "description": "Name of the Scene class to render (e.g., 'MyScene')",
                    },
                    "filepath": {
                        "type": "string",
                        "description": "Path to Python file containing the scene. Either filepath or code must be provided.",
                    },
                    "code": {
                        "type": "string",
                        "description": "Python code defining the scene. Either filepath or code must be provided.",
                    },
                    "quality": {
                        "type": "string",
                        "enum": ["low", "medium", "hd", "4k"],
                        "description": "Render quality. Default is medium (720p).",
                    },
                    "format": {
                        "type": "string",
                        "enum": ["mp4", "gif", "png"],
                        "description": "Output format. Default is mp4.",
                    },
                    "output_dir": {
                        "type": "string",
                        "description": "Directory for output files. Defaults to ./manim_output/",
                    },
                },
                "required": ["scene_name"],
            },
        ),
        Tool(
            name="list_scenes",
            description="List all Scene classes defined in a Python file or code snippet.",
            inputSchema={
                "type": "object",
                "properties": {
                    "filepath": {
                        "type": "string",
                        "description": "Path to Python file to analyze. Either filepath or code must be provided.",
                    },
                    "code": {
                        "type": "string",
                        "description": "Python code to analyze. Either filepath or code must be provided.",
                    },
                },
            },
        ),
        Tool(
            name="get_scene_preview",
            description="""Render just the final frame of a scene as a preview image.

Faster than full rendering - useful for checking the end state of an animation.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "scene_name": {
                        "type": "string",
                        "description": "Name of the Scene class to preview",
                    },
                    "filepath": {
                        "type": "string",
                        "description": "Path to Python file containing the scene",
                    },
                    "code": {"type": "string", "description": "Python code defining the scene"},
                    "return_base64": {
                        "type": "boolean",
                        "description": "If true, return the image as base64. Default is true for inline viewing.",
                    },
                },
                "required": ["scene_name"],
            },
        ),
        # === Interactive Session Tools ===
        Tool(
            name="start_interactive_session",
            description="""Start an interactive ManimGL session for live coding.

This opens a persistent ManimGL environment where you can:
- Create and manipulate mobjects in real-time
- Play animations incrementally
- Capture screenshots to see the current state
- Experiment with different approaches

Returns a session_id to use with other session commands.

Available shortcuts in session:
- add(*mobjects) - Add objects to scene
- remove(*mobjects) - Remove objects
- play(*animations) - Play animations
- wait(duration) - Pause
- clear() - Clear all objects""",
            inputSchema={
                "type": "object",
                "properties": {
                    "initial_code": {
                        "type": "string",
                        "description": "Optional initial code to run in the session (e.g., 'circle = Circle()')",
                    },
                    "headless": {
                        "type": "boolean",
                        "description": "If true, use virtual framebuffer (for servers without display). Default is auto-detect.",
                    },
                },
            },
        ),
        Tool(
            name="run_in_session",
            description="""Execute code in an active interactive ManimGL session.

Use this to:
- Create mobjects: 'circle = Circle(color=BLUE)'
- Add to scene: 'add(circle)'
- Play animations: 'play(ShowCreation(circle))'
- Transform objects: 'play(circle.animate.scale(2))'

Set capture_after=true to automatically take a screenshot after execution.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {
                        "type": "string",
                        "description": "The session ID returned from start_interactive_session",
                    },
                    "code": {
                        "type": "string",
                        "description": "Python code to execute in the session",
                    },
                    "capture_after": {
                        "type": "boolean",
                        "description": "If true, capture a screenshot after code execution. Default is false.",
                    },
                },
                "required": ["session_id", "code"],
            },
        ),
        Tool(
            name="capture_screenshot",
            description="""Capture a screenshot of the current scene state in an interactive session.

Returns the image as base64 by default, allowing you to see what's currently displayed.
Use this frequently to verify your animations look correct.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {"type": "string", "description": "The session ID"},
                    "return_base64": {
                        "type": "boolean",
                        "description": "If true (default), return image as base64. Otherwise return file path.",
                    },
                },
                "required": ["session_id"],
            },
        ),
        Tool(
            name="get_scene_state",
            description="""Get information about the current scene state in an interactive session.

Returns:
- Number of mobjects in scene
- Current scene time
- Basic info about each mobject""",
            inputSchema={
                "type": "object",
                "properties": {"session_id": {"type": "string", "description": "The session ID"}},
                "required": ["session_id"],
            },
        ),
        Tool(
            name="get_session_history",
            description="Get the command history for an interactive session.",
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {"type": "string", "description": "The session ID"},
                    "last_n": {
                        "type": "integer",
                        "description": "If provided, only return the last N commands",
                    },
                },
                "required": ["session_id"],
            },
        ),
        Tool(
            name="end_session",
            description="End an interactive ManimGL session and clean up resources.",
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {"type": "string", "description": "The session ID to end"}
                },
                "required": ["session_id"],
            },
        ),
        Tool(
            name="list_sessions",
            description="List all active interactive sessions with their status and statistics.",
            inputSchema={"type": "object", "properties": {}},
        ),
        # === Help & Documentation Tools ===
        Tool(
            name="get_manim_help",
            description="""Get help information about ManimGL classes, methods, or concepts.

Available topics:
- scene: How to create and structure scenes
- mobjects: Available shapes and objects
- animations: Animation types and usage
- colors: Color constants and usage
- coordinates: Coordinate system and positioning
- interactive: Interactive mode shortcuts
- transforms: Transformation animations
- text: Text and LaTeX rendering
- graphs: Graphing and plotting
- 3d: 3D scenes and objects""",
            inputSchema={
                "type": "object",
                "properties": {
                    "topic": {
                        "type": "string",
                        "description": "Topic to get help on (e.g., 'animations', 'mobjects', 'scene')",
                    }
                },
                "required": ["topic"],
            },
        ),
        Tool(
            name="get_example_code",
            description="""Get example ManimGL code for common animation patterns.

Examples include:
- basic_shapes: Creating and animating shapes
- transform_demo: Morphing between objects
- text_latex: Text and LaTeX rendering
- graph_animation: Plotting functions with moving dot
- value_tracker: Animated values with updaters
- 3d_demo: 3D scenes with camera rotation
- updater_pattern: Using updaters for dynamic relationships
- interactive_session: Commands for interactive mode""",
            inputSchema={
                "type": "object",
                "properties": {
                    "example_name": {
                        "type": "string",
                        "description": "Name of the example to retrieve",
                    }
                },
                "required": ["example_name"],
            },
        ),
        Tool(
            name="get_full_documentation",
            description="""Get comprehensive ManimGL API documentation.

Returns detailed documentation covering:
- All Scene classes and their methods
- All Mobject types (shapes, text, graphs, 3D)
- All Animation types with parameters
- Color constants and coordinate system
- Updaters and ValueTrackers
- Interactive mode commands

Use this when you need complete information about ManimGL capabilities.
The documentation is ~30KB of detailed API reference.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "section": {
                        "type": "string",
                        "description": "Optional: specific section (scene, mobjects, animations, etc.) or 'all' for everything",
                    }
                },
            },
        ),
        # === RAG Tools ===
        Tool(
            name="search_examples",
            description="""Search indexed ManimGL scene examples using semantic search.

Use this to find relevant code examples for a concept you want to animate.
Returns matching examples with full source code.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Natural language description of what you want to animate",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum results (default: 5)",
                    },
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="search_api",
            description="""Search ManimGL API signatures and documentation.

Find classes, methods, and their parameters. Use category filter for:
- animation: Animation classes (Transform, FadeIn, etc.)
- mobject: Object classes (Circle, Text, Graph, etc.)
- scene: Scene classes
- util: Utility functions""",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "What API you're looking for",
                    },
                    "category": {
                        "type": "string",
                        "enum": ["animation", "mobject", "scene", "util"],
                        "description": "Filter by category",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum results (default: 10)",
                    },
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="search_patterns",
            description="""Search animation pattern templates.

Find reusable patterns for common visualizations like:
- Riemann sums, derivatives, integrals (calculus)
- Matrix transforms, vector fields (linear_algebra)
- Projectile motion, waves (physics)
- Graph traversal, sorting (algorithms)
- 3D surfaces, rotations (3d)""",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "What kind of animation pattern you need",
                    },
                    "category": {
                        "type": "string",
                        "description": "Filter by category (calculus, linear_algebra, physics, etc.)",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum results (default: 5)",
                    },
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="validate_api_call",
            description="""Validate a ManimGL API call against indexed signatures.

Checks if parameters are valid for a given class. Returns suggestions for invalid calls.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "class_name": {
                        "type": "string",
                        "description": "Name of the class (e.g., 'Circle', 'Transform')",
                    },
                    "parameters": {
                        "type": "object",
                        "description": "Parameters being passed to the class",
                    },
                },
                "required": ["class_name", "parameters"],
            },
        ),
        Tool(
            name="index_manimgl",
            description="""Build or rebuild the RAG index from ManimGL source code.

Indexes API signatures, examples, and documentation for semantic search.
Run this once after installation or to update the index.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "source_path": {
                        "type": "string",
                        "description": "Path to manimlib source (auto-detected if not provided)",
                    },
                },
            },
        ),
        # === Multi-Agent Pipeline ===
        Tool(
            name="generate_animation",
            description="""Generate ManimGL animation code using multi-agent pipeline.

This is the recommended way to create animations. It:
1. Analyzes your concept
2. Plans the scene with multiple animation patterns
3. Generates production-quality code
4. Reviews and fixes issues

Returns complete, working ManimGL code.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "description": {
                        "type": "string",
                        "description": "Describe what you want to animate",
                    },
                    "style": {
                        "type": "string",
                        "enum": ["educational", "artistic", "technical", "minimal"],
                        "description": "Animation style (default: educational)",
                    },
                    "duration": {
                        "type": "integer",
                        "description": "Target duration in seconds (default: 15)",
                    },
                },
                "required": ["description"],
            },
        ),
        # === Audio Narration ===
        Tool(
            name="generate_narration",
            description="""Generate audio narration for an animation.

Creates TTS audio that can be synced to video. Uses edge-tts (free, no API key needed).""",
            inputSchema={
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Narration text",
                    },
                    "voice": {
                        "type": "string",
                        "description": "Voice to use (en-US-male, en-US-female, en-GB-male, etc.)",
                    },
                    "output_path": {
                        "type": "string",
                        "description": "Output file path (auto-generated if not provided)",
                    },
                },
                "required": ["text"],
            },
        ),
        Tool(
            name="add_narration_to_video",
            description="""Add narration audio to a rendered video.

Automatically syncs audio duration to video length.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "video_path": {
                        "type": "string",
                        "description": "Path to video file",
                    },
                    "audio_path": {
                        "type": "string",
                        "description": "Path to audio file",
                    },
                    "output_path": {
                        "type": "string",
                        "description": "Output path (auto-generated if not provided)",
                    },
                },
                "required": ["video_path", "audio_path"],
            },
        ),
        # === Self-Learning ===
        Tool(
            name="store_error_fix",
            description="""Store an error pattern and its fix for future reference.

The system learns from errors to provide better suggestions.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "error_pattern": {
                        "type": "string",
                        "description": "The error message pattern",
                    },
                    "fix": {
                        "type": "string",
                        "description": "How to fix this error",
                    },
                    "example": {
                        "type": "string",
                        "description": "Optional code example of the fix",
                    },
                },
                "required": ["error_pattern", "fix"],
            },
        ),
        Tool(
            name="get_similar_errors",
            description="""Find similar errors and their fixes from learned patterns.""",
            inputSchema={
                "type": "object",
                "properties": {
                    "error_message": {
                        "type": "string",
                        "description": "The error message you encountered",
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum results (default: 3)",
                    },
                },
                "required": ["error_message"],
            },
        ),
        Tool(
            name="get_rag_stats",
            description="""Get statistics about the RAG index.

Shows number of indexed examples, API signatures, patterns, and errors.""",
            inputSchema={
                "type": "object",
                "properties": {},
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent | ImageContent]:
    """Handle tool calls."""
    try:
        if name == "write_scene":
            result = await scene_manager.write_scene(
                code=arguments["code"],
                filename=arguments.get("filename"),
                output_dir=arguments.get("output_dir"),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "render_scene":
            result = await scene_manager.render_scene(
                scene_name=arguments["scene_name"],
                filepath=arguments.get("filepath"),
                code=arguments.get("code"),
                quality=arguments.get("quality", "medium"),
                format=arguments.get("format", "mp4"),
                output_dir=arguments.get("output_dir"),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "list_scenes":
            result = await scene_manager.list_scenes(
                filepath=arguments.get("filepath"), code=arguments.get("code")
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_scene_preview":
            result = await scene_manager.get_scene_preview(
                scene_name=arguments["scene_name"],
                filepath=arguments.get("filepath"),
                code=arguments.get("code"),
                return_base64=arguments.get(
                    "return_base64", True
                ),  # Default to True for inline viewing
            )
            if result.get("base64"):
                return [ImageContent(type="image", data=result["base64"], mimeType="image/png")]
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "start_interactive_session":
            result = await interactive_manager.start_session(
                initial_code=arguments.get("initial_code"),
                headless=arguments.get("headless", False),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "run_in_session":
            result = await interactive_manager.run_code(
                session_id=arguments["session_id"],
                code=arguments["code"],
                capture_after=arguments.get("capture_after", False),
            )
            # If screenshot was captured, return it as image
            if result.get("screenshot", {}).get("base64"):
                return [
                    TextContent(
                        type="text",
                        text=json.dumps(
                            {k: v for k, v in result.items() if k != "screenshot"}, indent=2
                        ),
                    ),
                    ImageContent(
                        type="image", data=result["screenshot"]["base64"], mimeType="image/png"
                    ),
                ]
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "capture_screenshot":
            result = await interactive_manager.capture_screenshot(
                session_id=arguments["session_id"],
                return_base64=arguments.get("return_base64", True),
            )
            if result.get("base64"):
                return [ImageContent(type="image", data=result["base64"], mimeType="image/png")]
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_scene_state":
            result = await interactive_manager.get_scene_state(session_id=arguments["session_id"])
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_session_history":
            result = await interactive_manager.get_session_history(
                session_id=arguments["session_id"], last_n=arguments.get("last_n")
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "end_session":
            result = await interactive_manager.end_session(session_id=arguments["session_id"])
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "list_sessions":
            result = await interactive_manager.list_sessions()
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_manim_help":
            result = get_topic_help(arguments["topic"])
            return [TextContent(type="text", text=result)]

        elif name == "get_example_code":
            result = get_example_template(arguments["example_name"])
            return [TextContent(type="text", text=result)]

        elif name == "get_full_documentation":
            section = arguments.get("section", "all")
            if section and section.lower() != "all":
                result = get_topic_help(section)
            else:
                result = get_full_documentation()
            return [TextContent(type="text", text=result)]

        # === RAG Tools ===
        elif name == "search_examples":
            from .rag import search_examples

            result = await search_examples(
                query=arguments["query"],
                limit=arguments.get("limit", 5),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "search_api":
            from .rag import search_api

            result = await search_api(
                query=arguments["query"],
                category=arguments.get("category"),
                limit=arguments.get("limit", 10),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "search_patterns":
            from .rag import search_patterns

            result = await search_patterns(
                query=arguments["query"],
                category=arguments.get("category"),
                limit=arguments.get("limit", 5),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "validate_api_call":
            from .rag.search import validate_api_call

            result = await validate_api_call(
                class_name=arguments["class_name"],
                parameters=arguments.get("parameters", {}),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "index_manimgl":
            from pathlib import Path

            from .rag.indexer import index_manimgl

            source_path = arguments.get("source_path")
            if not source_path:
                # Try to find manimlib in common locations
                import manimlib

                source_path = Path(manimlib.__file__).parent

            result = await index_manimgl(source_path)
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_rag_stats":
            from .rag.store import get_store

            store = get_store()
            result = store.get_stats()
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        # === Multi-Agent Pipeline ===
        elif name == "generate_animation":
            from .llm import run_multi_agent_pipeline
            from .rag import search_examples

            # Get relevant examples for context
            examples = await search_examples(arguments["description"], limit=3)
            example_codes = [ex.get("code", "") for ex in examples if ex.get("code")]

            result = await run_multi_agent_pipeline(
                user_request=arguments["description"],
                examples=example_codes,
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        # === Audio Narration ===
        elif name == "generate_narration":
            from .audio import generate_narration

            result = await generate_narration(
                text=arguments["text"],
                voice=arguments.get("voice", "en-US-male"),
                output_path=arguments.get("output_path"),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "add_narration_to_video":
            from .audio import sync_audio_to_video

            result = sync_audio_to_video(
                video_path=arguments["video_path"],
                audio_path=arguments["audio_path"],
                output_path=arguments.get("output_path"),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        # === Self-Learning ===
        elif name == "store_error_fix":
            from .rag.search import store_error_fix

            result = await store_error_fix(
                error_pattern=arguments["error_pattern"],
                fix=arguments["fix"],
                example=arguments.get("example"),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        elif name == "get_similar_errors":
            from .rag.search import get_similar_errors

            result = await get_similar_errors(
                error_message=arguments["error_message"],
                limit=arguments.get("limit", 3),
            )
            return [TextContent(type="text", text=json.dumps(result, indent=2))]

        else:
            return [TextContent(type="text", text=f"Unknown tool: {name}")]

    except Exception as e:
        logger.exception(f"Error executing tool {name}")
        return [
            TextContent(
                type="text", text=json.dumps({"error": str(e), "type": type(e).__name__}, indent=2)
            )
        ]


# Keep the old functions as aliases for backward compatibility
def get_example_code(example_name: str) -> str:
    """Get example code for common animation patterns."""
    legacy_aliases = {
        "basic_circle": "basic_shapes",
        "shape_transform": "transform_demo",
        "text_animation": "text_latex",
        "graph_plot": "graph_animation",
        "3d_scene": "3d_demo",
    }
    normalized = example_name.lower().replace("-", "_").replace(" ", "_")
    mapped = legacy_aliases.get(normalized, example_name)
    return get_example_template(mapped)


def get_manim_help(topic: str) -> str:
    """Get help information about ManimGL concepts."""
    result = get_topic_help(topic)
    normalized = topic.lower().strip()

    # Backward-compatible augmentation for tests and legacy callers
    if normalized == "mobjects" and "Circle" not in result and "Square" not in result:
        result = (
            result
            + "\n\n## Common Mobject Types\n"
            + "- Circle()\n"
            + "- Square()\n"
            + "- Rectangle()\n"
            + "- Triangle()\n"
        )

    return result


async def run_server():
    """Run the MCP server."""
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


def main():
    """Main entry point."""
    asyncio.run(run_server())


if __name__ == "__main__":
    main()
