import asyncio
import json
import logging
import platform
import sys
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Union

from dotenv import load_dotenv
from mcp.server import FastMCP
from mcp.types import TextContent
from pydantic import Field, BaseModel
from pydantic.fields import FieldInfo

from aworld.logs.util import Color
from .base import ActionResponse, _color_log

load_dotenv()
#workspace = Path("/tmp/project")
workspace = Path.home()

command_history: list[dict] = []
max_history_size = 50

# Define dangerous commands for safety
dangerous_commands = [
    "rm -rf /",
    "mkfs",
    "dd if=",
    ":(){ :|:& };:",  # Unix
    "del /f /s /q",
    "format",
    "diskpart",  # Windows
    "sudo rm",
    "sudo dd",
    "sudo mkfs",  # Sudo variants
]

# Get current platform info
platform_info = {
    "system": platform.system(),
    "platform": platform.platform(),
    "architecture": platform.architecture()[0],
}


class CommandResult(BaseModel):
    """Individual command execution result with structured data."""

    command: str
    success: bool
    stdout: str
    stderr: str
    return_code: int
    duration: str
    timestamp: str


class TerminalMetadata(BaseModel):
    """Metadata for terminal operation results."""

    command: str
    platform: str
    working_directory: str
    timeout_seconds: int
    execution_time: float | None = None
    return_code: int | None = None
    safety_check_passed: bool = True
    error_type: str | None = None
    history_count: int | None = None
    output_data: str | None = None


mcp = FastMCP("terminal-server", instructions="""
Terminal MCP Server

This module provides MCP server functionality for executing terminal commands safely.
It supports command execution with timeout controls and returns LLM-friendly formatted results.

Key features:
- Execute terminal commands with configurable timeouts
- Cross-platform command execution support
- Command history tracking and retrieval
- Safety checks for dangerous commands
- LLM-optimized output formatting

Main functions:
- mcp_execute_command: Execute terminal commands with safety checks
- mcp_get_command_history: Retrieve recent command execution history
- mcp_get_terminal_capabilities: Get terminal service capabilities
""")


@mcp.tool(description="""
Execute a terminal command with safety checks and timeout controls.

        This tool provides secure command execution with:
        - Cross-platform compatibility (Windows, macOS, Linux)
        - Configurable timeout controls
        - Safety checks for dangerous commands
        - LLM-optimized result formatting
        - Command history tracking

        Specialized Feature:
        - Execute Python code and output the result to stdout
            - Example (Directly execute simple Python code): `python -c "nums = [1, 2, 3, 4]\nsum_of_nums = sum(nums)\nprint(f'{sum_of_nums=}')"`
            - Example (Execute code from a file): `python my_script.py`
""")
async def execute_command(
        command: str = Field(description="Terminal command to execute"),
        timeout: int = Field(default=30, description="Command timeout in seconds (default: 30)"),
        output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> Union[str, TextContent]:
    if isinstance(command, FieldInfo):
        command = command.default
    if isinstance(timeout, FieldInfo):
        timeout = timeout.default
    if isinstance(output_format, FieldInfo):
        output_format = output_format.default

    output_data=""

    try:
        # Safety check
        is_safe, safety_reason = _check_command_safety(command)
        if not is_safe:
            action_response = ActionResponse(
                success=False,
                message=f"Command rejected for security reasons: {safety_reason}",
                metadata=TerminalMetadata(
                    command=command,
                    platform=platform_info["system"],
                    working_directory=str(workspace),
                    timeout_seconds=timeout,
                    safety_check_passed=False,
                    error_type="security_violation",
                    output_data=safety_reason,
                ).model_dump(),
            )
            return TextContent(
                type="text",
                text=json.dumps(action_response.model_dump()),  # Empty string instead of None
                **{"metadata": {}}  # Pass as additional fields
            )

        _color_log(f"🔧 Executing command: {command}", Color.cyan)

        # Execute command
        start_time = time.time()
        result = await _execute_command_async(command, timeout)
        execution_time = time.time() - start_time

        # Format output
        formatted_output = _format_command_output(result, output_format)

        outputs = []
        if result.stderr:
            outputs.append(result.stderr)
        if result.stdout:
            outputs.append(result.stdout)
        output_data = "\n".join(outputs)


        # Create metadata
        metadata = TerminalMetadata(
            command=command,
            platform=platform_info["system"],
            working_directory=str(workspace),
            timeout_seconds=timeout,
            execution_time=execution_time,
            return_code=result.return_code,
            safety_check_passed=True,
            output_data=output_data,
        )

        if result.success:
            _color_log("✅ Command completed successfully", Color.green)
        else:
            _color_log(f"❌ Command failed with return code {result.return_code}", Color.red)
            metadata.error_type = "execution_failure"

        action_response = ActionResponse(success=result.success, message=formatted_output,
                                         metadata=metadata.model_dump())
        return TextContent(
            type="text",
            text=json.dumps(action_response.model_dump()),  # Empty string instead of None
            **{"metadata": {}}  # Pass as additional fields
        )

    except Exception as e:
        error_msg = f"Failed to execute command: {str(e)}"
        logging.error(f"Command execution error: {traceback.format_exc()}")

        action_response = ActionResponse(
            success=False,
            message=error_msg,
            metadata=TerminalMetadata(
                command=command,
                platform=platform_info["system"],
                working_directory=str(workspace),
                timeout_seconds=timeout,
                safety_check_passed=True,
                error_type="internal_error",
            ).model_dump(),
        )
        return TextContent(
            type="text",
            text=json.dumps(action_response.model_dump()),  # Empty string instead of None
            **{"metadata": {}}  # Pass as additional fields
        )


@mcp.tool(description="""
Retrieve recent command execution history.
""")
async def get_command_history(
        count: int = Field(default=10, description="Number of recent commands to return (default: 10)"),
        output_format: str = Field(default="markdown", description="Output format: 'markdown', 'json', or 'text'"),
) -> Union[str, TextContent]:
    if isinstance(count, FieldInfo):
        count = count.default
    if isinstance(output_format, FieldInfo):
        output_format = output_format.default

    try:
        # Get recent history
        recent_history = command_history[-count:] if command_history else []

        if not recent_history:
            message = "No command history available."
        else:
            if output_format == "json":
                message = json.dumps(recent_history, indent=2)
            elif output_format == "text":
                history_lines = []
                for i, entry in enumerate(recent_history, 1):
                    status = "SUCCESS" if entry["success"] else "FAILED"
                    history_lines.append(
                        f"{i}. [{entry['timestamp']}] {entry['command']} - {status}"
                        f" ({entry.get('duration', 'N/A')})"
                    )
                message = "\n".join(history_lines)
            else:  # markdown
                history_lines = ["# Command History", f"Showing {len(recent_history)} recent commands:\n"]

                for i, entry in enumerate(recent_history, 1):
                    status_emoji = "✅" if entry["success"] else "❌"
                    history_lines.extend(
                        [
                            f"## {i}. {status_emoji} `{entry['command']}`",
                            f"- **Timestamp:** {entry['timestamp']}",
                            f"- **Duration:** {entry.get('duration', 'N/A')}",
                            "",
                        ]
                    )

                message = "\n".join(history_lines)

        metadata = TerminalMetadata(
            command="get_command_history",
            platform=platform_info["system"],
            working_directory=str(workspace),
            timeout_seconds=0,
            history_count=len(recent_history),
        )

        action_response = ActionResponse(success=True, message=message, metadata=metadata.model_dump())
        return TextContent(
            type="text",
            text=json.dumps(action_response.model_dump()),  # Empty string instead of None
            **{"metadata": {}}  # Pass as additional fields
        )

    except Exception as e:
        error_msg = f"Failed to retrieve command history: {str(e)}"
        logging.error(f"History retrieval error: {traceback.format_exc()}")

        action_response = ActionResponse(
            success=False,
            message=error_msg,
            metadata=TerminalMetadata(
                command="get_command_history",
                platform=platform_info["system"],
                working_directory=str(workspace),
                timeout_seconds=0,
                error_type="internal_error",
            ).model_dump(),
        )
        return TextContent(
            type="text",
            text=json.dumps(action_response.model_dump()),  # Empty string instead of None
            **{"metadata": {}}  # Pass as additional fields
        )


@mcp.tool(description="""
Get information about terminal service capabilities and configuration.
""")
async def get_terminal_capabilities(
) -> Union[str, TextContent]:
    capabilities = {
        "platform_info": platform_info,
        "supported_features": [
            "Cross-platform command execution",
            "Configurable timeout controls",
            "Command history tracking",
            "Safety checks for dangerous commands",
            "Multiple output formats (markdown, json, text)",
            "LLM-optimized result formatting",
            "Async command execution",
        ],
        "supported_formats": ["markdown", "json", "text"],
        "configuration": {
            "max_history_size": max_history_size,
            "current_history_count": len(command_history),
            "working_directory": str(workspace),
            "dangerous_commands_count": len(dangerous_commands),
        },
        "safety_features": [
            "Dangerous command detection",
            "Timeout controls",
            "Error handling and logging",
            "Command validation",
        ],
    }

    formatted_info = f"""# Terminal Service Capabilities

            ## Platform Information
            - **System:** {platform_info["system"]}
            - **Platform:** {platform_info["platform"]}
            - **Architecture:** {platform_info["architecture"]}

            ## Features
            {chr(10).join(f"- {feature}" for feature in capabilities["supported_features"])}

            ## Supported Output Formats
            {chr(10).join(f"- {fmt}" for fmt in capabilities["supported_formats"])}

            ## Current Configuration
            - **Max History Size:** {capabilities["configuration"]["max_history_size"]}
            - **Current History Count:** {capabilities["configuration"]["current_history_count"]}
            - **Working Directory:** {capabilities["configuration"]["working_directory"]}
            - **Dangerous Commands Monitored:** {capabilities["configuration"]["dangerous_commands_count"]}

            ## Safety Features
            {chr(10).join(f"- {feature}" for feature in capabilities["safety_features"])}
            """

    action_response = ActionResponse(success=True, message=formatted_info, metadata=capabilities)
    return TextContent(
        type="text",
        text=json.dumps(action_response.model_dump()),  # Empty string instead of None
        **{"metadata": {}}  # Pass as additional fields
    )


def _check_command_safety(command: str) -> tuple[bool, str | None]:
    """Check if command is safe to execute.

    Args:
        command: Command string to check

    Returns:
        Tuple of (is_safe, reason_if_unsafe)
    """
    command_lower = command.lower().strip()

    for dangerous_cmd in dangerous_commands:
        if dangerous_cmd.lower() in command_lower:
            return False, f"Command contains dangerous pattern: {dangerous_cmd}"

    return True, None


def _format_command_output(result: CommandResult, output_format: str = "markdown") -> str:
    """Format command execution results for LLM consumption.

    Args:
        result: Command execution result
        output_format: Format type ('markdown', 'json', 'text')

    Returns:
        Formatted string suitable for LLM consumption
    """
    if output_format == "json":
        return json.dumps(result.model_dump(), indent=2)

    elif output_format == "text":
        output_parts = [
            f"Command: {result.command}",
            f"Status: {'SUCCESS' if result.success else 'FAILED'}",
            f"Duration: {result.duration}",
            f"Return Code: {result.return_code}",
        ]

        if result.stdout:
            output_parts.extend(["\nOutput:", result.stdout])

        if result.stderr:
            output_parts.extend(["\nErrors/Warnings:", result.stderr])

        return "\n".join(output_parts)

    else:  # markdown (default)
        status_emoji = "✅" if result.success else "❌"

        output_parts = [
            f"# Terminal Command Execution {status_emoji}",
            f"**Command:** `{result.command}`",
            f"**Status:** {'SUCCESS' if result.success else 'FAILED'}",
            f"**Duration:** {result.duration}",
            f"**Return Code:** {result.return_code}",
            f"**Timestamp:** {result.timestamp}",
        ]

        if result.stdout:
            output_parts.extend(["\n## Output", "```", result.stdout.strip(), "```"])

        if result.stderr:
            output_parts.extend(["\n## Errors/Warnings", "```", result.stderr.strip(), "```"])

        return "\n".join(output_parts)


async def _execute_command_async(command: str, timeout: int) -> CommandResult:
    """Execute command asynchronously with timeout.

    Args:
        command: Command to execute
        timeout: Timeout in seconds

    Returns:
        CommandResult with execution details
    """
    start_time = datetime.now()

    try:
        # Create appropriate subprocess for platform
        if platform_info["system"] == "Windows":
            process = await asyncio.create_subprocess_shell(
                command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, shell=True
            )
        else:
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                shell=True,
                executable="/bin/bash",
            )

        try:
            stdout, stderr = await asyncio.wait_for(process.communicate(), timeout)
            stdout = stdout.decode("utf-8", errors="replace")
            stderr = stderr.decode("utf-8", errors="replace")
            return_code = process.returncode

        except asyncio.TimeoutError:
            try:
                process.kill()
            except Exception:
                pass

            duration = str(datetime.now() - start_time)
            return CommandResult(
                command=command,
                success=False,
                stdout="",
                stderr=f"Command timed out after {timeout} seconds",
                return_code=-1,
                duration=duration,
                timestamp=start_time.isoformat(),
            )

        duration = str(datetime.now() - start_time)
        result = CommandResult(
            command=command,
            success=return_code == 0,
            stdout=stdout,
            stderr=stderr,
            return_code=return_code,
            duration=duration,
            timestamp=start_time.isoformat(),
        )

        # Add to history
        command_history.append(
            {
                "timestamp": start_time.isoformat(),
                "command": command,
                "success": return_code == 0,
                "duration": duration,
            }
        )

        # Maintain history size limit
        if len(command_history) > max_history_size:
            command_history.pop(0)

        return result

    except Exception as e:
        duration = str(datetime.now() - start_time)
        return CommandResult(
            command=command,
            success=False,
            stdout="",
            stderr=f"Error executing command: {str(e)}",
            return_code=-1,
            duration=duration,
            timestamp=start_time.isoformat(),
        )


def main():
    from dotenv import load_dotenv

    load_dotenv(override=True)

    print("Starting  MCP terminal-server...", file=sys.stderr)
    mcp.run(transport="stdio")


# Make the module callable
def __call__():
    """
    Make the module callable for uvx.
    This function is called when the module is executed directly.
    """
    main()


sys.modules[__name__].__call__ = __call__

if __name__ == "__main__":
    main()
