import json
import subprocess
from io import StringIO
from contextlib import redirect_stdout, redirect_stderr

from fastmcp import FastMCP
from wizlib.parser import WizParser

from dyngle.command import DyngleCommand
from dyngle.model.context import Context
from dyngle.error import DyngleError


class McpCommand(DyngleCommand):
    """Run an MCP server exposing Dyngle operations as tools"""

    name = "mcp"

    @classmethod
    def add_args(cls, parser: WizParser):
        super().add_args(parser)
        parser.add_argument(
            "--host", default="127.0.0.1", help="Host to bind the server to"
        )
        parser.add_argument(
            "--port", type=int, default=8000, help="Port to bind the server to"
        )
        parser.add_argument(
            "--transport",
            default="stdio",
            choices=["stdio", "http", "sse"],
            help="Transport protocol to use",
        )

    def create_server(self) -> FastMCP:
        """Create and configure the FastMCP server with tools for each
        operation"""
        mcp = FastMCP("Dyngle MCP Server")

        # Register each operation as a tool
        for op_name, operation in self.app.dyngleverse.operations.items():
            # Create a closure to capture the operation name and instance
            def make_tool(name: str, op):
                async def tool_func(
                    data: dict = None, args: list = None
                ) -> str:
                    """Execute a Dyngle operation

                    Args:
                        data: Input data dictionary to pass to the operation
                        args: List of arguments to pass to the operation

                    Returns:
                        JSON string with stdout, stderr, success, and optional
                        exception
                    """
                    data = data or {}
                    args = args or []

                    # Create context from data
                    context = Context(data)

                    # Capture stdout and stderr
                    stdout_capture = StringIO()
                    stderr_capture = StringIO()

                    result = {
                        "stdout": "",
                        "stderr": "",
                        "success": True,
                        "exception": None,
                    }

                    try:
                        # Patch subprocess.run to capture outputs
                        original_run = subprocess.run

                        def patched_run(cmd, **kwargs):
                            # Check if the operation explicitly requested
                            # stdout capture
                            explicitly_capturing_stdout = (
                                "stdout" in kwargs
                                and kwargs["stdout"] == subprocess.PIPE
                            )

                            # Always capture both stdout and stderr
                            kwargs["stdout"] = subprocess.PIPE
                            kwargs["stderr"] = subprocess.PIPE
                            kwargs["text"] = True

                            run_result = original_run(cmd, **kwargs)

                            # Always collect both outputs for our result
                            if run_result.stdout:
                                stdout_capture.write(run_result.stdout)
                            if run_result.stderr:
                                stderr_capture.write(run_result.stderr)

                            # If the operation was explicitly capturing stdout,
                            # preserve that in the result object so the
                            # operation can use it
                            if not explicitly_capturing_stdout:
                                # Create a new result with stdout set to None
                                # so the operation doesn't try to use it
                                run_result = subprocess.CompletedProcess(
                                    args=run_result.args,
                                    returncode=run_result.returncode,
                                    stdout=run_result.stdout,
                                    stderr=run_result.stderr,
                                )

                            return run_result

                        # Temporarily replace subprocess.run
                        subprocess.run = patched_run

                        try:
                            # Execute the operation
                            op.run(context, args)
                        finally:
                            # Restore original subprocess.run
                            subprocess.run = original_run

                        result["stdout"] = stdout_capture.getvalue()
                        result["stderr"] = stderr_capture.getvalue()

                    except DyngleError as e:
                        result["success"] = False
                        result["exception"] = str(e)
                        result["stdout"] = stdout_capture.getvalue()
                        result["stderr"] = stderr_capture.getvalue()
                    except Exception as e:
                        result["success"] = False
                        result["exception"] = (
                            f"Unexpected error: {type(e).__name__}: {str(e)}"
                        )
                        result["stdout"] = stdout_capture.getvalue()
                        result["stderr"] = stderr_capture.getvalue()

                    return json.dumps(result)

                # Set the function name for the tool
                tool_func.__name__ = name
                return tool_func

            # Register the tool
            tool = make_tool(op_name, operation)
            mcp.tool(tool)

        return mcp

    @DyngleCommand.wrap
    def execute(self):
        """Start the MCP server"""
        server = self.create_server()

        # Build run arguments based on transport
        run_kwargs = {"transport": self.transport}
        
        # Only add host and port for non-stdio transports
        if self.transport != "stdio":
            run_kwargs["host"] = self.host
            run_kwargs["port"] = self.port

        # Run the server with the specified transport
        server.run(**run_kwargs)

        return f"MCP server started on {self.transport}"
