#!/usr/bin/env python3
"""
shell-mcp (mediabox) — run shell commands on the media box host via MCP/SSE.

Adapted from thethreemagi/shell-mcp, which was written for and built on the
arrsstack Pi 5 (arm64). Differences that matter:

  * Runs on 10.0.1.20 (amd64), NOT on arrsstack. The tool description below is
    rewritten accordingly — if it still claimed to be the Pi, every session
    would start with the wrong mental model of which host it is touching.
  * Listens on 8103 by default, not 8085, so NPM proxy host 42 can simply be
    repointed from mediabox-mcp:8103 to 10.0.1.20:8103. Same URL, same cert,
    same connector entry.
  * PORT is read from the environment instead of being hardcoded.
  * Dependencies are pinned (see requirements.txt). The original installed
    `mcp starlette uvicorn` unpinned; starlette has since gone 1.x. An
    unpinned install inside a first-boot script is a time bomb.
  * 2026-08-02: commands run in a worker thread (run_in_executor), not
    inline in the async handler. The original blocked uvicorn's event loop
    for the duration of every command, serializing ALL requests behind the
    slowest one — discovered mid-migration when a long du wedged the server.

Commands run in the host namespaces via nsenter — full root on the media box.
"""
import asyncio
import os
import shlex
import subprocess

from mcp.server import Server
from mcp.server.sse import SseServerTransport
from mcp.types import Tool, TextContent
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route
import uvicorn

BEARER_TOKEN = os.environ["BEARER_TOKEN"]
PORT = int(os.environ.get("PORT", "8103"))

# ── MCP server ──────────────────────────────────────────────────────────────

mcp = Server("shell-mcp")


@mcp.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="run_command",
            description=(
                "Run a shell command on the MEDIA BOX host (10.0.1.20, hostname "
                "'mediabox') with full root access. This is the Plex / Docker / "
                "GPU host — an amd64 machine with an Intel i7-8700K, UHD 630 "
                "Quick Sync, and an NVIDIA GTX 1070. It is NOT arrsstack; that "
                "is a separate Raspberry Pi connector. Commands run in host "
                "namespaces via nsenter, so filesystem, network, processes and "
                "systemd services are all the real host. Docker CLI available. "
                "The NAS is mounted under /mnt/nas/. Use for container "
                "management, service control, log inspection, file read/write, "
                "and general system administration."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "Shell command to execute on the media box host (bash -c)",
                    },
                    "working_directory": {
                        "type": "string",
                        "description": "Directory on the host to run the command in (optional)",
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds, default 60, max 600",
                        "default": 60,
                    },
                },
                "required": ["command"],
            },
        )
    ]


@mcp.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    if name != "run_command":
        return [TextContent(type="text", text=f"Unknown tool: {name}")]

    command = arguments["command"]
    working_dir = arguments.get("working_directory")
    timeout = min(int(arguments.get("timeout", 60)), 600)

    inner = f"cd {shlex.quote(working_dir)} && {command}" if working_dir else command

    host_cmd = (
        "nsenter --target 1 --mount --uts --ipc --net --pid -- "
        f"bash -c {shlex.quote(inner)}"
    )

    # Run the blocking subprocess in the default thread pool. Doing this inline
    # would block uvicorn's event loop and serialize every request behind the
    # slowest command — a long du froze the whole server for minutes once.
    def _run() -> subprocess.CompletedProcess:
        return subprocess.run(
            host_cmd,
            shell=True,
            executable="/bin/bash",
            capture_output=True,
            text=True,
            timeout=timeout,
        )

    try:
        result = await asyncio.get_running_loop().run_in_executor(None, _run)
        parts = []
        if result.stdout:
            parts.append(result.stdout.rstrip())
        if result.stderr:
            parts.append(f"[stderr]\n{result.stderr.rstrip()}")
        if result.returncode != 0:
            parts.append(f"[exit code: {result.returncode}]")
        return [TextContent(type="text", text="\n".join(parts) or "(no output)")]

    except subprocess.TimeoutExpired:
        return [TextContent(type="text", text=f"[timed out after {timeout}s]")]
    except Exception as e:
        return [TextContent(type="text", text=f"[error: {e}]")]


# ── SSE transport ───────────────────────────────────────────────────────────

sse = SseServerTransport("/messages/")


async def health_endpoint(request: Request):
    return JSONResponse({"status": "ok", "host": "mediabox"})


starlette_app = Starlette(
    routes=[
        Route("/health", endpoint=health_endpoint),
        Mount("/messages/", app=sse.handle_post_message),
    ],
)


async def app(scope, receive, send):
    if scope.get("type") != "http":
        await starlette_app(scope, receive, send)
        return

    path = scope.get("path", "")
    method = scope.get("method", "").upper()

    if path == "/sse":
        # Only GET establishes an SSE stream; reject everything else
        if method != "GET":
            await send({"type": "http.response.start", "status": 405,
                        "headers": [(b"content-type", b"text/plain"),
                                    (b"allow", b"GET")]})
            await send({"type": "http.response.body", "body": b"Method Not Allowed",
                        "more_body": False})
            return

        # Auth gate. Claude.ai's connector UI has no bearer-token field, so the
        # token may arrive as ?token=; the Authorization header is also accepted.
        query_string = scope.get("query_string", b"").decode()
        token = None
        for part in query_string.split("&"):
            if part.startswith("token="):
                token = part[6:]
                break

        auth_header = ""
        for name, value in scope.get("headers", []):
            if name.lower() == b"authorization":
                auth_header = value.decode()
                break

        if token != BEARER_TOKEN and auth_header != f"Bearer {BEARER_TOKEN}":
            await send({"type": "http.response.start", "status": 401,
                        "headers": [(b"content-type", b"text/plain")]})
            await send({"type": "http.response.body", "body": b"Unauthorized",
                        "more_body": False})
            return

        # Handle SSE directly — bypasses Starlette Route, no NoneType crash on close
        async with sse.connect_sse(scope, receive, send) as (read_stream, write_stream):
            await mcp.run(read_stream, write_stream, mcp.create_initialization_options())
        return

    await starlette_app(scope, receive, send)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")