#!/usr/bin/env python3
"""Small read-only Athena knowledge MCP.

The normal entry point is ATHENA.md.  Large historical documentation remains
available through bounded search/read tools but is never loaded automatically.
"""

from __future__ import annotations

import json
import os
import re
import sys
from pathlib import Path
from typing import Any


VERSION = "2.0.0"
REPO_ROOT = Path(os.environ.get("ATHENA_REPO_ROOT", "/knowledge/repo")).resolve()
RUNTIME_FILE = Path(os.environ.get("ATHENA_RUNTIME_FILE", "/runtime/runtime.json"))
MAX_OVERVIEW_CHARS = 14_000
MAX_READ_LINES = 160
MAX_SEARCH_RESULTS = 8
ALLOWED_SUFFIXES = {".md", ".json", ".yaml", ".yml", ".txt"}
BLOCKED_PARTS = {".git", "secrets", "private", "credentials"}

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_get_overview",
        "description": (
            "START HERE for Athena architecture or administration. Returns the compact, "
            "authoritative ATHENA.md. Do not read additional platform documents unless a "
            "specific unresolved question remains."
        ),
        "inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
    },
    {
        "name": "athena_get_current_state",
        "description": "Return the compact generated runtime snapshot: active profile, containers, GPUs, source commit and recovery status.",
        "inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
    },
    {
        "name": "athena_get_external_services",
        "description": "List known services outside Athena so an existing Unraid or home-network backend is reused instead of duplicated.",
        "inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
    },
    {
        "name": "athena_search_reference",
        "description": (
            "Search ATHENA.md and documentation for one concrete term. Returns at most eight "
            "short excerpts. Use only when ATHENA.md did not answer the question."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"query": {"type": "string", "minLength": 2, "maxLength": 120}},
            "required": ["query"],
            "additionalProperties": False,
        },
    },
    {
        "name": "athena_read_reference",
        "description": (
            "Read a bounded line range from one known documentation file. Missing paths are "
            "reported as a normal not-found result and must not be retried by guessing."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "path": {"type": "string", "pattern": "^[A-Za-z0-9_.+/-]{1,200}$"},
                "start_line": {"type": "integer", "minimum": 1, "maximum": 1000000, "default": 1},
                "line_count": {"type": "integer", "minimum": 1, "maximum": MAX_READ_LINES, "default": 80},
            },
            "required": ["path"],
            "additionalProperties": False,
        },
    },
]


def result_error(message: str, **details: Any) -> dict[str, Any]:
    return {"ok": False, "error": message, "retry": False, **details}


def safe_path(relative: str) -> Path | None:
    if not relative or relative.startswith("/"):
        return None
    candidate = Path(relative)
    if ".." in candidate.parts or any(part.lower() in BLOCKED_PARTS for part in candidate.parts):
        return None
    target = (REPO_ROOT / candidate).resolve(strict=False)
    try:
        target.relative_to(REPO_ROOT)
    except ValueError:
        return None
    if candidate.name != "ATHENA.md" and (not candidate.parts or candidate.parts[0] != "docs"):
        return None
    if target.suffix.lower() not in ALLOWED_SUFFIXES:
        return None
    return target


def read_text(path: Path, limit: int | None = None) -> str:
    text = path.read_text(encoding="utf-8", errors="replace")
    return text if limit is None else text[:limit]


def overview() -> dict[str, Any]:
    path = REPO_ROOT / "ATHENA.md"
    try:
        content = read_text(path, MAX_OVERVIEW_CHARS)
    except (OSError, PermissionError) as exc:
        return result_error("ATHENA.md is unavailable", path="ATHENA.md", detail=str(exc))
    return {"ok": True, "source": "ATHENA.md", "content": content, "truncated": path.stat().st_size > len(content.encode())}


def current_state() -> dict[str, Any]:
    try:
        value = json.loads(RUNTIME_FILE.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        return result_error("runtime snapshot is unavailable", detail=str(exc))
    containers = value.get("containers") or []
    return {
        "ok": True,
        "generated_at": value.get("generated_at"),
        "hostname": value.get("hostname"),
        "active_inference_profiles": value.get("active_inference_profiles") or [],
        "source_commit": value.get("source_commit"),
        "gpus": value.get("gpus") or [],
        "containers": containers,
        "container_count": len(containers),
        "recovery_kit": value.get("recovery_kit") or {"present": False},
    }


def external_services() -> dict[str, Any]:
    path = REPO_ROOT / "config/service-catalog.json"
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        return result_error("service catalog is unavailable", detail=str(exc))
    services = []
    for item in value.get("services", []):
        services.append({key: item.get(key) for key in ("id", "name", "host", "address", "port", "protocol", "purpose") if item.get(key) is not None})
    return {"ok": True, "services": services, "count": len(services)}


def reference_files() -> list[Path]:
    files = [REPO_ROOT / "ATHENA.md"]
    docs = REPO_ROOT / "docs"
    try:
        files.extend(sorted(path for path in docs.glob("*.md") if path.is_file()))
    except OSError:
        pass
    return files


def search_reference(arguments: dict[str, Any]) -> dict[str, Any]:
    query = str(arguments.get("query", "")).strip()
    if len(query) < 2:
        return result_error("query must contain at least two characters")
    pattern = re.compile(re.escape(query), re.IGNORECASE)
    matches: list[dict[str, Any]] = []
    for path in reference_files():
        try:
            lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
        except OSError:
            continue
        for number, line in enumerate(lines, 1):
            if pattern.search(line):
                matches.append({
                    "path": str(path.relative_to(REPO_ROOT)),
                    "line": number,
                    "excerpt": line.strip()[:280],
                })
                if len(matches) >= MAX_SEARCH_RESULTS:
                    return {"ok": True, "query": query, "matches": matches, "truncated": True}
    return {"ok": True, "query": query, "matches": matches, "truncated": False}


def read_reference(arguments: dict[str, Any]) -> dict[str, Any]:
    relative = str(arguments.get("path", ""))
    path = safe_path(relative)
    if path is None:
        return result_error("path is not an allowed documentation path", path=relative)
    if not path.is_file():
        return result_error("documentation file not found", path=relative)
    start = max(1, int(arguments.get("start_line", 1)))
    count = min(MAX_READ_LINES, max(1, int(arguments.get("line_count", 80))))
    try:
        lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError as exc:
        return result_error("documentation file is unreadable", path=relative, detail=str(exc))
    selected = lines[start - 1:start - 1 + count]
    return {
        "ok": True,
        "path": relative,
        "start_line": start,
        "end_line": start + len(selected) - 1 if selected else start - 1,
        "total_lines": len(lines),
        "content": "\n".join(selected),
        "truncated": start - 1 + len(selected) < len(lines),
    }


def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
    if name == "athena_get_overview":
        return overview()
    if name == "athena_get_current_state":
        return current_state()
    if name == "athena_get_external_services":
        return external_services()
    if name == "athena_search_reference":
        return search_reference(arguments)
    if name == "athena_read_reference":
        return read_reference(arguments)
    return result_error("unknown tool", tool=name)


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-platform-context", "version": VERSION},
        })
    elif method == "tools/list":
        emit(request_id, {"tools": TOOLS})
    elif method == "tools/call":
        params = message.get("params") or {}
        value = call_tool(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,
        })
    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()