"""MCP server for Cielara Agent Memory — a Claude Connector over the Agent
Memory API, with an optional World Model bridge.

This is the memory slice of the unified Cielara MCP server. It always exposes
the three Agent Memory tools (search_memories, ask_memories,
list_memory_collections). When WORLD_MODEL_URL is configured it additionally
exposes the Agent Fabric World Model surface (query_world_model plus the
multi-turn inquiry tools start_inquiry / continue_inquiry / end_inquiry).
Ontology tools are intentionally omitted.

It is a thin HTTP client over those APIs:

  MEMORY_API_URL       (required) — e.g. http://agent-memory:8001
  CIELARA_BACKEND_URL  (optional) — only used to ship tool-call audit events
                                    to /api/mcp-audit/. Audit is skipped when
                                    unset.
  WORLD_MODEL_URL      (optional) — Agent Fabric wm-agent endpoint, e.g.
                                    http://agent-fabric:3001. When unset, the
                                    World Model and inquiry tools are not
                                    registered and the server is memory-only.

It does NOT import or embed any upstream code. All data access goes through the
published HTTP API. MEMORY_API_URL is required — the server fails fast on import
if it is unset.
"""

import functools
import json
import logging
import os
import threading
import time
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs

import requests
from mcp.server.fastmcp import FastMCP
from mcp.types import ToolAnnotations
from prometheus_client import CONTENT_TYPE_LATEST, Counter, Histogram, generate_latest

logger = logging.getLogger(__name__)

# Thread-local used by tool functions to pass token counts to the audit wrapper
# without changing return types.
_audit_ctx = threading.local()


def _set_token_counts(input_tokens: int | None, output_tokens: int | None) -> None:
    _audit_ctx.input_tokens = input_tokens
    _audit_ctx.output_tokens = output_tokens


def _pop_token_counts() -> tuple[int | None, int | None]:
    inp = getattr(_audit_ctx, "input_tokens", None)
    out = getattr(_audit_ctx, "output_tokens", None)
    _audit_ctx.input_tokens = None
    _audit_ctx.output_tokens = None
    return inp, out


_OPEN_PATHS = frozenset({
    "/live",
    "/metrics",
    "/.well-known/oauth-authorization-server",
    "/oauth/token",
})

# Opt-in scrape token. When set, /metrics requires
# `Authorization: Bearer <METRICS_AUTH_TOKEN>`; when unset, /metrics is
# anonymous (matches the other data-plane services).
_METRICS_AUTH_TOKEN = os.environ.get("METRICS_AUTH_TOKEN", "").strip()


_http_request_duration = Histogram(
    "http_request_duration_seconds",
    "Duration of HTTP requests in seconds",
    ["method", "route", "status_code"],
)

_http_requests_total = Counter(
    "http_requests_total",
    "Total number of HTTP requests",
    ["method", "route", "status_code", "status_class"],
)


def _status_class(status_code: int) -> str:
    if 200 <= status_code < 300:
        return "2xx"
    if 400 <= status_code < 500:
        return "4xx"
    if status_code >= 500:
        return "5xx"
    return "other"


# MCP streamable-http session paths embed UUIDs; collapse them so we don't
# explode label cardinality. Anything not in this set is bucketed as "other".
_KNOWN_ROUTES = frozenset({
    "/live",
    "/metrics",
    "/oauth/token",
    "/.well-known/oauth-authorization-server",
    "/mcp",
    "/mcp/",
    "/",
})


def _normalize_route(path: str) -> str:
    if path in _KNOWN_ROUTES:
        return path
    if path.startswith("/mcp"):
        return "/mcp/*"
    return "other"


class OAuthBearerApp:
    """ASGI app wrapping FastMCP with OAuth 2.0 discovery + Bearer auth.

    Implements just enough OAuth 2.0 for Claude.ai custom connectors:
      GET  /.well-known/oauth-authorization-server  — discovery metadata
      POST /oauth/token                             — client_credentials grant
      *    everything else                          — Bearer token auth gate

    Streaming-safe: never buffers the MCP response body.
    """

    def __init__(self, app: Any, token: str | None) -> None:
        self._app = app
        self._token = token

    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
        if scope["type"] != "http":
            await self._app(scope, receive, send)
            return

        path: str = scope.get("path", "")
        method: str = scope.get("method", "GET")
        start = time.monotonic()
        status_holder = {"code": 500}

        async def send_wrapper(message: Any) -> None:
            if message.get("type") == "http.response.start":
                status_holder["code"] = int(message.get("status", 500))
            await send(message)

        try:
            if path == "/live":
                status_holder["code"] = 200
                await self._live(send_wrapper)
                return

            if path == "/metrics":
                if not self._metrics_authorized(scope):
                    status_holder["code"] = 401
                    body = b'{"error":"unauthorized","error_description":"metrics auth required"}'
                    await send_wrapper({
                        "type": "http.response.start",
                        "status": 401,
                        "headers": [
                            (b"content-type", b"application/json"),
                            (b"content-length", str(len(body)).encode()),
                            (b"www-authenticate", b'Bearer realm="metrics"'),
                        ],
                    })
                    await send_wrapper({"type": "http.response.body", "body": body, "more_body": False})
                    return
                status_holder["code"] = 200
                await self._metrics(send_wrapper)
                return

            if path == "/.well-known/oauth-authorization-server":
                status_holder["code"] = 200
                await self._oauth_metadata(scope, send_wrapper)
                return

            if path == "/oauth/token":
                await self._oauth_token(scope, receive, send_wrapper)
                return

            if self._token is not None and path not in _OPEN_PATHS:
                headers = {k.lower(): v for k, v in scope.get("headers", [])}
                auth = headers.get(b"authorization", b"").decode("utf-8", "replace")
                if not auth.startswith("Bearer ") or auth[7:] != self._token:
                    status_holder["code"] = 401
                    body = b'{"error":"unauthorized","error_description":"Valid Bearer token required"}'
                    await send_wrapper({
                        "type": "http.response.start",
                        "status": 401,
                        "headers": [
                            (b"content-type", b"application/json"),
                            (b"content-length", str(len(body)).encode()),
                            (b"www-authenticate", b'Bearer realm="cielara-memory-mcp"'),
                        ],
                    })
                    await send_wrapper({"type": "http.response.body", "body": body, "more_body": False})
                    return

            await self._app(scope, receive, send_wrapper)
        finally:
            duration = time.monotonic() - start
            status = status_holder["code"]
            route = _normalize_route(path)
            status_str = str(status)
            _http_request_duration.labels(method, route, status_str).observe(duration)
            _http_requests_total.labels(method, route, status_str, _status_class(status)).inc()

    def _metrics_authorized(self, scope: Any) -> bool:
        """Validate the bearer token presented to /metrics.

        Open when METRICS_AUTH_TOKEN is unset; otherwise require
        `Authorization: Bearer <METRICS_AUTH_TOKEN>`.
        """
        if not _METRICS_AUTH_TOKEN:
            return True
        import secrets as _secrets
        headers = {k.lower(): v for k, v in scope.get("headers", [])}
        auth_header = headers.get(b"authorization", b"").decode("utf-8", "replace")
        if not auth_header.startswith("Bearer "):
            return False
        presented = auth_header[len("Bearer "):]
        return _secrets.compare_digest(presented, _METRICS_AUTH_TOKEN)

    async def _metrics(self, send: Any) -> None:
        body = generate_latest()
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (b"content-type", CONTENT_TYPE_LATEST.encode()),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body, "more_body": False})

    async def _live(self, send: Any) -> None:
        body = b'{"status":"ok"}'
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body, "more_body": False})

    async def _oauth_metadata(self, scope: Any, send: Any) -> None:
        headers = dict(scope.get("headers", []))
        host = headers.get(b"host", b"localhost").decode("utf-8", "replace")
        scheme = scope.get("scheme", "https")
        base = f"{scheme}://{host}"
        body = json.dumps({
            "issuer": base,
            "token_endpoint": f"{base}/oauth/token",
            "token_endpoint_auth_methods_supported": ["client_secret_post"],
            "grant_types_supported": ["client_credentials"],
            "scopes_supported": [],
            "response_types_supported": ["token"],
        }).encode()
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body, "more_body": False})

    async def _oauth_token(self, scope: Any, receive: Any, send: Any) -> None:
        chunks: list[bytes] = []
        while True:
            msg = await receive()
            chunks.append(msg.get("body", b""))
            if not msg.get("more_body", False):
                break
        params = parse_qs(b"".join(chunks).decode("utf-8", "replace"))

        grant_type = params.get("grant_type", [""])[0]
        client_secret = params.get("client_secret", [""])[0]

        if grant_type != "client_credentials" or client_secret != self._token:
            body = b'{"error":"invalid_client","error_description":"Invalid client credentials"}'
            await send({
                "type": "http.response.start",
                "status": 401,
                "headers": [
                    (b"content-type", b"application/json"),
                    (b"content-length", str(len(body)).encode()),
                ],
            })
            await send({"type": "http.response.body", "body": body, "more_body": False})
            return

        body = json.dumps({
            "access_token": self._token,
            "token_type": "bearer",
            "expires_in": 3600,
        }).encode()
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body, "more_body": False})


def _connector_config_path() -> Path:
    configured = os.environ.get("CIELARA_CONNECTOR_CONFIG", "").strip()
    if configured:
        return Path(configured).expanduser()
    return Path.home() / ".cielara" / "connector.json"


def _connector_connection() -> dict[str, Any]:
    path = _connector_config_path()
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        logger.warning("Could not read Cielara connector config %s: %s", path, exc)
        return {}
    connections = data.get("connections")
    if isinstance(connections, dict):
        active = data.get("active") or "default"
        conn = connections.get(active)
        return conn if isinstance(conn, dict) else {}
    return data if isinstance(data, dict) else {}


def _connector_config_value(name: str) -> str:
    conn = _connector_connection()
    if name == "CIELARA_BACKEND_URL":
        return str(conn.get("backendUrl") or conn.get("baseUrl") or "").strip().rstrip("/")
    if name == "MEMORY_API_URL":
        memory_url = str(conn.get("memoryApiUrl") or "").strip().rstrip("/")
        if memory_url:
            return memory_url
        base_url = str(conn.get("baseUrl") or conn.get("backendUrl") or "").strip().rstrip("/")
        return f"{base_url}/api/agent-memory" if base_url else ""
    if name == "WORLD_MODEL_URL":
        return str(conn.get("worldModelUrl") or "").strip().rstrip("/")
    return ""


def _require_env(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        value = _connector_config_value(name)
    if not value:
        raise RuntimeError(
            f"{name} is required but not set. Configure it in the service env "
            "(e.g. mcp-server-static.env in production, or claude.json `env` "
            "for local stdio use) — no hardcoded default is provided to avoid "
            "silently hitting an unexpected endpoint."
        )
    return value.rstrip("/")


MEMORY_API_URL = _require_env("MEMORY_API_URL")
# Backend URL is optional here — only used to ship audit events. Audit is
# skipped when it is unset.
CIELARA_BACKEND_URL = (
    os.environ.get("CIELARA_BACKEND_URL", "").strip() or _connector_config_value("CIELARA_BACKEND_URL")
).rstrip("/")

MEMORY_API_TIMEOUT = float(os.environ.get("MEMORY_API_TIMEOUT", "120"))

# World Model is optional. When WORLD_MODEL_URL is unset, the world model and
# inquiry tools are not registered (the server stays memory-only).
WORLD_MODEL_URL = (
    os.environ.get("WORLD_MODEL_URL", "").strip() or _connector_config_value("WORLD_MODEL_URL")
).rstrip("/")
WORLD_MODEL_TIMEOUT = float(os.environ.get("WORLD_MODEL_TIMEOUT", "300"))
INQUIRY_TIMEOUT = float(os.environ.get("INQUIRY_TIMEOUT", "600"))

AUDIT_ENABLED = (
    os.environ.get("AUDIT_ENABLED", "true").lower() not in ("false", "0", "no")
    and bool(CIELARA_BACKEND_URL)
)

_SOURCE = os.environ.get("CIELARA_MCP_SOURCE", "").strip()


def _detect_source() -> str:
    if _SOURCE:
        return _SOURCE
    # Infer from well-known env vars set by each client
    if os.environ.get("CURSOR_TRACE_ID") or os.environ.get("CURSOR_SESSION_ID"):
        return "cursor"
    if os.environ.get("CLAUDE_CODE_ENTRYPOINT") or os.environ.get("CLAUDE_CLI_VERSION"):
        return "claude_code"
    if os.environ.get("CLAUDE_DESKTOP_VERSION"):
        return "claude_desktop"
    return "mcp_stdio"


_RESOLVED_SOURCE = _detect_source()


def _post_audit(payload: dict) -> None:
    """Fire-and-forget POST to the audit ingest endpoint. Runs in a daemon thread."""
    if not CIELARA_BACKEND_URL:
        return
    try:
        url = f"{CIELARA_BACKEND_URL}/api/mcp-audit/"
        requests.post(url, json=payload, timeout=5, headers=_auth_headers())
    except Exception as exc:  # noqa: BLE001
        logger.debug("audit post failed (non-critical): %s", exc)


def _audit_tool_call(fn: Any, args: tuple, kwargs: dict) -> str:
    """Wrap a tool function call with timing and fire-and-forget audit logging."""
    start = time.monotonic()
    result: str = ""
    status = "success"
    error_msg = ""
    try:
        result = fn(*args, **kwargs)
        return result
    except Exception as exc:
        status = "error"
        error_msg = str(exc)
        raise
    finally:
        if AUDIT_ENABLED:
            latency_ms = int((time.monotonic() - start) * 1000)
            input_tokens, output_tokens = _pop_token_counts()
            query_text = ""
            if args:
                query_text = str(args[0])
            elif kwargs:
                query_text = str(next(iter(kwargs.values())))

            payload: dict[str, Any] = {
                "tool_name": fn.__name__,
                "query_text": query_text,
                "status": status,
                "response_text": result,
                "error_msg": error_msg,
                "latency_ms": latency_ms,
                "source": _RESOLVED_SOURCE,
                "user_id": "",
                "metadata": {},
            }
            if input_tokens is not None:
                payload["input_tokens"] = input_tokens
            if output_tokens is not None:
                payload["output_tokens"] = output_tokens
            if input_tokens is not None and output_tokens is not None:
                payload["total_tokens"] = input_tokens + output_tokens
            t = threading.Thread(target=_post_audit, args=(payload,), daemon=True)
            t.start()


def _make_audited(fn: Any) -> Any:
    @functools.wraps(fn)
    def wrapper(*args: Any, **kwargs: Any) -> str:
        return _audit_tool_call(fn, args, kwargs)
    return wrapper


def _access_token() -> str:
    token_file = os.environ.get("CIELARA_MCP_TOKEN_FILE", "").strip()
    if token_file:
        try:
            token = Path(token_file).expanduser().read_text(encoding="utf-8").strip()
            if token:
                return token
        except Exception:
            pass
    token = os.environ.get("CIELARA_ACCESS_TOKEN", "").strip()
    if token:
        return token
    conn = _connector_connection()
    return str(conn.get("accessToken") or conn.get("access_token") or "").strip()


def _auth_headers(extra: dict[str, str] | None = None) -> dict[str, str]:
    headers = dict(extra or {})
    token = _access_token()
    if token and "Authorization" not in headers:
        headers["Authorization"] = f"Bearer {token}"
    return headers


def _post(base: str, path: str, payload: dict, timeout: float) -> dict:
    url = f"{base}{path}"
    resp = requests.post(url, json=payload, timeout=timeout, headers=_auth_headers())
    resp.raise_for_status()
    return resp.json()


def _get(base: str, path: str, params: dict | None = None, *, timeout: float) -> dict:
    url = f"{base}{path}"
    resp = requests.get(url, params=params or {}, timeout=timeout, headers=_auth_headers())
    resp.raise_for_status()
    return resp.json()


def _delete(base: str, path: str, *, timeout: float) -> dict:
    url = f"{base}{path}"
    resp = requests.delete(url, timeout=timeout, headers=_auth_headers())
    resp.raise_for_status()
    return resp.json()


def _classify_error(e: Exception) -> tuple[str, bool]:
    """Return (user-friendly message, is_expected)."""
    name = type(e).__name__
    msg = str(e)
    if "ReadTimeout" in name or "Read timed out" in msg:
        return "Service timed out — it may be overloaded or processing a slow upstream request. Try again shortly.", True
    if "ConnectTimeout" in name or "connect timeout" in msg.lower():
        return "Service is not reachable (connection timeout). Check that it is running.", True
    if "ConnectionError" in name or "Connection refused" in msg:
        return "Service is down or unreachable (connection refused). Check that it is running.", True
    if "HTTPError" in name:
        return f"Service returned an error: {msg}", False
    return f"Unexpected error: {msg}", False


# ===========================================================================
# Agent Memory tools
# ===========================================================================


_corpus_cache: dict[str, tuple[float, int]] = {}
_CORPUS_CACHE_TTL = 300  # 5 minutes


def _get_corpus_memory_count(collection: str) -> int:
    """Return total memory count for the collection, cached for 5 minutes."""
    now = time.monotonic()
    cached = _corpus_cache.get(collection)
    if cached and now - cached[0] < _CORPUS_CACHE_TTL:
        return cached[1]
    try:
        data = _get(
            MEMORY_API_URL, "/documents",
            {"collection": collection, "limit": "1"},
            timeout=MEMORY_API_TIMEOUT,
        )
        count = (data.get("stats") or {}).get("memories", 0)
        _corpus_cache[collection] = (now, count)
        return count
    except Exception:
        return cached[1] if cached else 0


def search_memories(query: str, collection: str = "default") -> str:
    """Search the Agent Memory store for memories relevant to a query.

    Returns memory titles and full text content, ranked by relevance. Uses
    hybrid search (semantic + BM25 + cue anchors) with policy-driven
    iterative retrieval.

    Args:
        query: Natural language search query.
        collection: Memory collection to search in.
    """
    try:
        data = _post(
            MEMORY_API_URL, "/retrieve",
            {"query": query, "collection": collection},
            timeout=MEMORY_API_TIMEOUT,
        )
        usage = data.get("usage") or {}
        _set_token_counts(usage.get("prompt_tokens") or usage.get("input_tokens"),
                          usage.get("completion_tokens") or usage.get("output_tokens"))
        memories = data.get("memories", [])
        if not memories:
            return f"No memories found for query: {query}"

        parts = [f"## {m['title']}\n{m['text']}" for m in memories]
        header = (
            f"Found {len(memories)} memories "
            f"({data.get('steps_taken', 0)} retrieval steps):\n\n"
        )
        retrieved_chars = sum(len(m.get("text", "")) for m in memories)
        corpus_memories = _get_corpus_memory_count(collection)
        metrics = {
            "retrieved_memories": len(memories),
            "retrieved_chars": retrieved_chars,
            "corpus_memories": corpus_memories,
        }
        metrics_block = f"\n\n<!-- metrics:{json.dumps(metrics)} -->"
        return header + "\n\n---\n\n".join(parts) + metrics_block
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("search_memories: %s", msg)
        else:
            logger.error("search_memories failed: %s", e, exc_info=True)
        return f"Error searching memories: {msg}"


def _post_ask_and_poll(question: str, collection: str) -> dict:
    """POST /ask, then poll /ask/{job_id} until completion. Returns the completed job dict."""
    data = _post(
        MEMORY_API_URL, "/ask",
        {"question": question, "collection": collection, "include_memories": True},
        timeout=MEMORY_API_TIMEOUT,
    )
    job_id = data.get("job_id")
    if not job_id:
        return data

    deadline = time.monotonic() + MEMORY_API_TIMEOUT
    delay = 0.5
    while time.monotonic() < deadline:
        time.sleep(delay)
        delay = min(delay * 1.5, 3.0)
        result = _get(MEMORY_API_URL, f"/ask/{job_id}", timeout=MEMORY_API_TIMEOUT)
        status = result.get("status", "")
        if status not in ("queued", "running"):
            return result
    raise TimeoutError(f"ask job {job_id} did not complete within {MEMORY_API_TIMEOUT}s")


def ask_memories(question: str, collection: str = "default") -> str:
    """Ask a question and get an LLM-generated answer grounded in stored memories.

    The system retrieves relevant memories, builds context, and uses an LLM
    to synthesize an answer. Returns the answer followed by source citations.

    Args:
        question: Natural language question.
        collection: Memory collection to search in.
    """
    try:
        data = _post_ask_and_poll(question, collection)
        usage = data.get("usage") or {}
        _set_token_counts(usage.get("input_tokens"), usage.get("output_tokens"))
        answer_section = f"## Answer\n\n{data.get('answer', '')}"

        nodes = (data.get("graph") or {}).get("nodes", [])
        sources_section = ""
        if nodes:
            source_lines = [f"- {n['title']}" for n in nodes]
            sources_section = (
                f"\n\n## Sources ({data.get('context_memories', 0)} memories, "
                f"{data.get('context_chars', 0)} chars, "
                f"{data.get('retrieval_steps', 0)} retrieval steps)\n\n"
                + "\n".join(source_lines)
            )

        metrics = {
            "context_chars": data.get("context_chars", 0),
            "context_memories": data.get("context_memories", 0),
            "retrieved_chars": data.get("retrieved_chars", 0),
            "retrieved_memories": data.get("retrieved_memories", 0),
        }
        metrics_block = f"\n\n<!-- metrics:{json.dumps(metrics)} -->"
        return answer_section + sources_section + metrics_block
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("ask_memories: %s", msg)
        else:
            logger.error("ask_memories failed: %s", e, exc_info=True)
        return f"Error answering question: {msg}"


def list_memory_collections() -> str:
    """List all available Agent Memory collections.

    Collections are created during document ingestion. Each collection
    contains its own set of memories, segments, and source documents.
    """
    try:
        data = _get(MEMORY_API_URL, "/collections", timeout=MEMORY_API_TIMEOUT)
        collections = data.get("collections", [])
        if not collections:
            return "No memory collections found. Ingest documents first."
        return "Available memory collections:\n" + "\n".join(f"- {c}" for c in collections)
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("list_memory_collections: %s", msg)
        else:
            logger.error("list_memory_collections failed: %s", e, exc_info=True)
        return f"Error listing collections: {msg}"


# ===========================================================================
# World Model tools (Agent Fabric) — only registered when WORLD_MODEL_URL is set
# ===========================================================================


def query_world_model(
    question: str,
    verbose: bool = False,
    scope: str = "default",
    provider: str = "",
    connection_id: str = "",
) -> str:
    """Ask the Agent Fabric World Model a natural-language question.

    The World Model is a skill-orchestrated reasoning layer that runs CLI/API
    procedures (kubectl, gcloud, terraform, GitHub) against the live
    operational topology graph and returns a synthesized answer. Use this for
    questions that need fresh live-system state or skill-driven workflow
    orchestration — distinct from the memory substrate (ingested
    documents/decisions) and the local workload-topology snapshot.

    Args:
        question: Natural language question or task for the World Model.
        verbose: When true, include iteration count and duration metadata.
        scope: Connection scope: "default", "all", or "workload".
        provider: Optional provider selector such as "gcp", "aws", "azure", or "kubernetes".
        connection_id: Optional exact Cielara connection ID for scoped live queries.
    """
    if not WORLD_MODEL_URL:
        return (
            "World Model is not configured on this MCP server. "
            "Set WORLD_MODEL_URL to the Agent Fabric wm-agent endpoint "
            "(e.g. http://agent-fabric:3001) to enable this tool."
        )
    try:
        prompt = _world_model_prompt_with_connection_scope(
            question,
            scope=scope,
            provider=provider,
            connection_id=connection_id,
        )
        data = _wm_start_inquiry_session(prompt, timeout=WORLD_MODEL_TIMEOUT)
        answer = (data.get("answer") or "").strip()
        err = (data.get("error") or "").strip()
        session_id = data.get("session_id", "") or ""
        parts = []
        if answer:
            parts.append(f"## Answer\n\n{answer}")
        else:
            parts.append("## Answer\n\n(empty)")
        if err:
            parts.append(f"\n## Error\n\n{err}")
        if verbose:
            meta = []
            if "elapsed_sec" in data:
                meta.append(f"- elapsed_sec: {data['elapsed_sec']}")
            if session_id:
                meta.append(f"- session_id: {session_id}")
            if "session_type" in data:
                meta.append(f"- session_type: {data['session_type']}")
            if meta:
                parts.append("\n## Metadata\n\n" + "\n".join(meta))
        return "\n".join(parts)
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("query_world_model: %s", msg)
        else:
            logger.error("query_world_model failed: %s", e, exc_info=True)
        return f"Error querying world model: {msg}"


def _world_model_prompt_with_connection_scope(
    question: str,
    *,
    scope: str,
    provider: str,
    connection_id: str,
) -> str:
    scope = (scope or "default").strip().lower()
    if scope not in {"default", "all", "workload"}:
        scope = "default"

    selectors = []
    if scope:
        selectors.append(f"scope={scope}")
    if provider.strip():
        selectors.append(f"provider={provider.strip().lower()}")
    if connection_id.strip():
        selectors.append(f"connection_id={connection_id.strip()}")

    if not selectors:
        return question

    return (
        "Connection selection directive for Agent Fabric: "
        + ", ".join(selectors)
        + ". Use the connection-context skill. For scope=all, fan out read-only "
        "queries across matching connections and group results by connection. "
        "For mutating operations, require an exact connection_id.\n\n"
        + question
    )


def _wm_start_inquiry_session(prompt: str, *, timeout: float) -> dict:
    """POST wm-agent /session/start in inquiry mode and return the parsed body.

    Shared by query_world_model (one-shot) and start_inquiry (multi-turn) — both
    world-model queries and code-inquiry sessions live on the same wm-agent
    endpoint, so the wire call is identical.
    """
    return _post(
        WORLD_MODEL_URL,
        "/api/wm-agent/v1/session/start",
        {"session_type": "inquiry", "prompt": prompt},
        timeout=timeout,
    )


def _wm_end_inquiry_session(session_id: str, *, timeout: float) -> dict:
    """DELETE wm-agent /session/{id}. Used to release one-shot world-model
    sessions and to satisfy explicit end_inquiry calls."""
    return _delete(
        WORLD_MODEL_URL,
        f"/api/wm-agent/v1/session/{session_id}",
        timeout=timeout,
    )


def _format_inquiry_response(data: dict) -> str:
    answer = (data.get("answer") or "").strip()
    session_id = data.get("session_id", "")
    err = (data.get("error") or "").strip()
    parts = []
    if answer:
        parts.append(f"## Answer\n\n{answer}")
    elif err:
        parts.append(f"## Error\n\n{err}")
    else:
        parts.append("## Answer\n\n(empty)")
    if session_id:
        parts.append(f"\n_Session: `{session_id}` (type: {data.get('session_type', 'inquiry')})_")
    if data.get("elapsed_sec"):
        parts.append(f"\n_Elapsed: {data['elapsed_sec']:.1f}s_")
    return "\n".join(parts)


def start_inquiry(question: str) -> str:
    """Start an inquiry session for questions about code, workloads, and infrastructure.

    Opens a new conversational session on the Agent Fabric wm-agent in
    inquiry mode. The agent has access to cloned repositories, CLI skills
    (kubectl, gcloud, terraform), and can read/analyze source code. Returns
    an answer and a session_id — pass the session_id to continue_inquiry
    for follow-up questions in the same conversation.

    Use this instead of query_world_model when you need multi-turn
    conversation about code or want the agent to explore repositories.

    Args:
        question: Natural language question or task.
    """
    if not WORLD_MODEL_URL:
        return (
            "Inquiry sessions require WORLD_MODEL_URL to be configured "
            "(Agent Fabric endpoint, e.g. http://agent-fabric:3001)."
        )
    try:
        data = _wm_start_inquiry_session(question, timeout=INQUIRY_TIMEOUT)
        return _format_inquiry_response(data)
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("start_inquiry: %s", msg)
        else:
            logger.error("start_inquiry failed: %s", e, exc_info=True)
        return f"Error starting inquiry: {msg}"


def continue_inquiry(session_id: str, question: str) -> str:
    """Continue an existing inquiry session with a follow-up question.

    Sends a follow-up message to an inquiry session previously started
    with start_inquiry. The session retains context from prior turns.

    Args:
        session_id: Session ID returned by start_inquiry.
        question: Follow-up question or task.
    """
    if not WORLD_MODEL_URL:
        return "Inquiry sessions require WORLD_MODEL_URL to be configured."
    try:
        data = _post(
            WORLD_MODEL_URL,
            "/api/wm-agent/v1/session/continue",
            {"session_id": session_id, "prompt": question},
            timeout=INQUIRY_TIMEOUT,
        )
        return _format_inquiry_response(data)
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("continue_inquiry: %s", msg)
        else:
            logger.error("continue_inquiry failed: %s", e, exc_info=True)
        return f"Error continuing inquiry: {msg}"


def end_inquiry(session_id: str) -> str:
    """End an inquiry session and release server-side resources.

    Call this when the conversation is finished. After ending, the
    session_id can no longer be used with continue_inquiry.

    Args:
        session_id: Session ID to clean up.
    """
    if not WORLD_MODEL_URL:
        return "Inquiry sessions require WORLD_MODEL_URL to be configured."
    try:
        data = _wm_end_inquiry_session(session_id, timeout=WORLD_MODEL_TIMEOUT)
        status = data.get("status", "unknown")
        return f"Session `{session_id}` ended (status: {status})."
    except Exception as e:
        msg, expected = _classify_error(e)
        if expected:
            logger.warning("end_inquiry: %s", msg)
        else:
            logger.error("end_inquiry failed: %s", e, exc_info=True)
        return f"Error ending inquiry: {msg}"


# ---------------------------------------------------------------------------
# Server factory
# ---------------------------------------------------------------------------


def create_server(host: str = "0.0.0.0", port: int = 8888) -> FastMCP:
    """Create and return a fully configured Cielara Agent Memory MCP server."""
    server = FastMCP(
        "cielara-memory",
        instructions=(
            "Cielara Agent Memory MCP server. Reasoning substrates:\n"
            " - Agent memory: search_memories for relevance-ranked retrieval, "
            "ask_memories for LLM-grounded answers with citations, "
            "list_memory_collections to enumerate corpora.\n"
            " - World model: query_world_model for skill-orchestrated operational "
            "reasoning against the live topology graph and CLI/API skills "
            "(kubectl, gcloud, terraform, GitHub). Only registered when "
            "WORLD_MODEL_URL is configured.\n"
            " - Inquiry sessions: start_inquiry / continue_inquiry / end_inquiry "
            "for multi-turn conversations about code, workloads, and "
            "infrastructure. The agent explores cloned repositories and uses "
            "CLI skills. Only registered when WORLD_MODEL_URL is configured.\n"
            "All tools except inquiry sessions are read-only."
        ),
        host=host,
        port=port,
        stateless_http=True,
    )
    _read_only = ToolAnnotations(readOnlyHint=True, destructiveHint=False)

    server.tool(annotations=_read_only)(_make_audited(search_memories))
    server.tool(annotations=_read_only)(_make_audited(ask_memories))
    server.tool(annotations=_read_only)(_make_audited(list_memory_collections))

    if WORLD_MODEL_URL:
        server.tool(annotations=_read_only)(_make_audited(query_world_model))
        server.tool()(_make_audited(start_inquiry))
        server.tool()(_make_audited(continue_inquiry))
        server.tool()(_make_audited(end_inquiry))
    else:
        logger.info(
            "WORLD_MODEL_URL not set — world model and inquiry tools disabled"
        )

    return server


def create_asgi_app(
    token: str | None = None,
    host: str = "0.0.0.0",
    port: int = 8888,
) -> Any:
    """Return an ASGI app for the Cielara Agent Memory MCP server.

    If *token* is provided, wraps the FastMCP app with OAuthBearerApp so
    every HTTP request must carry ``Authorization: Bearer <token>``.
    Pass ``None`` to disable auth (local dev only).
    """
    server = create_server(host=host, port=port)
    app: Any = server.streamable_http_app()
    return OAuthBearerApp(app, token=token)


mcp = create_server()