from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
import docker
import psutil
import subprocess
import os
from typing import Optional

# Configure transport security to avoid 421 errors (DNS rebinding protection)
enable_protection = os.getenv("MCP_ENABLE_DNS_PROTECTION", "false").lower() == "true"
if enable_protection:
    allowed_hosts = os.getenv("MCP_ALLOWED_HOSTS", "localhost:*,127.0.0.1:*,0.0.0.0:*").split(",")
    allowed_origins = os.getenv("MCP_ALLOWED_ORIGINS", "http://localhost:*,http://127.0.0.1:*,http://0.0.0.0:*").split(",")
    transport_security = TransportSecuritySettings(
        enable_dns_rebinding_protection=True,
        allowed_hosts=allowed_hosts,
        allowed_origins=allowed_origins,
    )
else:
    transport_security = TransportSecuritySettings(
        enable_dns_rebinding_protection=False,
    )

mcp = FastMCP("homelab", transport_security=transport_security)

DOCKER_CLIENT = docker.from_env()


@mcp.tool()
def container_status(container_name: str) -> dict:
    """Get status of a Docker container"""
    try:
        container = DOCKER_CLIENT.containers.get(container_name)
        return {
            "status": container.status,
            "running": container.status == "running",
            "image": container.image.tags[0] if container.image.tags else container.image.id,
            "ports": container.ports,
            "health": container.attrs.get("State", {}).get("Health", {}).get("Status", "none")
        }
    except docker.errors.NotFound:
        return {"error": f"Container {container_name} not found"}


@mcp.tool()
def list_containers(all: bool = False) -> list[dict]:
    """List Docker containers"""
    containers = DOCKER_CLIENT.containers.list(all=all)
    return [
        {
            "name": c.name,
            "status": c.status,
            "image": c.image.tags[0] if c.image.tags else c.image.id[:12],
            "ports": c.ports
        }
        for c in containers
    ]


@mcp.tool()
def start_container(container_name: str) -> dict:
    """Start a Docker container"""
    try:
        container = DOCKER_CLIENT.containers.get(container_name)
        container.start()
        return {"success": True, "message": f"Started {container_name}"}
    except docker.errors.NotFound:
        return {"error": f"Container {container_name} not found"}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
def stop_container(container_name: str, timeout: int = 10) -> dict:
    """Stop a Docker container"""
    try:
        container = DOCKER_CLIENT.containers.get(container_name)
        container.stop(timeout=timeout)
        return {"success": True, "message": f"Stopped {container_name}"}
    except docker.errors.NotFound:
        return {"error": f"Container {container_name} not found"}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
def restart_container(container_name: str) -> dict:
    """Restart a Docker container"""
    try:
        container = DOCKER_CLIENT.containers.get(container_name)
        container.restart()
        return {"success": True, "message": f"Restarted {container_name}"}
    except docker.errors.NotFound:
        return {"error": f"Container {container_name} not found"}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
def container_logs(container_name: str, lines: int = 100) -> str:
    """Get logs from a Docker container"""
    try:
        container = DOCKER_CLIENT.containers.get(container_name)
        return container.logs(tail=lines).decode("utf-8")
    except docker.errors.NotFound:
        return f"Container {container_name} not found"


@mcp.tool()
def system_resources() -> dict:
    """Get system resource usage"""
    return {
        "cpu_percent": psutil.cpu_percent(interval=1),
        "memory": {
            "total": psutil.virtual_memory().total // (1024 * 1024),
            "available": psutil.virtual_memory().available // (1024 * 1024),
            "percent": psutil.virtual_memory().percent
        },
        "disk": {
            "total": psutil.disk_usage("/").total // (1024 * 1024 * 1024),
            "used": psutil.disk_usage("/").used // (1024 * 1024 * 1024),
            "percent": psutil.disk_usage("/").percent
        }
    }


@mcp.tool()
def run_command(command: str, cwd: Optional[str] = None) -> dict:
    """Run a shell command (use carefully)"""
    try:
        result = subprocess.run(
            command,
            shell=True,
            cwd=cwd,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "returncode": result.returncode
        }
    except subprocess.TimeoutExpired:
        return {"error": "Command timed out after 30s"}
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
def docker_compose_action(
    compose_file: str,
    action: str,
    service: Optional[str] = None
) -> dict:
    """Run docker-compose action (up, down, restart, pull)"""
    if action not in ["up", "down", "restart", "pull"]:
        return {"error": f"Invalid action: {action}"}
    
    cmd = f"docker-compose -f {compose_file} {action}"
    if service:
        cmd += f" {service}"
    
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=120
        )
        return {
            "success": result.returncode == 0,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {"error": "Command timed out after 120s"}
    except Exception as e:
        return {"error": str(e)}


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