"""
Interactive Session Manager for ManimGL MCP Server

Manages headless IPython sessions via pexpect.  manimlib is pre-loaded into
the REPL so all animation classes (Circle, Transform, …) are immediately
available without ever opening a display window.
"""

import asyncio
import os
import re
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

_ANSI_RE = re.compile(r"\x1b\[[\x30-\x3f]*[\x20-\x2f]*[\x40-\x7e]|\x1b\][^\x07]*\x07|\x1b[@-Z\\-_]")

try:
    import pexpect

    PEXPECT_AVAILABLE = True
except ImportError:
    PEXPECT_AVAILABLE = False


@dataclass
class InteractiveSession:
    """Represents an interactive ManimGL session."""

    id: str
    process: Any  # pexpect.spawn
    created_at: float = field(default_factory=time.time)
    last_activity: float = field(default_factory=time.time)
    history: list[dict] = field(default_factory=list)
    screenshot_dir: str = field(
        default_factory=lambda: tempfile.mkdtemp(prefix="manim_screenshots_")
    )
    screenshot_count: int = 0

    def update_activity(self):
        self.last_activity = time.time()

    def get_screenshot_path(self) -> str:
        """Get path for next screenshot."""
        self.screenshot_count += 1
        return os.path.join(self.screenshot_dir, f"screenshot_{self.screenshot_count:04d}.png")


class InteractiveSessionManager:
    """
    Manages headless interactive sessions backed by IPython + manimlib.

    Sessions spawn `python -m IPython` with `from manimlib import *`
    pre-loaded so all animation classes are immediately available.
    No display, no window — fully headless on every platform.
    """

    def __init__(self, timeout: int = 60, max_sessions: int = 5):
        self.sessions: dict[str, InteractiveSession] = {}
        self.timeout = timeout
        self.max_sessions = max_sessions
        # IPython prompt pattern - matches "In [1]:", "In [2]:", etc.
        # Must be broad enough to match through ANSI color codes that wrap the digits
        self.ipython_prompt = r"In \["
        # Also match continuation prompts
        self.continuation_prompt = r"\.\.\.\.:|\s+\.\.\.\."

    async def start_session(
        self,
        initial_code: str | None = None,
        headless: bool = False,  # noqa: ARG002
        resolution: tuple[int, int] = (1280, 720),  # noqa: ARG002
    ) -> dict[str, Any]:
        """
        Start a new headless interactive session (IPython + manimlib).

        Args:
            initial_code: Optional Python code to run immediately after startup.
            headless: Ignored — sessions are always headless (no window).
            resolution: Ignored — no window is opened.

        Returns:
            Dict with session_id, status, and available commands.
        """
        if not PEXPECT_AVAILABLE:
            return {
                "error": "pexpect not available",
                "hint": "Install pexpect: pip install pexpect",
                "fix": "uv add pexpect",
            }

        # Check session limit
        if len(self.sessions) >= self.max_sessions:
            return {
                "error": f"Maximum sessions ({self.max_sessions}) reached",
                "hint": "End an existing session first using end_session",
                "active_sessions": list(self.sessions.keys()),
            }

        session_id = str(uuid.uuid4())[:8]
        screenshot_dir = tempfile.mkdtemp(prefix=f"manim_session_{session_id}_")

        try:
            import shutil
            import sys

            # Spawn a headless IPython process — manimlib depends on IPython so it's
            # always available.  We deliberately do NOT invoke manimgl, which would open
            # an OpenGL window.  Instead we get a plain IPython REPL and pre-load
            # manimlib so all classes (Circle, Transform, …) are available without any
            # Scene being constructed or any window being created.
            # Use the python interpreter from the same venv as manimgl.
            # We run `python -m IPython` rather than the `ipython` binary because
            # manimlib overwrites the `ipython` entry-point with its own CLI wrapper.
            manimgl_bin = shutil.which("manimgl") or "manimgl"
            venv_bin_dir = os.path.dirname(os.path.abspath(manimgl_bin))

            python_bin = (
                os.path.join(venv_bin_dir, "python")
                if os.path.exists(os.path.join(venv_bin_dir, "python"))
                else os.path.join(venv_bin_dir, "python3")
                if os.path.exists(os.path.join(venv_bin_dir, "python3"))
                else shutil.which("python3") or sys.executable
            )

            # Try IPython first (rich prompt, %cpaste support for multiline)
            try:
                import subprocess as _sp

                _sp.run([python_bin, "-c", "import IPython"], check=True, capture_output=True)
                spawn_cmd = python_bin
                spawn_args = [
                    "-m",
                    "IPython",
                    "--no-banner",
                    "--no-confirm-exit",
                    "--no-autoindent",
                ]
            except Exception:
                # Fall back to plain python -i
                spawn_cmd = python_bin
                spawn_args = ["-i"]

            env = {
                **os.environ,
                "TERM": "dumb",
                "PYTHONUNBUFFERED": "1",
            }

            # Start the process — no display or window required
            process = pexpect.spawn(
                spawn_cmd,
                args=spawn_args,
                encoding="utf-8",
                timeout=self.timeout,
                env=env,
                dimensions=(50, 200),  # rows, cols for terminal
            )

            # Wait for the first IPython prompt — startup is fast (no GL init)
            startup_timeout = 30
            try:
                process.expect(self.ipython_prompt, timeout=startup_timeout)
                startup_output = process.before or ""
            except pexpect.TIMEOUT:
                process.terminate(force=True)
                return {
                    "error": "Timeout waiting for IPython to start",
                    "hint": "Ensure IPython is installed in the same environment as manimgl",
                    "partial_output": process.before if process.before else "",
                }
            except pexpect.EOF:
                return {
                    "error": "IPython process terminated unexpectedly on startup",
                    "output": process.before if process.before else "",
                }

            # Reset sys.argv so manimlib's argparse doesn't choke on IPython's own
            # startup flags, then import all manimlib names into the REPL namespace.
            process.sendline("import sys; sys.argv = ['ipython']")
            try:
                process.expect(self.ipython_prompt, timeout=10)
            except (pexpect.TIMEOUT, pexpect.EOF):
                pass

            process.sendline("from manimlib import *")
            try:
                process.expect(self.ipython_prompt, timeout=30)
            except (pexpect.TIMEOUT, pexpect.EOF):
                pass  # best-effort; import may still have worked

            session = InteractiveSession(
                id=session_id, process=process, screenshot_dir=screenshot_dir
            )
            self.sessions[session_id] = session

            result = {
                "success": True,
                "session_id": session_id,
                "message": "Interactive session started (headless IPython + manimlib)",
                "startup_output": startup_output.strip(),
                "screenshot_directory": screenshot_dir,
                "available_shortcuts": {
                    "play": "play(*animations) - Play animations",
                    "wait": "wait(duration) - Pause for duration seconds",
                    "add": "add(*mobjects) - Add mobjects to scene",
                    "remove": "remove(*mobjects) - Remove mobjects from scene",
                    "clear": "clear() - Remove all mobjects",
                    "save_state": "save_state() - Save current state for undo",
                    "undo": "undo() - Undo last change",
                    "redo": "redo() - Redo undone change",
                },
                "tips": [
                    "Use 'circle = Circle()' to create objects",
                    "Use 'add(circle)' to show objects",
                    "Use 'play(ShowCreation(circle))' to animate",
                    "Use 'capture_screenshot' tool to see the current state",
                ],
            }

            # Run initial code if provided
            if initial_code:
                code_result = await self.run_code(session_id, initial_code)
                result["initial_code_result"] = code_result

            return result

        except pexpect.exceptions.ExceptionPexpect as e:
            return {
                "error": f"Failed to start IPython session: {str(e)}",
                "hint": "Is IPython installed in the same environment as manimgl?",
                "fix": "pip install ipython",
            }
        except FileNotFoundError:
            return {
                "error": "IPython/Python interpreter not found",
                "hint": "IPython is not installed or not in PATH",
                "fix": "pip install ipython",
            }
        except Exception as e:
            return {
                "error": str(e),
                "type": type(e).__name__,
                "hint": "Unexpected error starting session",
            }

    async def run_code(
        self,
        session_id: str,
        code: str,
        capture_after: bool = False,
        wait_for_animation: bool = True,
    ) -> dict[str, Any]:
        """
        Execute code in an interactive session.

        Args:
            session_id: The session ID
            code: Python code to execute
            capture_after: If True, capture a screenshot after execution
            wait_for_animation: If True, wait longer for animations to complete

        Returns:
            Dict with execution result, output, and optional screenshot
        """
        session = self.sessions.get(session_id)
        if not session:
            return {
                "error": f"Session not found: {session_id}",
                "available_sessions": list(self.sessions.keys()),
                "hint": "Start a new session with start_interactive_session",
            }

        if not session.process.isalive():
            del self.sessions[session_id]
            return {
                "error": "Session is no longer active",
                "hint": "Start a new session with start_interactive_session",
            }

        try:
            # Prepare the code - handle multi-line properly
            code = code.strip()
            lines = code.split("\n")

            full_output = []
            errors = []

            # Determine timeout based on code content
            timeout = self.timeout
            if wait_for_animation and any(kw in code for kw in ["play(", "wait(", "animate"]):
                timeout = max(timeout, 120)  # At least 2 minutes for animations

            # Send code to IPython.
            # For multi-line input use %cpaste so the whole block is treated as
            # one logical cell, avoiding the "first prompt consumed early" bug.
            if len(lines) == 1:
                session.process.sendline(code)
            else:
                session.process.sendline("%cpaste -q")
                await asyncio.sleep(0.05)
                for line in lines:
                    session.process.sendline(line)
                    await asyncio.sleep(0.02)
                # Two dashes ends %cpaste mode
                session.process.sendline("--")

            # Wait for the prompt
            try:
                index = session.process.expect(
                    [self.ipython_prompt, pexpect.TIMEOUT, pexpect.EOF], timeout=timeout
                )

                output = session.process.before or ""

                # Clean up the output
                output = self._clean_output(output, code)

                if index == 0:  # Got prompt
                    if output.strip():
                        full_output.append(output.strip())
                elif index == 1:  # Timeout
                    full_output.append("(execution timed out - animation may still be running)")
                elif index == 2:  # EOF — IPython process exited
                    self.sessions.pop(session_id, None)
                    return {
                        "error": "Session ended unexpectedly (IPython process exited)",
                        "output": _ANSI_RE.sub("", output),
                        "hint": "Start a new session with start_interactive_session",
                    }

            except Exception as e:
                errors.append(str(e))

            # Check for Python errors — use ANSI-stripped text so color codes don't hide keywords
            error_patterns = [
                r"Traceback \(most recent call last\):",
                r"\bError:",
                r"\bException:",
                r"\bSyntaxError:",
                r"\bNameError:",
                r"\bTypeError:",
                r"\bValueError:",
                r"\bAttributeError:",
            ]

            combined_output = "\n".join(full_output)
            has_error = any(re.search(pattern, combined_output) for pattern in error_patterns)

            session.update_activity()

            # Record in history
            history_entry = {
                "code": code,
                "output": combined_output,
                "timestamp": time.time(),
                "had_error": has_error,
            }
            session.history.append(history_entry)

            result = {
                "success": not has_error,
                "output": combined_output,
                "session_id": session_id,
                "command_number": len(session.history),
            }

            if has_error:
                result["error_detected"] = True
                result["hint"] = "Check the output for error details"

            if errors:
                result["warnings"] = errors

            # Capture screenshot if requested
            if capture_after:
                screenshot_result = await self.capture_screenshot(session_id)
                if screenshot_result.get("success"):
                    result["screenshot"] = screenshot_result

            return result

        except pexpect.EOF:
            self.sessions.pop(session_id, None)
            return {
                "error": "Session ended unexpectedly (IPython process exited)",
                "hint": "Start a new session with start_interactive_session",
            }
        except Exception as e:
            return {"error": str(e), "type": type(e).__name__}

    async def capture_screenshot(
        self,
        session_id: str,
        return_base64: bool = True,  # noqa: ARG002
    ) -> dict[str, Any]:
        """
        Not supported in headless IPython sessions.

        Interactive sessions run inside a plain IPython REPL — there is no live
        OpenGL scene to snapshot.  To get a visual of your work, write the scene
        to a file with `write_scene` and render a PNG preview with
        `get_scene_preview`.

        Args:
            session_id: The session ID (checked for existence)
            return_base64: Ignored

        Returns:
            Dict explaining the limitation with suggested alternatives.
        """
        session = self.sessions.get(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        if not session.process.isalive():
            return {"error": "Session is no longer active"}

        return {
            "success": False,
            "error": "Screenshots are not available in headless IPython sessions",
            "hint": (
                "Interactive sessions are fully headless — no OpenGL scene exists to capture. "
                "To preview your work: (1) write the scene code to a file with write_scene, "
                "then (2) call get_scene_preview to render the final frame as a PNG."
            ),
        }

    async def end_session(self, session_id: str) -> dict[str, Any]:
        """
        End an interactive session and clean up resources.

        Args:
            session_id: The session ID to end

        Returns:
            Dict with confirmation and session summary
        """
        session = self.sessions.get(session_id)
        if not session:
            return {
                "error": f"Session not found: {session_id}",
                "available_sessions": list(self.sessions.keys()),
            }

        try:
            # Gather session statistics
            session_duration = time.time() - session.created_at
            commands_run = len(session.history)
            screenshots_taken = session.screenshot_count

            # Try graceful exit first
            if session.process.isalive():
                session.process.sendline("exit()")
                try:
                    session.process.expect(pexpect.EOF, timeout=5)
                except pexpect.TIMEOUT:
                    session.process.terminate(force=True)

            # Clean up screenshot directory
            screenshot_dir = session.screenshot_dir

            del self.sessions[session_id]

            return {
                "success": True,
                "message": f"Session {session_id} ended successfully",
                "summary": {
                    "duration_seconds": round(session_duration, 2),
                    "commands_executed": commands_run,
                    "screenshots_taken": screenshots_taken,
                    "screenshot_directory": screenshot_dir,
                },
                "hint": "Screenshots are preserved in the screenshot directory",
            }

        except Exception as e:
            # Force kill if needed
            try:
                session.process.terminate(force=True)
            except Exception:
                pass
            if session_id in self.sessions:
                del self.sessions[session_id]
            return {
                "success": True,
                "message": f"Session {session_id} forcefully terminated",
                "warning": str(e),
            }

    async def list_sessions(self) -> dict[str, Any]:
        """
        List all active sessions with detailed information.

        Returns:
            Dict with session information
        """
        sessions_info = []
        inactive_sessions = []

        for session_id, session in list(self.sessions.items()):
            if not session.process.isalive():
                inactive_sessions.append(session_id)
                continue

            age = time.time() - session.created_at
            idle_time = time.time() - session.last_activity

            sessions_info.append(
                {
                    "session_id": session_id,
                    "created_at": session.created_at,
                    "last_activity": session.last_activity,
                    "commands_run": len(session.history),
                    "screenshots_taken": session.screenshot_count,
                    "age_seconds": round(age, 1),
                    "idle_seconds": round(idle_time, 1),
                    "screenshot_directory": session.screenshot_dir,
                }
            )

        # Clean up dead sessions
        for session_id in inactive_sessions:
            del self.sessions[session_id]

        return {
            "success": True,
            "active_sessions": sessions_info,
            "count": len(sessions_info),
            "max_sessions": self.max_sessions,
            "available_slots": self.max_sessions - len(sessions_info),
            "cleaned_up": len(inactive_sessions),
        }

    async def get_session_history(
        self, session_id: str, last_n: int | None = None
    ) -> dict[str, Any]:
        """
        Get the command history for a session.

        Args:
            session_id: The session ID
            last_n: If provided, only return the last N commands

        Returns:
            Dict with command history
        """
        session = self.sessions.get(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        history = session.history
        if last_n and last_n > 0:
            history = history[-last_n:]

        return {
            "success": True,
            "session_id": session_id,
            "history": history,
            "total_commands": len(session.history),
            "returned_commands": len(history),
        }

    async def get_scene_state(self, session_id: str) -> dict[str, Any]:
        """
        Get information about the current scene state (mobjects, etc.).

        Args:
            session_id: The session ID

        Returns:
            Dict with scene state information
        """
        session = self.sessions.get(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        if not session.process.isalive():
            return {"error": "Session is no longer active"}

        try:
            # In a headless IPython session there is no live Scene object, so we
            # introspect the REPL namespace instead: list all Mobject instances.
            query_code = (
                "import manimlib.mobject.mobject as _mm; "
                "_mvars = [(k, type(v).__name__) for k, v in list(locals().items()) "
                "if isinstance(v, _mm.Mobject)]; "
                "print('MOBJECT_VARS:', len(_mvars), _mvars)"
            )
            session.process.sendline(query_code)
            session.process.expect(self.ipython_prompt, timeout=10)
            output = _ANSI_RE.sub("", session.process.before or "")

            mobject_count = 0
            mobject_vars: list = []
            m = re.search(r"MOBJECT_VARS:\s*(\d+)\s*(\[.*\])", output)
            if m:
                mobject_count = int(m.group(1))
                try:
                    import ast as _ast

                    mobject_vars = _ast.literal_eval(m.group(2))
                except Exception:
                    pass

            return {
                "success": True,
                "session_id": session_id,
                "mobject_count": mobject_count,
                "mobject_vars": mobject_vars,
                "note": "Counts Mobject instances visible in the IPython local namespace",
            }

        except Exception as e:
            return {"error": str(e), "hint": "Could not query session namespace"}

    def _clean_output(self, output: str, code: str) -> str:
        """Clean up IPython output by removing ANSI codes, echoed commands, and prompts."""
        if not output:
            return ""

        # Strip all ANSI escape sequences first
        output = _ANSI_RE.sub("", output)
        # Normalise line endings
        output = output.replace("\r\n", "\n").replace("\r", "\n")

        lines = output.split("\n")
        cleaned_lines = []
        code_lines = [cl.strip() for cl in code.strip().split("\n")]

        for line in lines:
            stripped = line.strip()
            # Skip empty lines at start
            if not cleaned_lines and not stripped:
                continue
            # Skip echoed input lines (IPython repeats what was typed)
            if stripped in code_lines:
                continue
            # Skip the %cpaste terminator echoed back
            if stripped == "--":
                continue
            # Skip partial prompt tails like "1]: some_code" that pexpect captures
            if re.match(r"^\d+\]:", stripped):
                continue
            # Skip full IPython input prompts
            if re.match(r"^In \[", stripped):
                continue
            # For Out[N]: lines, keep the value part but drop angle-bracket reprs
            m = re.match(r"^Out\[\d+\]:\s*(.*)", stripped)
            if m:
                val = m.group(1).strip()
                # Skip bare object reprs like <Circle object at 0x...>
                if re.match(r"^<.*>$", val):
                    continue
                cleaned_lines.append(val)
                continue
            # Skip bare repr output from expressions (angle-bracket objects)
            if re.match(r"^<.*>$", stripped):
                continue
            if stripped.startswith("..."):
                continue
            cleaned_lines.append(line)

        return "\n".join(cleaned_lines).strip()

    def cleanup_all(self):
        """Clean up all sessions."""
        for session_id in list(self.sessions.keys()):
            try:
                session = self.sessions[session_id]
                if session.process.isalive():
                    session.process.terminate(force=True)
            except Exception:
                pass
        self.sessions.clear()
