import sys
import io
import traceback
import base64
import os
import tempfile
import signal
from contextlib import redirect_stdout, redirect_stderr
from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationCapabilities
from mcp.server.stdio import stdio_server
import mcp.types as types

server = Server("python-runner")

def run_code(code: str) -> dict:
    """Execute Python code and return stdout, stderr, and images."""
    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    images = []
    tmpdir = tempfile.mkdtemp()
    original_dir = os.getcwd()
    os.chdir(tmpdir)

    # Prepare a globals dict with matplotlib import
    exec_globals = {"__builtins__": __builtins__}
    try:
        import matplotlib
        matplotlib.use('Agg')  # non-interactive backend
        exec_globals["plt"] = __import__("matplotlib.pyplot")
    except ImportError:
        pass

    def save_fig():
        """Helper to save current matplotlib figure."""
        try:
            import matplotlib.pyplot as plt
            fig = plt.gcf()
            if fig.get_axes():
                buf = io.BytesIO()
                fig.savefig(buf, format='png', bbox_inches='tight')
                buf.seek(0)
                img_b64 = base64.b64encode(buf.read()).decode()
                images.append(img_b64)
                plt.close(fig)
        except Exception:
            pass

    # Timeout handler
    def timeout_handler(signum, frame):
        raise TimeoutError("Code execution timed out (10s)")

    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(10)

    try:
        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
            exec(code, exec_globals)
        save_fig()
        output = stdout_capture.getvalue()
        errors = stderr_capture.getvalue()
        signal.alarm(0)
    except Exception as e:
        save_fig()
        output = stdout_capture.getvalue()
        errors = stderr_capture.getvalue() + "\n" + traceback.format_exc()
        signal.alarm(0)
    finally:
        os.chdir(original_dir)
        # Clean up temp dir
        for f in os.listdir(tmpdir):
            os.unlink(os.path.join(tmpdir, f))
        os.rmdir(tmpdir)

    return {
        "stdout": output,
        "stderr": errors,
        "images": images
    }

@server.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="execute_code",
            description="Execute Python code and return output and any matplotlib figures as base64 PNG images.",
            inputSchema={
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to run"}
                },
                "required": ["code"]
            }
        )
    ]

@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent | types.ImageContent]:
    if name == "execute_code":
        result = run_code(arguments["code"])
        contents = []
        if result["stdout"]:
            contents.append(types.TextContent(type="text", text=f"STDOUT:\n{result['stdout']}"))
        if result["stderr"]:
            contents.append(types.TextContent(type="text", text=f"STDERR:\n{result['stderr']}"))
        for img in result["images"]:
            contents.append(types.ImageContent(type="image", data=img, mimeType="image/png"))
        if not contents:
            contents.append(types.TextContent(type="text", text="(no output)"))
        return contents
    raise ValueError(f"Unknown tool: {name}")

async def main():
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationCapabilities(
                sampling=None,
                experimental=None,
                roots=None
            ),
            notification_options=NotificationOptions(),
        )

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())