#!/usr/bin/env python3
"""Compact MCP facade for the Athena host operator."""

from __future__ import annotations

import json
import os
import socket
import sys
from typing import Any


VERSION = "3.0.0"
SOCKET_PATH = os.environ.get("ATHENA_OPERATOR_SOCKET", "/operator/operator.sock")

if hasattr(sys.stdin, "reconfigure"):
    sys.stdin.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")


TOOLS = [
    {
        "name": "athena_operator_inspect",
        "description": "START HERE once for Athena work. Inspect the requested live area without changing it.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "subject": {"type": "string", "enum": ["overview", "git_status", "containers", "models", "jobs"], "default": "overview"},
            },
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_operator_search_source",
        "description": "Find the exact source path and line before reading or editing it. Returns at most 20 matches.",
        "inputSchema": {
            "type": "object",
            "properties": {"query": {"type": "string", "minLength": 1, "maxLength": 120}},
            "required": ["query"],
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_operator_read_source",
        "description": (
            "Read only the needed section of one known source file. Default 80 and maximum "
            "200 lines. A not-found result must not be retried by guessing more paths."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "pattern": "^[A-Za-z0-9_.+/-]{1,240}$"},
                "start_line": {"type": "integer", "minimum": 1, "maximum": 1000000, "default": 1},
                "line_count": {"type": "integer", "minimum": 1, "maximum": 200, "default": 80},
            },
            "required": ["path"],
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_operator_terminal",
        "description": (
            "General bounded terminal on Athena for focused Docker, Git, file, HTTP, model "
            "and test work when no structured operation fits. Limit noisy output yourself. "
            "Power control and Athena reachability changes are blocked."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "command": {"type": "string", "minLength": 1, "maxLength": 8000},
                "cwd": {"type": "string", "maxLength": 500, "default": "/opt/mike-ai/stack"},
                "timeout_seconds": {"type": "integer", "minimum": 1, "maximum": 3600, "default": 300},
                "max_output_chars": {"type": "integer", "minimum": 1000, "maximum": 30000, "default": 8000},
            },
            "required": ["command"],
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_operator_change",
        "description": (
            "Apply one durable change that the user has requested. Operations: patch_update, "
            "file_update, mcp_release, run_checks, compose_deploy, container_action, "
            "openwebui_sync, git_publish, model_download, benchmark, recovery. Use a small "
            "patch for edits and file_update only for a new or intentionally replaced file. "
            "mcp_release can perform the complete MCP build/test/deploy/publish workflow."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["patch_update", "file_update", "mcp_release", "run_checks", "compose_deploy", "container_action", "openwebui_sync", "git_publish", "model_download", "benchmark", "recovery"],
                },
                "payload": {"type": "object"},
            },
            "required": ["operation", "payload"],
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_operator_job",
        "description": "Poll one asynchronous change returned by athena_operator_change. Do not start a duplicate job.",
        "inputSchema": {
            "type": "object",
            "properties": {"job_id": {"type": "string", "pattern": "^[0-9a-f]{32}$"}},
            "required": ["job_id"],
            "additionalProperties": False,
        },
    },
]


def request(action: str, arguments: dict[str, Any]) -> dict[str, Any]:
    payload = json.dumps({"action": action, "arguments": arguments}, ensure_ascii=False, separators=(",", ":")).encode() + b"\n"
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
        client.settimeout(3700 if action == "change" else 30)
        client.connect(SOCKET_PATH)
        client.sendall(payload)
        chunks = bytearray()
        while not chunks.endswith(b"\n"):
            part = client.recv(65536)
            if not part:
                break
            chunks.extend(part)
            if len(chunks) > 2_000_000:
                raise RuntimeError("operator response exceeds limit")
    response = json.loads(chunks)
    if not response.get("ok"):
        raise RuntimeError(response.get("error", "operator request failed"))
    return response["result"]


def call(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
    mapping = {
        "athena_operator_inspect": "inspect",
        "athena_operator_search_source": "search_source",
        "athena_operator_read_source": "read_source",
        "athena_operator_terminal": "terminal",
        "athena_operator_change": "change",
        "athena_operator_job": "job",
    }
    action = mapping.get(name)
    if action is None:
        return {"ok": False, "error": "unknown tool", "retry": False}
    return request(action, arguments)


def emit(request_id: Any, result: Any = None, error: dict[str, Any] | None = None) -> None:
    message = {"jsonrpc": "2.0", "id": request_id}
    message["error" if error else "result"] = error or result
    sys.stdout.write(json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n")
    sys.stdout.flush()


def handle(message: dict[str, Any]) -> None:
    method, request_id = message.get("method"), message.get("id")
    if method == "initialize":
        emit(request_id, {
            "protocolVersion": message.get("params", {}).get("protocolVersion", "2024-11-05"),
            "capabilities": {"tools": {"listChanged": False}},
            "serverInfo": {"name": "mike-ai-athena-operator", "version": VERSION},
        })
    elif method == "tools/list":
        emit(request_id, {"tools": TOOLS})
    elif method == "tools/call":
        params = message.get("params") or {}
        try:
            value = call(str(params.get("name", "")), params.get("arguments") or {})
            emit(request_id, {
                "content": [{"type": "text", "text": json.dumps(value, ensure_ascii=False, separators=(",", ":"))}],
                "structuredContent": value,
                "isError": False,
            })
        except Exception as exc:
            emit(request_id, {"content": [{"type": "text", "text": f"ERROR: {exc}"}], "isError": True})
    elif request_id is not None:
        emit(request_id, error={"code": -32601, "message": "method not found"})


def main() -> None:
    for line in sys.stdin:
        try:
            if line.strip():
                handle(json.loads(line))
        except Exception as exc:
            sys.stderr.write(f"MCP input error: {exc}\n")
            sys.stderr.flush()


if __name__ == "__main__":
    main()