#!/usr/bin/env python3
"""Read-only MCP server over the symbol atlas and the figure half, for agents
building on Babelon.

Same wire style as bridger's mcp/server.py: stdio JSON-RPC, standard library
only, nothing to install. Point an MCP client at it:

  {"mcpServers": {"babelon-atlas": {
     "command": "python3",
     "args": ["/path/to/babelon-atlas/mcp_server.py"],
     "env": {"BABELON_ATLAS": "/path/to/symbol-atlas.json",
             "BABELON_DATA": "/path/to/downloaded/babelon",
             "BABELON_ROOMS": "/path/to/invisible-rooms/rooms"}}}}

It never writes. It cannot reach jawnomicon's Neo4j or bridger's CanonGraph --
it reads the built atlas file, which is the point: the graphs stay where they
are and only the joined, tiered export travels.

The figure tools read the files jawnomicon.com publishes under /babelon/,
because an agent that could resolve `rainbow` and not `Korrigan` was answering
half the question. They fetch on demand over HTTPS and hold the result for the
life of the process; BABELON_DATA points at downloaded copies instead, and
BABELON_OFFLINE=1 refuses the network so a sealed host fails with a sentence
rather than a stack.
"""

from __future__ import annotations

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

import atlas as A

PROTOCOL_VERSION = "2024-11-05"

ATLAS_PATH = pathlib.Path(os.environ.get("BABELON_ATLAS", A.DEFAULT_ATLAS))
ROOMS_PATH = os.environ.get("BABELON_ROOMS")

# One reader for the whole process. A server answering forty relation lookups
# must not download descriptions.json forty times, and this is a cache for the
# life of the process, not a copy on disk: restart it and you get the current
# export.
PUBLISHED = A.Published(A.DEFAULT_DATA_DIR,
                        offline=os.environ.get("BABELON_OFFLINE", "") not in ("", "0", "false"))

TOOLS = [
    {
        "name": "symbol_lookup",
        "description": (
            "Look up one symbol by name or slug. Returns its gloss, how many "
            "sources attest it, any caveats (a homonym or multi-referent row "
            "holds more than one referent), its symbol-domain placement (world, "
            "subcategory, sense and furniture counts), and any creatures whose "
            "description cites it, with the citing sentence."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
            "additionalProperties": False,
        },
    },
    {
        "name": "symbol_search",
        "description": "Search symbol names and glosses for a keyword.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 20},
            },
            "required": ["query"],
            "additionalProperties": False,
        },
    },
    {
        "name": "figure_lookup",
        "description": (
            "Resolve a figure — by name, by any alias, by canon id, or by the "
            "name of a figure ruled to be the same being (an absorbed id). "
            "Returns traditions, regions, classes, affinity family, symbols "
            "borne, the reference works that name it, the lede, and the page "
            "URL. Use it rather than building a path: the route is /creature/, "
            "a slug carries a batch prefix (Korrigan is g3_korrigan), and an "
            "apostrophe becomes a DOUBLED hyphen, so no name yields its own URL."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
            "additionalProperties": False,
        },
    },
    {
        "name": "figure_search",
        "description": (
            "Search figure names, aliases and ledes, ranked — a name match "
            "before an alias before a lede, and close matches when nothing "
            "matches as a substring, because a miss is usually a misspelling "
            "rather than an absence."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 10},
            },
            "required": ["query"],
            "additionalProperties": False,
        },
    },
    {
        "name": "figure_relations",
        "description": (
            "One figure's relations — traditions, regions, classes, motifs, "
            "powers — with each warrant resolved to the text it cites. A row "
            "carries no sentence: it carries which description field the "
            "warrant sits in and [start, end) character offsets into it, and "
            "this tool fetches descriptions.json and takes the slice. What it "
            "returns is therefore JAWNOMICON'S OWN MACHINE-WRITTEN PROSE, never "
            "a quotation from a reference work — what is attested is which "
            "works name the figure, not their wording. Rows with a null span "
            "did something else and keep a literal note; they come back marked "
            "'literal-note' and must not be read as our prose."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
            "additionalProperties": False,
        },
    },
    {
        "name": "figure_profile",
        "description": (
            "One figure's 29-axis eidogen scores, each returned WITH the "
            "sentence its band carries and the two poles it sits between. A "
            "bare 0.2 on 'Cognition' says nothing on its own — near which end, "
            "meaning what — so this reads the number back into the written "
            "band a rater was choosing between. Scores are machine ratings of "
            "jawnomicon's own description by two model families, not a "
            "source's judgement of the figure. Axes returned under "
            "'not_applicable' carry a SENTINEL, not a low score: the axis does "
            "not apply to this being, and treating it as 0 reads every honest "
            "gap as an extreme position."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"name": {"type": "string"}},
            "required": ["name"],
            "additionalProperties": False,
        },
    },
    {
        "name": "room_seed",
        "description": (
            "Draw one or more room seeds: a symbol that carries a domain, with "
            "its furniture counts and any linked creature. Use it to start a "
            "room from attested material rather than from a mood."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "world": {"type": "string", "description": "e.g. 'Creation and Cosmos'"},
                "linked": {"type": "boolean", "default": False,
                           "description": "only symbols that have a creature"},
                "count": {"type": "integer", "minimum": 1, "maximum": 10, "default": 1},
                "seed": {"type": "integer", "description": "reproducible draw"},
            },
            "additionalProperties": False,
        },
    },
    {
        "name": "convergence",
        "description": (
            "The tic watch. Reads a Babelon rooms/ directory and reports which "
            "words the rooms already share and how many, plus a per-room "
            "mechanism check (present, and one sentence). Moods are why twenty "
            "rooms turn into one room; this measures the drift directly."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "rooms_dir": {"type": "string",
                              "description": "defaults to $BABELON_ROOMS"},
                "limit": {"type": "integer", "minimum": 1, "maximum": 200, "default": 25},
            },
            "additionalProperties": False,
        },
    },
]


def _atlas() -> dict:
    return A.load_atlas(ATLAS_PATH)


def _symbol_payload(atlas: dict, symbol: dict) -> dict:
    out = dict(symbol)
    creatures = []
    for link in symbol.get("creatures") or []:
        creature = A.creature_by_id(atlas, link["id"])
        creatures.append({**link, **({"creature": creature} if creature else {})})
    if creatures:
        out["creatures"] = creatures
    return out


def _symbol_lookup(name: str) -> dict:
    atlas = _atlas()
    symbol = A.find(atlas, name)
    if not symbol:
        return {"found": False, "name": name}
    return {"found": True, "symbol": _symbol_payload(atlas, symbol)}


def _symbol_search(query: str, limit: int) -> dict:
    atlas = _atlas()
    needle = query.lower()
    hits = [
        s for s in atlas["symbols"]
        if needle in s["name"].lower() or needle in (s.get("gloss") or "").lower()
    ]
    return {
        "total": len(hits),
        "results": [
            {
                "name": s["name"],
                "slug": s["slug"],
                "register": s["register"],
                "gloss": s.get("gloss"),
                "caveats": s.get("caveats"),
                "in": s["in"],
                "world": (s.get("domain") or {}).get("world"),
            }
            for s in hits[:limit]
        ],
    }


# What a figure row means, said with the row rather than left in a spec. The
# published file carries a longer `conventions` block; these are the three a
# caller gets wrong without being told.
FIGURE_CONVENTIONS = {
    "attested": ("recorded_in names the reference works whose entries name this "
                 "figure. Their text is never reproduced — that is the attested "
                 "part, and the only attested part."),
    "ours": "lede is jawnomicon's own machine-written prose, not a source's wording.",
    "absent_means": ("An empty list or a null family means NOT YET SURVEYED, never "
                     "'has none'. Coverage differs sharply per field."),
    "url": ("Use url or slug. A slug is the canon id with / \\ : ' ’ each replaced "
            "by a hyphen, so 2_will-o'-the-wisp is served at 2_will-o--the-wisp; "
            "it is not derivable from a name."),
}


def _figures() -> dict:
    return PUBLISHED.get("figures.json")


def _figure_lookup(name: str) -> dict:
    try:
        found = A.find_figures(_figures(), name)
    except A.DataUnavailable as exc:
        return {"error": str(exc)}
    if not found:
        return {"found": False, "name": name,
                "note": "no figure answers to that name, alias, canon id or absorbed"
                        " name — try figure_search, which matches substrings and"
                        " near-misses"}
    if len(found) > 1:
        # The file promises no uniqueness across every alias it carries, and
        # picking one quietly is the move this project refuses everywhere else.
        return {"found": False, "name": name, "ambiguous": [
            {"id": f["id"], "name": f["name"], "url": f["url"]} for f in found]}
    return {"found": True, "figure": found[0], "conventions": FIGURE_CONVENTIONS}


def _figure_search(query: str, limit: int) -> dict:
    try:
        hits = A.search_figures(_figures(), query)
    except A.DataUnavailable as exc:
        return {"error": str(exc)}
    return {
        "total": len(hits),
        "results": [
            {
                "name": hit["figure"]["name"],
                "id": hit["figure"]["id"],
                "slug": hit["figure"]["slug"],
                "url": hit["figure"]["url"],
                "matched": hit["how"],
                "aliases": hit["figure"].get("aliases"),
                "lede": hit["figure"].get("lede"),
            }
            for hit in hits[:limit]
        ],
    }


def _figure_profile(name: str) -> dict:
    try:
        found = A.find_figures(_figures(), name)
        if not found:
            return {"found": False, "name": name,
                    "note": "no figure answers to that name — try figure_search"}
        if len(found) > 1:
            return {"found": False, "name": name, "ambiguous": [
                {"id": f["id"], "name": f["name"], "url": f["url"]} for f in found]}
        return {"found": True, **A.figure_profile(PUBLISHED, found[0])}
    except A.DataUnavailable as exc:
        return {"error": str(exc)}


def _figure_relations(name: str) -> dict:
    try:
        found = A.find_figures(_figures(), name)
        if not found:
            return {"found": False, "name": name,
                    "note": "no figure answers to that name — try figure_search"}
        if len(found) > 1:
            return {"found": False, "name": name, "ambiguous": [
                {"id": f["id"], "name": f["name"], "url": f["url"]} for f in found]}
        return {"found": True, **A.figure_relations(PUBLISHED, found[0])}
    except A.DataUnavailable as exc:
        return {"error": str(exc)}


def _room_seed(world: str | None, linked: bool, count: int, seed: int | None) -> dict:
    import random

    atlas = _atlas()
    pool = [s for s in atlas["symbols"] if s.get("domain")]
    if world:
        pool = [s for s in pool if s["domain"]["world"].lower().startswith(world.lower())]
    if linked:
        pool = [s for s in pool if s.get("creatures")]
    if not pool:
        return {"seeds": [], "note": "no symbol matches that filter"}
    drawn = random.Random(seed).sample(pool, min(count, len(pool)))
    return {"seeds": [_symbol_payload(atlas, s) for s in drawn]}


def _convergence(rooms_dir: str | None, limit: int) -> dict:
    target = rooms_dir or ROOMS_PATH
    if not target:
        return {"error": "no rooms_dir given and BABELON_ROOMS is unset"}
    path = pathlib.Path(target)
    paths = sorted(p for p in path.glob("*.md") if p.stem != "TEMPLATE")
    if not paths:
        return {"error": f"no rooms in {target}"}

    import collections

    rooms = [A.parse_room(p) for p in paths]
    use: dict[str, set[str]] = collections.defaultdict(set)
    for room in rooms:
        body = room["fields"].get("THE CITY", room["text"]).lower()
        for word in set(A.WORD.findall(body)):
            if word not in A.STOP:
                use[word].add(room["slug"])

    n = len(rooms)
    shared = sorted(
        ((w, sorted(r)) for w, r in use.items() if len(r) > 1),
        key=lambda item: (-len(item[1]), item[0]),
    )
    import re as _re

    mechanisms = []
    for room in rooms:
        mech = " ".join((room["fields"].get("MECHANISM") or "").split())
        mechanisms.append({
            "room": room["slug"],
            "author": room["meta"].get("author"),
            "present": bool(mech),
            "words": len(mech.split()),
            "sentences": len([s for s in _re.split(r"(?<=[.!?])\s+", mech) if s.strip()]),
        })
    return {
        "rooms": [r["slug"] for r in rooms],
        "shared_word_count": len(shared),
        "in_every_room": [w for w, r in shared if len(r) == n],
        "shared": [{"word": w, "breadth": len(r), "rooms": r} for w, r in shared[:limit]],
        "mechanisms": mechanisms,
        "incomplete": [
            {"room": room["slug"], "missing": field}
            for room in rooms
            for field in ("MECHANISM", "ACTION", "CONSEQUENCE", "DISCOVERABLE TEXT", "THE CITY")
            if not room["fields"].get(field)
        ],
    }


def _result(payload: Any) -> dict:
    """A tool's answer. A payload that is only {"error": ...} is a tool that could
    not do its job (data unreachable, no rooms to read), so it is marked isError:
    a client can then tell a failure from data. "Not found" and "ambiguous" stay
    ordinary results: they are answers, and carry no "error" key."""
    result = {"content": [{"type": "text",
                           "text": json.dumps(payload, ensure_ascii=False, indent=2)}]}
    if isinstance(payload, dict) and set(payload) == {"error"}:
        result["isError"] = True
    return result


def _error(request_id: Any, message: str, code: int = -32602) -> dict:
    return {"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}}


def _call(name: str, arguments: dict, request_id: Any) -> dict:
    if name == "symbol_lookup":
        if not isinstance(arguments.get("name"), str):
            return _error(request_id, "symbol_lookup requires a string 'name'")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_symbol_lookup(arguments["name"]))}
    if name == "symbol_search":
        if not isinstance(arguments.get("query"), str):
            return _error(request_id, "symbol_search requires a string 'query'")
        limit = arguments.get("limit", 20)
        if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= 100:
            return _error(request_id, "limit must be an integer from 1 to 100")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_symbol_search(arguments["query"], limit))}
    if name == "figure_lookup":
        if not isinstance(arguments.get("name"), str):
            return _error(request_id, "figure_lookup requires a string 'name'")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_figure_lookup(arguments["name"]))}
    if name == "figure_search":
        if not isinstance(arguments.get("query"), str):
            return _error(request_id, "figure_search requires a string 'query'")
        limit = arguments.get("limit", 10)
        if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= 100:
            return _error(request_id, "limit must be an integer from 1 to 100")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_figure_search(arguments["query"], limit))}
    if name == "figure_relations":
        if not isinstance(arguments.get("name"), str):
            return _error(request_id, "figure_relations requires a string 'name'")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_figure_relations(arguments["name"]))}
    if name == "figure_profile":
        if not isinstance(arguments.get("name"), str):
            return _error(request_id, "figure_profile requires a string 'name'")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_figure_profile(arguments["name"]))}
    if name == "room_seed":
        count = arguments.get("count", 1)
        if isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= 10:
            return _error(request_id, "count must be an integer from 1 to 10")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_room_seed(arguments.get("world"),
                                             bool(arguments.get("linked", False)),
                                             count, arguments.get("seed")))}
    if name == "convergence":
        limit = arguments.get("limit", 25)
        if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= 200:
            return _error(request_id, "limit must be an integer from 1 to 200")
        return {"jsonrpc": "2.0", "id": request_id,
                "result": _result(_convergence(arguments.get("rooms_dir"), limit))}
    return _error(request_id, f"unknown tool: {name}")


def handle(request: dict) -> dict | None:
    request_id = request.get("id")
    method = request.get("method")
    if request_id is None:
        return None
    if method == "initialize":
        return {"jsonrpc": "2.0", "id": request_id, "result": {
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "babelon-atlas", "version": "0.1.0"},
        }}
    if method == "ping":
        return {"jsonrpc": "2.0", "id": request_id, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": request_id, "result": {"tools": TOOLS}}
    if method == "tools/call":
        params = request.get("params") or {}
        arguments = params.get("arguments") or {}
        if not isinstance(arguments, dict):
            return _error(request_id, "tool arguments must be an object")
        return _call(params.get("name"), arguments, request_id)
    return _error(request_id, f"method not found: {method}", code=-32601)


def main() -> int:
    for line in sys.stdin:
        if not line.strip():
            continue
        request_id = None
        try:
            request = json.loads(line)
            if isinstance(request, dict):
                request_id = request.get("id")
            response = handle(request)
        except Exception:
            response = {"jsonrpc": "2.0", "id": request_id,
                        "error": {"code": -32603, "message": "internal error"}}
        if response is not None:
            sys.stdout.write(json.dumps(response, separators=(",", ":")) + "\n")
            sys.stdout.flush()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
