"""
Docker MCP Server
=================
MCP server providing Docker container and image management capabilities.
Supports listing, inspecting, starting, stopping, removing containers,
managing images, viewing logs, executing commands, and managing networks/volumes.
Connects to the Docker daemon via socket or TCP.
"""

import json
import os
from typing import Optional, List, Dict, Any

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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DOCKER_HOST = os.environ.get("DOCKER_HOST", "unix:///var/run/docker.sock")

# ---------------------------------------------------------------------------
# MCP Server
# ---------------------------------------------------------------------------

mcp = FastMCP("docker_mcp")

# ---------------------------------------------------------------------------
# Docker client
# ---------------------------------------------------------------------------


def _client() -> docker.DockerClient:
    """Get a Docker client instance."""
    return docker.DockerClient(base_url=DOCKER_HOST, timeout=30)


def _safe(func):
    """Wrapper for safe Docker API calls."""
    try:
        return func()
    except NotFound as e:
        return {"error": f"Not found: {str(e)}"}
    except APIError as e:
        return {"error": f"Docker API error: {str(e)}"}
    except DockerException as e:
        return {"error": f"Docker error: {str(e)}"}


# ---------------------------------------------------------------------------
# Container tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def list_containers(
    all: bool = True,
    filters: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    List Docker containers.

    Args:
        all: Show all containers (including stopped). Default True.
        filters: Optional filter dict (e.g., {"status": "running", "name": "my-app"})
    """
    client = _client()
    containers = client.containers.list(all=all, filters=filters or {})

    result = []
    for c in containers:
        result.append({
            "id": c.short_id,
            "name": c.name,
            "image": str(c.image.tags[0]) if c.image.tags else str(c.image.id[:12]),
            "status": c.status,
            "state": c.attrs.get("State", {}).get("Status"),
            "created": c.attrs.get("Created"),
            "ports": c.ports,
        })

    return {"containers": result, "total": len(result)}


@mcp.tool()
async def inspect_container(container_id: str) -> Dict[str, Any]:
    """
    Get detailed information about a container.

    Args:
        container_id: Container ID or name
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        attrs = c.attrs
        return {
            "id": c.short_id,
            "name": c.name,
            "image": str(c.image.tags[0]) if c.image.tags else str(c.image.id[:12]),
            "status": c.status,
            "state": attrs.get("State", {}),
            "config": {
                "env": attrs.get("Config", {}).get("Env", []),
                "cmd": attrs.get("Config", {}).get("Cmd"),
                "entrypoint": attrs.get("Config", {}).get("Entrypoint"),
                "working_dir": attrs.get("Config", {}).get("WorkingDir"),
            },
            "network": attrs.get("NetworkSettings", {}).get("Networks", {}),
            "mounts": attrs.get("Mounts", []),
            "ports": c.ports,
            "created": attrs.get("Created"),
            "restart_count": attrs.get("RestartCount", 0),
        }
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


@mcp.tool()
async def container_logs(
    container_id: str,
    tail: int = 100,
    since: Optional[str] = None,
    timestamps: bool = False,
) -> Dict[str, Any]:
    """
    Get logs from a container.

    Args:
        container_id: Container ID or name
        tail: Number of lines from the end (default 100)
        since: Show logs since timestamp (e.g., '2024-01-01T00:00:00')
        timestamps: Include timestamps in log output
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        kwargs: Dict[str, Any] = {
            "tail": tail,
            "timestamps": timestamps,
            "stdout": True,
            "stderr": True,
        }
        if since:
            kwargs["since"] = since

        logs = c.logs(**kwargs)
        log_text = logs.decode("utf-8", errors="replace")

        return {
            "container": container_id,
            "lines": tail,
            "logs": log_text,
        }
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


@mcp.tool()
async def start_container(container_id: str) -> Dict[str, Any]:
    """
    Start a stopped container.

    Args:
        container_id: Container ID or name
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        c.start()
        c.reload()
        return {"status": "started", "container": c.name, "state": c.status}
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


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

    Args:
        container_id: Container ID or name
        timeout: Seconds to wait before killing (default 10)
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        c.stop(timeout=timeout)
        c.reload()
        return {"status": "stopped", "container": c.name, "state": c.status}
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


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

    Args:
        container_id: Container ID or name
        timeout: Seconds to wait before killing (default 10)
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        c.restart(timeout=timeout)
        c.reload()
        return {"status": "restarted", "container": c.name, "state": c.status}
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


@mcp.tool()
async def remove_container(
    container_id: str,
    force: bool = False,
    v: bool = False,
) -> Dict[str, Any]:
    """
    Remove a container.

    Args:
        container_id: Container ID or name
        force: Force remove even if running
        v: Remove associated volumes
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        name = c.name
        c.remove(force=force, v=v)
        return {"status": "removed", "container": name}
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


@mcp.tool()
async def exec_in_container(
    container_id: str,
    command: str,
    workdir: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Execute a command inside a running container.

    Args:
        container_id: Container ID or name
        command: Command to execute (shell string)
        workdir: Optional working directory inside the container
    """
    client = _client()
    try:
        c = client.containers.get(container_id)
        kwargs: Dict[str, Any] = {"cmd": command, "stdout": True, "stderr": True}
        if workdir:
            kwargs["workdir"] = workdir

        exit_code, output = c.exec_run(**kwargs)
        return {
            "container": container_id,
            "command": command,
            "exit_code": exit_code,
            "output": output.decode("utf-8", errors="replace"),
        }
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}


@mcp.tool()
async def container_stats(container_id: str) -> Dict[str, Any]:
    """
    Get resource usage stats for a container.

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

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

        # Parse 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

        return {
            "container": container_id,
            "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_bytes": sum(
                v.get("rx_bytes", 0) for v in stats.get("networks", {}).values()
            ),
            "network_tx_bytes": sum(
                v.get("tx_bytes", 0) for v in stats.get("networks", {}).values()
            ),
            "pids": stats.get("pids_stats", {}).get("current", 0),
        }
    except NotFound:
        return {"error": f"Container '{container_id}' not found"}
    except (KeyError, ZeroDivisionError) as e:
        return {"error": f"Failed to parse stats: {str(e)}"}


# ---------------------------------------------------------------------------
# Image tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def list_images(
    name: Optional[str] = None,
    all: bool = False,
) -> Dict[str, Any]:
    """
    List Docker images.

    Args:
        name: Optional filter by image name
        all: Show all images including intermediate layers
    """
    client = _client()
    images = client.images.list(name=name, all=all)

    result = []
    for img in images:
        result.append({
            "id": img.short_id,
            "tags": img.tags,
            "size_mb": round(img.attrs.get("Size", 0) / (1024 * 1024), 2),
            "created": img.attrs.get("Created"),
        })

    return {"images": result, "total": len(result)}


@mcp.tool()
async def pull_image(
    image: str,
    tag: str = "latest",
) -> Dict[str, Any]:
    """
    Pull a Docker image from a registry.

    Args:
        image: Image name (e.g., 'nginx', 'python')
        tag: Image tag (default: 'latest')
    """
    client = _client()
    try:
        img = client.images.pull(image, tag=tag)
        return {
            "status": "pulled",
            "image": f"{image}:{tag}",
            "id": img.short_id,
            "size_mb": round(img.attrs.get("Size", 0) / (1024 * 1024), 2),
        }
    except APIError as e:
        return {"error": f"Failed to pull {image}:{tag}: {str(e)}"}


@mcp.tool()
async def remove_image(
    image: str,
    force: bool = False,
) -> Dict[str, Any]:
    """
    Remove a Docker image.

    Args:
        image: Image ID or name:tag
        force: Force removal
    """
    client = _client()
    try:
        client.images.remove(image, force=force)
        return {"status": "removed", "image": image}
    except NotFound:
        return {"error": f"Image '{image}' not found"}


# ---------------------------------------------------------------------------
# System tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def docker_system_info() -> Dict[str, Any]:
    """Get Docker system-wide information."""
    client = _client()
    info = client.info()
    return {
        "docker_version": info.get("ServerVersion"),
        "os": info.get("OperatingSystem"),
        "arch": info.get("Architecture"),
        "cpus": info.get("NCPU"),
        "memory_gb": round(info.get("MemTotal", 0) / (1024**3), 2),
        "containers_running": info.get("ContainersRunning"),
        "containers_stopped": info.get("ContainersStopped"),
        "containers_paused": info.get("ContainersPaused"),
        "images": info.get("Images"),
        "storage_driver": info.get("Driver"),
    }


@mcp.tool()
async def docker_disk_usage() -> Dict[str, Any]:
    """Get Docker disk usage summary."""
    client = _client()
    df = client.df()

    containers_size = sum(c.get("SizeRw", 0) for c in df.get("Containers", []))
    images_size = sum(i.get("Size", 0) for i in df.get("Images", []))
    volumes_size = sum(v.get("UsageData", {}).get("Size", 0) for v in df.get("Volumes", []))

    return {
        "containers_size_mb": round(containers_size / (1024 * 1024), 2),
        "images_size_mb": round(images_size / (1024 * 1024), 2),
        "volumes_size_mb": round(volumes_size / (1024 * 1024), 2),
        "total_mb": round((containers_size + images_size + volumes_size) / (1024 * 1024), 2),
        "images_count": len(df.get("Images", [])),
        "containers_count": len(df.get("Containers", [])),
        "volumes_count": len(df.get("Volumes", [])),
    }


# ---------------------------------------------------------------------------
# Network tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def list_networks() -> Dict[str, Any]:
    """List Docker networks."""
    client = _client()
    networks = client.networks.list()

    result = []
    for n in networks:
        result.append({
            "id": n.short_id,
            "name": n.name,
            "driver": n.attrs.get("Driver"),
            "scope": n.attrs.get("Scope"),
            "containers": len(n.attrs.get("Containers", {})),
        })

    return {"networks": result, "total": len(result)}


# ---------------------------------------------------------------------------
# Volume tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def list_volumes() -> Dict[str, Any]:
    """List Docker volumes."""
    client = _client()
    volumes = client.volumes.list()

    result = []
    for v in volumes:
        result.append({
            "name": v.name,
            "driver": v.attrs.get("Driver"),
            "mountpoint": v.attrs.get("Mountpoint"),
            "created": v.attrs.get("CreatedAt"),
        })

    return {"volumes": result, "total": len(result)}


@mcp.tool()
async def prune_system(
    containers: bool = True,
    images: bool = True,
    volumes: bool = False,
    networks: bool = True,
) -> Dict[str, Any]:
    """
    Prune unused Docker resources.

    Args:
        containers: Prune stopped containers
        images: Prune dangling images
        volumes: Prune unused volumes (CAUTION: data loss)
        networks: Prune unused networks
    """
    client = _client()
    results = {}

    if containers:
        r = client.containers.prune()
        results["containers_deleted"] = len(r.get("ContainersDeleted", []) or [])
        results["containers_space_mb"] = round(r.get("SpaceReclaimed", 0) / (1024 * 1024), 2)

    if images:
        r = client.images.prune()
        results["images_deleted"] = len(r.get("ImagesDeleted", []) or [])
        results["images_space_mb"] = round(r.get("SpaceReclaimed", 0) / (1024 * 1024), 2)

    if volumes:
        r = client.volumes.prune()
        results["volumes_deleted"] = len(r.get("VolumesDeleted", []) or [])
        results["volumes_space_mb"] = round(r.get("SpaceReclaimed", 0) / (1024 * 1024), 2)

    if networks:
        r = client.networks.prune()
        results["networks_deleted"] = len(r.get("NetworksDeleted", []) or [])

    return results