"""Kept as an MCP server (stdio): any MCP client (Claude Desktop, Claude Code, agent frameworks with MCP support) gets three
tools: kept_remember, kept_recall, kept_credit. No dependencies beyond the Kept SDK file next to it (kept.py).

Claude Code:     claude mcp add kept -e KEPT_AGENT=my-agent -e KEPT_API_KEY=kept_... -- python /path/to/kept_mcp.py
Claude Desktop:  "mcpServers": {"kept": {"command": "python", "args": ["/path/to/kept_mcp.py"],
                                          "env": {"KEPT_AGENT": "my-agent", "KEPT_API_KEY": "kept_..."}}}
KEPT_AGENT defaults to the public demo agent (read-only: recall and credit work, remember doesn't)."""
from __future__ import annotations

import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    from kept import Kept, KeptError                                  # kept.py beside this file, or the installed package
except ImportError:  # pragma: no cover
    sys.stderr.write("put kept.py (from https://keptlayer.com/sdk/kept.py) next to kept_mcp.py\n")
    raise

MEM = Kept(agent=os.environ.get("KEPT_AGENT", "demo"), api_key=os.environ.get("KEPT_API_KEY") or None,
           url=os.environ.get("KEPT_URL", "https://keptlayer.com"))

TOOLS = [
    {"name": "kept_remember", "description": "Record something that happened (a tool call, an error, a decision, an outcome) in the agent's permanent memory.",
     "inputSchema": {"type": "object", "required": ["content"], "properties": {
         "content": {"type": "string", "description": "What happened, in plain words"},
         "type": {"type": "string", "description": "goal, decision, tool_call, command, error, outcome, file, observation, feedback, note"},
         "task_id": {"type": "string", "description": "Groups events of one task or session"},
         "metadata": {"type": "object", "description": "Optional fields such as project, topic, outcome, value"}}}},
    {"name": "kept_recall", "description": "Ask the agent's memory a question about its past. Returns only the memories worth loading, each with its original evidence.",
     "inputSchema": {"type": "object", "required": ["query"], "properties": {
         "query": {"type": "string"}, "budget": {"type": "integer", "description": "Max context tokens to return (default 1500)"}}}},
    {"name": "kept_credit", "description": "After finishing a task, report whether the recalled memories helped. This is the only thing that changes future recall.",
     "inputSchema": {"type": "object", "required": ["recall_id", "success"], "properties": {
         "recall_id": {"type": "string"}, "success": {"type": "boolean"},
         "used": {"type": "array", "items": {"type": "integer"}, "description": "Event ids the answer actually used"}}}},
]


def call(name: str, args: dict) -> str:
    if name == "kept_remember":
        r = MEM.remember({k: v for k, v in args.items() if k in ("content", "type", "task_id", "metadata")})
        return "remembered %d event(s); chain head %s" % (r["accepted"], r["head"][:16])
    if name == "kept_recall":
        r = MEM.recall(args["query"], budget=int(args.get("budget") or 1500))
        ev = "\n".join("- event %d (%s, %s, #%s): %s" % (e["event_id"], e["time"], e["type"], e["hash"][:8], e["content"]) for e in r["evidence"])
        warn = "".join("\n! records disagree in %s: %s" % (c["label"], ", ".join(c["subjects"])) for c in r.get("contradictions", []))
        return "recall_id: %s\nanswer: %s\nevidence:\n%s\ncontext:\n%s\n(read %d of %d events, %d tokens)%s" % (
            r["recall_id"], r["answer"], ev or "- none", r["context"], r["cost"]["events_read"], r["lifetime"], r["cost"]["context_tokens"], warn)
    if name == "kept_credit":
        r = MEM.credit(args["recall_id"], success=bool(args["success"]), used=args.get("used"))
        return "credited: %d branches strengthened, %d weakened" % (r["credited"], r["wasted"])
    raise KeyError(name)


def main():
    for line in sys.stdin:
        if not line.strip():
            continue
        msg = json.loads(line)
        mid, method = msg.get("id"), msg.get("method")
        if mid is None:
            continue                                                  # notifications need no answer
        try:
            if method == "initialize":
                res = {"protocolVersion": msg.get("params", {}).get("protocolVersion", "2024-11-05"),
                       "capabilities": {"tools": {}}, "serverInfo": {"name": "kept", "version": "0.5"}}
            elif method == "tools/list":
                res = {"tools": TOOLS}
            elif method == "tools/call":
                p = msg.get("params", {})
                try:
                    res = {"content": [{"type": "text", "text": call(p.get("name"), p.get("arguments") or {})}]}
                except (KeptError, KeyError, ValueError) as e:
                    res = {"content": [{"type": "text", "text": "Kept error: %s" % e}], "isError": True}
            elif method == "ping":
                res = {}
            else:
                raise NotImplementedError(method)
            out = {"jsonrpc": "2.0", "id": mid, "result": res}
        except NotImplementedError as e:
            out = {"jsonrpc": "2.0", "id": mid, "error": {"code": -32601, "message": "unknown method %s" % e}}
        sys.stdout.write(json.dumps(out) + "\n")
        sys.stdout.flush()


if __name__ == "__main__":
    main()