"""MCP server for Cielara Agent Memory — a Claude Connector over the Agent
Memory API only.

This is a memory-only slice of the unified Cielara MCP server. It exposes the
three Agent Memory tools (search_memories, ask_memories, list_memory_collections)
and nothing else. Ontology and World Model tools are intentionally omitted.

It is a thin HTTP client over the Agent Memory API:

  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.

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 ""
    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"))

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 _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
# ===========================================================================


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"
        )
        return header + "\n\n---\n\n".join(parts)
    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)
            )
        return answer_section + sources_section
    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}"


# ---------------------------------------------------------------------------
# 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. One reasoning substrate:\n"
            " - Agent memory: search_memories for relevance-ranked retrieval, "
            "ask_memories for LLM-grounded answers with citations, "
            "list_memory_collections to enumerate corpora.\n"
            "All tools 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))

    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()
