import docker
from docker.errors import APIError, NotFound
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("Docker MCP", instructions="MCP server for managing Docker containers.")

_client: docker.DockerClient | None = None


def get_client() -> docker.DockerClient:
    global _client
    if _client is None:
        _client = docker.from_env()
    return _client


def _container_summary(c) -> dict:
    return {
        "id": c.short_id,
        "name": c.name,
        "image": str(c.image.tags[0] if c.image.tags else c.image.short_id),
        "status": c.status,
        "ports": c.ports,
    }


@mcp.tool()
def list_containers(all: bool = False) -> list[dict]:
    """List Docker containers.

    Args:
        all: If True, include stopped containers. Default is only running.
    """
    client = get_client()
    containers = client.containers.list(all=all)
    return [_container_summary(c) for c in containers]


@mcp.tool()
def inspect_container(container_id: str) -> dict:
    """Get detailed information about a container.

    Args:
        container_id: Container ID or name.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
    except NotFound:
        return {"error": f"Container '{container_id}' not found."}

    return {
        "id": c.id,
        "short_id": c.short_id,
        "name": c.name,
        "status": c.status,
        "image": str(c.image.tags[0] if c.image.tags else c.image.short_id),
        "labels": c.labels,
        "ports": c.ports,
        "created": str(c.attrs.get("Created", "")),
        "platform": c.attrs.get("Platform", ""),
        "state": c.attrs.get("State", {}),
        "network_settings": {
            k: v.get("IPAddress", "") if isinstance(v, dict) else v
            for k, v in (c.attrs.get("NetworkSettings", {}).get("Networks", {})).items()
        },
        "mounts": [
            {
                "source": m.get("Source", ""),
                "destination": m.get("Destination", ""),
                "mode": m.get("Mode", ""),
            }
            for m in c.attrs.get("Mounts", [])
        ],
        "env": c.attrs.get("Config", {}).get("Env", []),
        "cmd": c.attrs.get("Config", {}).get("Cmd", []),
    }


@mcp.tool()
def start_container(container_id: str) -> str:
    """Start a stopped container.

    Args:
        container_id: Container ID or name.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        c.start()
        return f"Container '{c.name}' started."
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to start container: {e.explanation}"


@mcp.tool()
def stop_container(container_id: str, timeout: int = 10) -> str:
    """Stop a running container.

    Args:
        container_id: Container ID or name.
        timeout: Seconds to wait before killing. Default 10.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        c.stop(timeout=timeout)
        return f"Container '{c.name}' stopped."
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to stop container: {e.explanation}"


@mcp.tool()
def restart_container(container_id: str, timeout: int = 10) -> str:
    """Restart a container.

    Args:
        container_id: Container ID or name.
        timeout: Seconds to wait before killing during restart. Default 10.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        c.restart(timeout=timeout)
        return f"Container '{c.name}' restarted."
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to restart container: {e.explanation}"


@mcp.tool()
def remove_container(container_id: str, force: bool = False) -> str:
    """Remove a container.

    Args:
        container_id: Container ID or name.
        force: Force remove a running container. Default False.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        name = c.name
        c.remove(force=force)
        return f"Container '{name}' removed."
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to remove container: {e.explanation}"


@mcp.tool()
def container_logs(container_id: str, tail: int = 100, timestamps: bool = False) -> str:
    """Get logs from a container.

    Args:
        container_id: Container ID or name.
        tail: Number of lines from the end. Default 100.
        timestamps: Include timestamps in output. Default False.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        logs = c.logs(tail=tail, timestamps=timestamps)
        return logs.decode("utf-8", errors="replace")
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to get logs: {e.explanation}"


@mcp.tool()
def exec_in_container(container_id: str, command: str) -> str:
    """Execute a command inside a running container.

    Args:
        container_id: Container ID or name.
        command: The command to run (e.g. "ls -la /app").
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        exit_code, output = c.exec_run(command)
        decoded = output.decode("utf-8", errors="replace")
        return f"exit_code={exit_code}\n{decoded}"
    except NotFound:
        return f"Container '{container_id}' not found."
    except APIError as e:
        return f"Failed to exec command: {e.explanation}"


@mcp.tool()
def run_container(
    image: str,
    name: str | None = None,
    command: str | None = None,
    detach: bool = True,
    ports: dict[str, int] | None = None,
    environment: dict[str, str] | None = None,
    volumes: dict[str, dict[str, str]] | None = None,
) -> str:
    """Run a new container from an image.

    Args:
        image: Docker image to run (e.g. "nginx:latest").
        name: Optional container name.
        command: Optional command to run.
        detach: Run in background. Default True.
        ports: Port mapping, e.g. {"80/tcp": 8080}.
        environment: Environment variables, e.g. {"KEY": "value"}.
        volumes: Volume mounts, e.g. {"/host/path": {"bind": "/container/path", "mode": "rw"}}.
    """
    client = get_client()
    try:
        c = client.containers.run(
            image,
            command=command,
            name=name,
            detach=detach,
            ports=ports,
            environment=environment,
            volumes=volumes,
        )
        if detach:
            return f"Container '{c.name}' ({c.short_id}) started from image '{image}'."
        return c.decode("utf-8", errors="replace") if isinstance(c, bytes) else str(c)
    except APIError as e:
        return f"Failed to run container: {e.explanation}"


@mcp.tool()
def container_stats(container_id: str) -> dict:
    """Get resource usage statistics for a container (CPU, memory, network I/O).

    Args:
        container_id: Container ID or name.
    """
    client = get_client()
    try:
        c = client.containers.get(container_id)
        stats = c.stats(stream=False)

        # CPU
        cpu_delta = (
            stats["cpu_stats"]["cpu_usage"]["total_usage"]
            - stats["precpu_stats"]["cpu_usage"]["total_usage"]
        )
        system_delta = stats["cpu_stats"].get("system_cpu_usage", 0) - stats[
            "precpu_stats"
        ].get("system_cpu_usage", 0)
        num_cpus = stats["cpu_stats"].get("online_cpus", 1)
        cpu_percent = (
            (cpu_delta / system_delta) * num_cpus * 100.0 if system_delta > 0 else 0.0
        )

        # Memory
        mem_usage = stats["memory_stats"].get("usage", 0)
        mem_limit = stats["memory_stats"].get("limit", 1)
        mem_percent = (mem_usage / mem_limit) * 100.0

        # Network
        net_rx = 0
        net_tx = 0
        for iface_stats in stats.get("networks", {}).values():
            net_rx += iface_stats.get("rx_bytes", 0)
            net_tx += iface_stats.get("tx_bytes", 0)

        return {
            "container": c.name,
            "cpu_percent": round(cpu_percent, 2),
            "memory_usage_mb": round(mem_usage / (1024 * 1024), 2),
            "memory_limit_mb": round(mem_limit / (1024 * 1024), 2),
            "memory_percent": round(mem_percent, 2),
            "network_rx_mb": round(net_rx / (1024 * 1024), 2),
            "network_tx_mb": round(net_tx / (1024 * 1024), 2),
        }
    except NotFound:
        return {"error": f"Container '{container_id}' not found."}
    except APIError as e:
        return {"error": f"Failed to get stats: {e.explanation}"}
    except KeyError, ZeroDivisionError:
        return {"error": "Stats unavailable for this container."}


if __name__ == "__main__":
    mcp.run()