"""MCP (Model Context Protocol) endpoint for the Jarvis memory service.

Mounts a SINGLE streamable-HTTP MCP endpoint onto the same FastAPI app that
serves the REST API, so Claude Code and the claude.ai app can share the one
memory brain. The tools are THIN wrappers that call app.py's OWN functions in
process (never over HTTP to ourselves, no logic duplicated).

Two doors, same endpoint (auth enforced HERE, before the MCP handler runs):
  • POST /mcp            — Authorization: Bearer $MEMORY_API_KEY   (Claude Code --header)
  • POST /mcp/<token>    — token = first 32 hex of sha256(MEMORY_API_KEY)
                           (path-secret door for claude.ai custom connectors,
                           whose UI can't send custom headers; OAuth is overkill
                           for a personal service).
If MEMORY_API_KEY is empty, BOTH doors are open (matches app.py `_auth`).

Wiring (in app.py): `import mcp_server`, run `mcp_server.lifespan()` from the
FastAPI lifespan (the session manager needs its task group running), and call
`mcp_server.mount(app)` to attach the two routes. Nothing else changes.
"""
from __future__ import annotations

import hashlib
import hmac
import os
from contextlib import asynccontextmanager

import anyio
from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
from starlette.routing import Route

API_KEY = os.environ.get("MEMORY_API_KEY", "")
_BEARER = f"Bearer {API_KEY}"
# claude.ai custom connectors can't send custom headers, so the second door is a
# path secret: the first 32 hex chars of sha256(MEMORY_API_KEY). Print it locally
# with:  python3 -c "import hashlib,os;print(hashlib.sha256(os.environ['MEMORY_API_KEY'].encode()).hexdigest()[:32])"
PATH_TOKEN = hashlib.sha256(API_KEY.encode()).hexdigest()[:32] if API_KEY else ""

mcp = FastMCP(
    "jarvis-memory",
    stateless_http=True,
    json_response=True,
    # FastMCP auto-enables DNS-rebinding Host validation (localhost allowlist
    # only) whenever its `host` setting is 127.0.0.1 — behind Railway's proxy
    # that 421s every request ("Invalid Host header"). Our doors already gate
    # every call with a secret (bearer or path token), and rebinding protection
    # exists for UNauthenticated localhost servers, so disable it explicitly.
    transport_security=TransportSecuritySettings(
        enable_dns_rebinding_protection=False),
    instructions=(
        "Ahmed's shared long-term memory (a fact graph). Search it before "
        "answering from scratch, save durable facts, and checkpoint projects so "
        "any Claude — Code or app — can pick up where the last one stopped."),
)


# --- tools: thin wrappers over app.py's SAME internal functions ------------
# Each runs the (blocking) graph work in a worker thread so the MCP event loop
# (which also drives the session manager) never stalls. `import app` is lazy to
# avoid an import cycle — by call time app.py is fully loaded and started.

@mcp.tool()
async def memory_search(query: str, k: int = 8, when_from: str | None = None,
                        when_to: str | None = None,
                        as_of: str | None = None) -> dict:
    """Search shared memory: salience-ranked facts plus their connected context.

    Optional temporal filters (all ISO dates): `when_from`/`when_to` window the
    EVENT date a fact refers to; `as_of` answers "what was true THEN" bi-temporally.
    """
    import app
    return await anyio.to_thread.run_sync(
        lambda: app.search(q=query, k=k, group=None, hops=1,
                           include_stale=False, as_of=as_of,
                           when_from=when_from, when_to=when_to,
                           authorization=_BEARER))


@mcp.tool()
async def memory_remember(text: str, kind: str = "fact",
                          entities: list[str] | None = None,
                          when: str | None = None,
                          importance: int = 5) -> dict:
    """Save one durable fact; `entities` are names (people/companies/projects) to attach.

    `when` is the ISO date the fact REFERS to — event time, not now (e.g. a meeting
    date), used for temporal recall. `importance` (1-10, default 5) weights recall.
    """
    import app
    ents = [{"name": e} for e in (entities or []) if str(e).strip()]
    body = app.Remember(text=text, kind=(kind or "fact"), source="mcp",
                        entities=ents or None, when=when, importance=importance)
    return await anyio.to_thread.run_sync(
        lambda: app.remember(body, authorization=_BEARER))


@mcp.tool()
async def project_checkpoint(project: str, done: str, next_steps: str,
                             details: str | None = None) -> dict:
    """Save a 'where we stopped' checkpoint for a project: what's DONE and what's NEXT."""
    import app
    text = f"PROJECT {project} checkpoint: DONE {done}. NEXT {next_steps}."
    if details and details.strip():
        text += f" {details.strip()}"
    # stamp the checkpoint with NOW as its event date so temporal recall orders
    # a project's checkpoints chronologically.
    body = app.Remember(text=text, kind="checkpoint", source="mcp",
                        entities=[{"name": project, "type": "project"}],
                        when=app._now_iso())
    return await anyio.to_thread.run_sync(
        lambda: app.remember(body, authorization=_BEARER))


@mcp.tool()
async def entity_lookup(name: str) -> dict:
    """Pull up everything memory knows about one entity (person / company / project)."""
    import app
    return await anyio.to_thread.run_sync(
        lambda: app.entity(name=name, group=None, authorization=_BEARER))


@mcp.tool()
async def entity_profile(name: str) -> dict:
    """Full dossier of one entity (person/company/concept): summary, contacts, relations, timeline of every interaction, open tasks. Use when Ahmed asks 'tell me about / show me X' or before a meeting with X."""
    import app
    return await anyio.to_thread.run_sync(
        lambda: app.profile(name=name, group=None, authorization=_BEARER))


# Create the session manager (lazy — only exists after streamable_http_app()).
_ = mcp.streamable_http_app()
session_manager = mcp.session_manager


@asynccontextmanager
async def lifespan():
    """Enter from the FastAPI lifespan so the session manager's task group runs."""
    async with session_manager.run():
        yield


def _authorized(scope, token: str | None) -> bool:
    if not API_KEY:
        return True  # both doors open when no key is configured (matches _auth)
    if token is not None:                       # path-secret door
        return hmac.compare_digest(token, PATH_TOKEN)
    headers = {k.decode().lower(): v.decode()   # bearer-header door
               for k, v in scope.get("headers", [])}
    return hmac.compare_digest(headers.get("authorization", ""), _BEARER)


async def _reject(send) -> None:
    await send({"type": "http.response.start", "status": 401,
                "headers": [(b"content-type", b"application/json")]})
    await send({"type": "http.response.body", "body": b'{"error":"unauthorized"}'})


class _Door:
    """ASGI gate in front of the shared MCP handler. mode: 'header' or 'token'."""

    def __init__(self, mode: str):
        self.mode = mode

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await _reject(send)
            return
        token = (scope.get("path_params", {}).get("token")
                 if self.mode == "token" else None)
        if not _authorized(scope, token):
            await _reject(send)
            return
        await session_manager.handle_request(scope, receive, send)


def mount(app) -> None:
    """Attach both MCP doors to an existing FastAPI / Starlette app."""
    app.router.routes.append(Route("/mcp", _Door("header")))
    app.router.routes.append(Route("/mcp/{token}", _Door("token")))