import asyncio
import sys
from io import StringIO
import traceback
import builtins

from mcp.server.models import InitializationOptions
import mcp.types as types
from mcp.server import NotificationOptions, Server
from pydantic import AnyUrl
import mcp.server.stdio

server = Server("python_interpreter")

# Global state for REPL sessions
class ReplSession:
    def __init__(self):
        self.locals = {"__builtins__": builtins}
        self.history = []
        
    def execute(self, code: str) -> tuple[str, str]:
        stdout_capture = StringIO()
        stderr_capture = StringIO()
        
        old_stdout, old_stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = stdout_capture, stderr_capture
        
        try:
            # Try to evaluate as expression first
            try:
                result = eval(code, self.locals)
                if result is not None:
                    print(repr(result))
            except SyntaxError:
                # If not an expression, execute as statement
                exec(code, self.locals)
                
        except Exception:
            traceback.print_exc()
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr
            
        return stdout_capture.getvalue(), stderr_capture.getvalue()

sessions: dict[str, ReplSession] = {}

@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="python_repl",
            description="Interactive Python REPL that maintains session state. NOTE THE USER DOES NOT SEE THE STDOUT/STDERR OUTPUT, MAKE SURE TO FORMAT/SUMMARIZEc IT APPROPRIATELY IN YOUR RESPONSE TO THE USER",
            inputSchema={
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"},
                    "session_id": {"type": "string", "description": "Session identifier"},
                },
                "required": ["code", "session_id"],
            },
        )
    ]

@server.call_tool()
async def handle_call_tool(
    name: str, 
    arguments: dict | None
) -> list[types.TextContent]:
    if name != "python_repl":
        raise ValueError(f"Unknown tool: {name}")

    if not arguments or "code" not in arguments or "session_id" not in arguments:
        raise ValueError("Missing code or session_id argument")

    code = arguments["code"]
    session_id = arguments["session_id"]
    
    if session_id not in sessions:
        sessions[session_id] = ReplSession()
    
    session = sessions[session_id]
    stdout, stderr = session.execute(code)
    session.history.append({"code": code, "stdout": stdout, "stderr": stderr})
    
    return [types.TextContent(type="text", text=f"NOTE THE USER DOES NOT SEE THIS OUTPUT, MAKE SURE TO FORMAT IT APPROPRIATELY IN YOUR RESPONSE TO THE USER\n\n{stdout}{stderr}")]

@server.list_resources()
async def handle_list_resources() -> list[types.Resource]:
    return [
        types.Resource(
            uri=AnyUrl(f"repl://{session_id}/history"),
            name=f"REPL Session {session_id}",
            description="REPL session history",
            mimeType="text/plain",
        )
        for session_id in sessions
    ]

@server.read_resource()
async def handle_read_resource(uri: AnyUrl) -> str:
    if uri.scheme != "repl":
        raise ValueError(f"Unsupported URI scheme: {uri.scheme}")

    session_id = uri.host
    if session_id not in sessions:
        raise ValueError(f"Session not found: {session_id}")
        
    history = sessions[session_id].history
    return "\n\n".join(
        f"In [{i}]: {entry['code']}\n"
        f"Out[{i}]:\n{entry['stdout']}{entry['stderr']}"
        for i, entry in enumerate(history)
    )

async def main():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="python_interpreter",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )

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