# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""MCP (Model Context Protocol) endpoint for OpenViking server.

Exposes tools to Claude Code (or any MCP client) via streamable HTTP:
  find, search, read, list, remember, add_resource, grep, glob,
  code_outline, code_search, code_expand, forget, health

Mounted on the FastAPI app at /mcp. The MCP session manager lifecycle is
tied to the FastAPI app lifespan (not a sub-app lifespan) so the task group
is always initialized before requests arrive.

Identity headers (X-OpenViking-Account, X-OpenViking-User)
are extracted from HTTP request scope and propagated via contextvars.
"""

from __future__ import annotations

import asyncio
import contextvars
import os
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any, List, Literal, Optional
from urllib.parse import quote

from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import BaseModel, Field
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.types import ASGIApp, Receive, Scope, Send

from openviking.parse.parsers.code.ast.code_tools import (
    CODE_SEARCH_CONCURRENCY,
    expand_symbol,
    filter_code_uris,
    outline_file,
    search_symbols,
)
from openviking.server.auth import resolve_actor_peer_headers, resolve_identity
from openviking.server.dependencies import get_server_config, get_service
from openviking.server.identity import RequestContext
from openviking.server.local_input_guard import (
    TEMP_FILE_ID_RE,
    is_remote_resource_source,
)
from openviking.server.resource_ingest import ingest_temp_upload
from openviking.server.temp_upload_store import TempUploadStore
from openviking.server.upload_token_store import upload_token_store
from openviking_cli.exceptions import (
    InvalidArgumentError,
    PermissionDeniedError,
    UnauthenticatedError,
)
from openviking_cli.session.user_id import UserIdentifier
from openviking_cli.utils import get_logger

# Backwards-compatible alias for existing tests that import this private name.
_filter_code_uris = filter_code_uris

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Identity propagation via contextvars
# ---------------------------------------------------------------------------

_mcp_ctx: contextvars.ContextVar[Optional[RequestContext]] = contextvars.ContextVar(
    "_mcp_ctx", default=None
)

# URL hints from the incoming request, captured by middleware so MCP tools can
# reconstruct the agent-facing public URL without knowing about ASGI/Starlette.
# Only used as a fallback when neither OPENVIKING_PUBLIC_BASE_URL nor
# ServerConfig.public_base_url is set.
_request_url_ctx: contextvars.ContextVar[Optional[dict]] = contextvars.ContextVar(
    "_request_url_ctx", default=None
)


def _get_ctx() -> RequestContext:
    ctx = _mcp_ctx.get()
    if ctx is None:
        raise UnauthenticatedError("MCP request identity not set")
    return ctx


def _scope_to_origin(scope: Scope) -> Optional[str]:
    """Derive the public-facing origin (scheme://host) from an ASGI scope.

    Resolution order matches openviking.server.oauth.router._public_origin:
      1. ``OPENVIKING_PUBLIC_BASE_URL`` environment variable
      2. ``app.state.oauth_config.issuer`` (if OAuth enabled)
      3. ``X-Forwarded-Proto`` / ``X-Forwarded-Host``
      4. scope's own scheme + Host header
    """
    import os as _os

    env_value = _os.environ.get("OPENVIKING_PUBLIC_BASE_URL", "").strip()
    if env_value:
        return env_value.rstrip("/")

    app = scope.get("app")
    if app is not None:
        cfg = getattr(app.state, "oauth_config", None)
        configured = getattr(cfg, "issuer", None) if cfg else None
        if configured:
            return configured.rstrip("/")

    headers = {
        k.decode("latin-1").lower(): v.decode("latin-1") for k, v in scope.get("headers", [])
    }
    proto = headers.get("x-forwarded-proto") or scope.get("scheme") or "http"
    proto = proto.split(",", 1)[0].strip()
    host = headers.get("x-forwarded-host") or headers.get("host")
    if not host:
        server = scope.get("server")
        if isinstance(server, (list, tuple)) and len(server) >= 2:
            host = f"{server[0]}:{server[1]}" if server[1] else str(server[0])
    if not host:
        return None
    host = host.split(",", 1)[0].strip()
    return f"{proto}://{host}"


def _oauth_enabled(scope: Scope) -> bool:
    """Return True if app.state has an oauth_provider (i.e. OAuth is configured)."""
    app = scope.get("app")
    if app is None:
        return False
    return getattr(app.state, "oauth_provider", None) is not None


class _IdentityASGIMiddleware:
    """ASGI middleware: delegates to auth.resolve_identity (the same function
    used by all REST API routes) so authentication logic is never duplicated."""

    def __init__(self, app: ASGIApp):
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        request = Request(scope)
        try:
            identity = await resolve_identity(
                request,
                x_api_key=request.headers.get("x-api-key"),
                authorization=request.headers.get("authorization"),
                x_openviking_account=request.headers.get("x-openviking-account"),
                x_openviking_user=request.headers.get("x-openviking-user"),
            )
            actor_peer_id, legacy_agent_id = resolve_actor_peer_headers(
                request.headers.get("x-openviking-actor-peer"),
                request.headers.get("x-openviking-agent"),
            )
        except (UnauthenticatedError, PermissionDeniedError, InvalidArgumentError) as exc:
            status = (
                401
                if isinstance(exc, UnauthenticatedError)
                else (403 if isinstance(exc, PermissionDeniedError) else 400)
            )
            headers: dict[str, str] = {}
            # When OAuth is enabled and the request is unauthenticated, advertise
            # the OAuth 2.0 protected resource metadata so MCP clients (Claude.ai,
            # Claude Desktop, etc.) can auto-discover the authorization server
            # per RFC 9728 §5.1.
            if status == 401 and _oauth_enabled(scope):
                origin = _scope_to_origin(scope)
                if origin:
                    headers["WWW-Authenticate"] = (
                        f'Bearer resource_metadata="{origin}/.well-known/oauth-protected-resource"'
                    )
            resp = JSONResponse(
                {"jsonrpc": "2.0", "id": None, "error": {"code": -32001, "message": str(exc)}},
                status_code=status,
                headers=headers,
            )
            return await resp(scope, receive, send)

        ctx = RequestContext(
            user=UserIdentifier(
                identity.account_id or "default",
                identity.user_id or "default",
            ),
            role=identity.role,
            actor_peer_id=actor_peer_id,
            legacy_agent_id=legacy_agent_id,
            from_oauth=identity.from_oauth,
        )
        url_info = {
            "x_forwarded_proto": request.headers.get("x-forwarded-proto"),
            "x_forwarded_host": request.headers.get("x-forwarded-host"),
            "host": request.headers.get("host"),
        }
        ctx_token = _mcp_ctx.set(ctx)
        url_token = _request_url_ctx.set(url_info)
        try:
            return await self.app(scope, receive, send)
        finally:
            _mcp_ctx.reset(ctx_token)
            _request_url_ctx.reset(url_token)


# ---------------------------------------------------------------------------
# MCP server tools (aligned with vikingbot/agent/tools/ov_file.py)
# ---------------------------------------------------------------------------

mcp = FastMCP(
    "openviking",
    transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)


# -- find / search ---------------------------------------------------------


@mcp.tool()
async def find(
    query: str,
    target_uri: str = "",
    limit: int = 10,
    min_score: float = 0.35,
    level: Optional[List[int]] = None,
) -> str:
    """Fast semantic retrieval without session context. Returns ranked memories, resources, and skills with URI, abstract, and score."""
    service = get_service()
    result = await service.search.find(
        query=query,
        ctx=_get_ctx(),
        target_uri=target_uri,
        limit=limit,
        score_threshold=min_score,
        level=level,
    )
    return _format_search_result(result)


@mcp.tool()
async def search(
    query: str,
    target_uri: str = "",
    session_id: Optional[str] = None,
    limit: int = 10,
    min_score: float = 0.35,
    level: Optional[List[int]] = None,
) -> str:
    """Deep semantic retrieval with optional session context and intent analysis. Returns ranked memories, resources, and skills with URI, abstract, and score."""
    service = get_service()
    ctx = _get_ctx()
    session = None
    if session_id:
        session = service.sessions.session(ctx, session_id)
        await session.load()
    result = await service.search.search(
        query=query,
        ctx=ctx,
        target_uri=target_uri,
        session=session,
        limit=limit,
        score_threshold=min_score,
        level=level,
    )
    return _format_search_result(result)


def _format_search_result(result) -> str:
    items = []
    for ctx_type, contexts in [
        ("memory", result.memories),
        ("resource", result.resources),
        ("skill", result.skills),
    ]:
        for m in contexts:
            items.append((ctx_type, m))

    if not items:
        return "No matching context found."

    lines = []
    for ctx_type, m in items:
        abstract = (
            getattr(m, "abstract", "") or getattr(m, "overview", "") or "(no abstract)"
        ).strip()
        score = getattr(m, "score", 0.0)
        lines.append(f"- [{ctx_type} {score * 100:.0f}%] {m.uri}\n    {abstract}")

    return (
        f"Found {len(items)} item(s):\n\n"
        + "\n".join(lines)
        + "\n\nUse the read tool to expand a URI."
    )


# -- read ------------------------------------------------------------------


@mcp.tool()
async def read(uris: str | list[str]) -> str:
    """Read full content from one or more viking:// file URIs. Pass a single URI string or a list for batch reads. For directory listing, use the list tool instead."""
    import asyncio

    service = get_service()
    ctx = _get_ctx()
    uri_list = uris if isinstance(uris, list) else [uris]
    semaphore = asyncio.Semaphore(10)

    async def _read_one(uri: str) -> str:
        async with semaphore:
            try:
                body = await service.fs.read(uri, ctx=ctx)
                if isinstance(body, str) and body.strip():
                    return body
            except Exception:
                pass
            return f"(nothing found at {uri})"

    if len(uri_list) == 1:
        return await _read_one(uri_list[0])

    results = await asyncio.gather(*[_read_one(u) for u in uri_list])
    parts = []
    for uri, text in zip(uri_list, results, strict=True):
        parts.append(f"=== {uri} ===\n{text}")
    return "\n\n".join(parts)


# -- list ------------------------------------------------------------------


@mcp.tool(name="list")
async def ls(uri: str, recursive: bool = False) -> str:
    """List files and subdirectories under a viking:// directory URI. Use recursive=true for deep listing."""
    service = get_service()
    ctx = _get_ctx()

    entries = await service.fs.ls(uri, ctx=ctx, recursive=recursive, output="original")
    if not entries:
        return f"(no entries under {uri})"

    lines = []
    for e in entries:
        name = e.get("name", "?") if isinstance(e, dict) else getattr(e, "name", "?")
        is_dir = e.get("isDir", False) if isinstance(e, dict) else getattr(e, "is_dir", False)
        entry_uri = e.get("uri", "") if isinstance(e, dict) else getattr(e, "uri", "")
        if recursive and entry_uri:
            lines.append(f"[{'dir' if is_dir else 'file'}] {entry_uri}")
        else:
            lines.append(f"[{'dir' if is_dir else 'file'}] {name}")
    return "\n".join(lines)


# -- remember --------------------------------------------------------------


class StoreMessage(BaseModel):
    role: Literal["user", "assistant"] = Field(description="Message role")
    content: str = Field(description="Message text content")


@mcp.tool()
async def remember(messages: list[StoreMessage]) -> str:
    """Store information into OpenViking long-term memory. Use when the user says 'remember this', shares preferences, important facts, or decisions worth persisting."""
    import uuid

    from openviking.message.part import TextPart

    service = get_service()
    ctx = _get_ctx()
    session_id = f"mcp-store-{uuid.uuid4().hex[:12]}"
    session = await service.sessions.get(session_id, ctx, auto_create=True)
    for msg in messages:
        if msg.content:
            session.add_message(
                msg.role,
                [TextPart(text=msg.content)],
            )
    await service.sessions.commit_async(session_id, ctx)
    return f"Stored {len(messages)} message(s) and committed for memory extraction."


# -- add_resource ----------------------------------------------------------


_DEFAULT_UPLOAD_TTL_SECONDS = 600


def _resolve_public_base_url() -> tuple[str, str]:
    """Pick the URL the agent should POST uploads to. Returns ``(base_url, source)``.

    Resolution order (first match wins):

    1. ``env`` — ``OPENVIKING_PUBLIC_BASE_URL`` environment variable. Operator-set,
       always wins.
    2. ``config`` — ``ServerConfig.public_base_url``. Operator-set baseline in ov.conf.
    3. ``forwarded`` — ``X-Forwarded-Host`` (+ ``X-Forwarded-Proto``) from the request.
       Set by reverse proxies (nginx, ALB, ingress controllers, MCP proxies). Reliable
       when the proxy chain forwards these headers, which is the standard default.
    4. ``host`` — the raw ``Host`` header from a direct connection. Reliable for
       same-host MCP clients (e.g. local Claude Code talking to localhost server).
    5. ``listen`` — ``http://{listen_host}:{listen_port}`` last-resort fallback.
       Only produces an agent-reachable URL when the server is bound to a routable
       address; commonly wrong behind reverse proxies.

    Sources 1 and 2 are "explicit" — operator vouched for the URL. Sources 3-5 are
    inferred and may be wrong when the proxy chain doesn't forward request headers.
    Callers should append a "set OPENVIKING_PUBLIC_BASE_URL if upload fails" hint
    in that case.
    """
    env_url = os.environ.get("OPENVIKING_PUBLIC_BASE_URL")
    if env_url:
        return env_url.rstrip("/"), "env"
    config = get_server_config()
    if config is not None and config.public_base_url:
        return config.public_base_url.rstrip("/"), "config"

    url_info = _request_url_ctx.get()
    if url_info:
        # X-Forwarded-Host / -Proto can be comma-separated lists when the request
        # crosses multiple proxy hops. Take the first (left-most original-client)
        # value, matching the normalization in openviking.server.oauth.router.
        def _first(value: Optional[str]) -> Optional[str]:
            if not value:
                return None
            head = value.split(",", 1)[0].strip()
            return head or None

        xfh = _first(url_info.get("x_forwarded_host"))
        xfp = _first(url_info.get("x_forwarded_proto"))
        host_hdr = _first(url_info.get("host"))
        if xfh:
            proto = xfp or "https"
            return f"{proto}://{xfh}", "forwarded"
        if host_hdr:
            return f"http://{host_hdr}", "host"

    if config is not None:
        from openviking.server.config import map_bind_host_to_loopback

        host = map_bind_host_to_loopback(config.host)
        return f"http://{host}:{config.port}", "listen"
    return "http://127.0.0.1:1933", "listen"


async def _maybe_sitemap_hint(path: str) -> str:
    """Best-effort sitemap/RSS suggestion for a single-page add (never crawls).

    Policy: only suggest when the added URL is the site root / homepage (path is
    empty or "/"). Adding a deep article URL implies intent for just that page, and
    gating to the root also keeps the (bounded, non-recursive) probe from re-running
    when many pages of the same site are added one-by-one.

    Gated by config (``webfeed.suggest_feed``) and fully exception-safe so it can
    never block or fail the add it is attached to.
    """
    try:
        from urllib.parse import urlparse

        if urlparse(path).path not in ("", "/"):
            return ""

        from openviking.parse.accessors import discover_feed_hint
        from openviking.utils.network_guard import ensure_public_remote_target
        from openviking_cli.utils.config import get_openviking_config

        webfeed = getattr(get_openviking_config(), "webfeed", None)
        if webfeed is not None and not getattr(webfeed, "suggest_feed", True):
            return ""
        timeout = float(getattr(webfeed, "suggest_timeout", 2.5)) if webfeed else 2.5
        hint = await discover_feed_hint(
            path, timeout=timeout, request_validator=ensure_public_remote_target
        )
        return hint or ""
    except Exception:
        return ""


@mcp.tool()
async def add_resource(
    path: str = "",
    temp_file_id: str = "",
    description: str = "",
    watch_interval: float = 0,
    to: str = "",
    args: Optional[dict[str, Any]] = None,
) -> str:
    """Add a resource to OpenViking. Asynchronous — processing happens in the background.

    Remote URL: pass ``path`` as an http(s)://, git@, ssh://, or git:// URL. A sitemap /
    RSS / Atom URL ingests the WHOLE site as one resource tree; pass ``args={"site": true}``
    to force whole-site ingestion from a bare domain.

    Local file: pass ``path`` as a filesystem path (e.g. ``/tmp/foo.pdf``). The response is an
    upload instruction — HTTP POST the file to the returned URL and the server ingests it
    automatically; you do NOT need to call this tool again.

    Args:
        path: Remote URL or local filesystem path. Required unless ``temp_file_id`` is set.
        temp_file_id: Server-minted upload id from a prior signed local-file upload.
        description: Optional human-readable reason for adding the resource.
        watch_interval: Auto-refresh cadence in minutes. 0 = no watch. Prefer >=1440 (24h)
            unless the source changes faster — every refresh re-embeds the whole resource.
            Only applies to remote-URL invocations.
        to: Target URI under viking://resources/ (e.g. "viking://resources/volcengine/OpenViking").
            Leave empty to derive a URI from the source.
        args: Parser-specific options, e.g. {"feishu_access_token": "..."} for Feishu imports,
            or {"site": true} for whole-site ingestion.
    """
    from openviking.server.local_input_guard import require_remote_resource_source

    service = get_service()
    ctx = _get_ctx()

    if watch_interval < 0:
        return (
            "Error: watch_interval must be >= 0. Use 0 for one-shot add (no watch); "
            "use a positive number of minutes (>=1440 recommended) to subscribe to auto-refresh."
        )

    # Branch 1: ingest by temp_file_id. Kept for backward compat / REST-style use — the
    # signed upload now auto-ingests server-side, so agents no longer need this second leg.
    if temp_file_id:
        from openviking.server.config import ServerConfig

        server_config = get_server_config() or ServerConfig()
        store = TempUploadStore.build(server_config)
        try:
            result = await ingest_temp_upload(
                store, temp_file_id, ctx, to=to, reason=description, args=args
            )
        except (PermissionDeniedError, InvalidArgumentError) as exc:
            return f"Error: {exc}"
        except Exception as exc:
            return f"Error adding resource: {exc}"
        # add_resource returns a business-error dict (no raise) for parse/finalize failures;
        # surface it instead of reporting a false success.
        if isinstance(result, dict) and result.get("status") == "error":
            errors = result.get("errors")
            detail = (
                result.get("message")
                or (errors[0] if isinstance(errors, list) and errors else None)
                or "resource processing failed"
            )
            return f"Error adding resource: {detail}"
        root_uri = result.get("root_uri", "") if isinstance(result, dict) else ""
        return (
            f"Resource added: {root_uri}"
            if root_uri
            else "Resource added (processing in background)."
        )

    if not path:
        return "Error: provide either 'path' (remote URL or local file) or 'temp_file_id'."

    # Branch 2: agent passed a temp_file_id-shaped string as `path` — guide them
    if TEMP_FILE_ID_RE.match(path):
        return (
            f"Error: '{path}' looks like a temp_file_id, not a path. "
            f'Pass it as the temp_file_id kwarg: add_resource(temp_file_id="{path}")'
        )

    # Branch 3: remote URL — same flow as before
    if is_remote_resource_source(path):
        try:
            path = require_remote_resource_source(path)
            result = await service.resources.add_resource(
                path=path,
                ctx=ctx,
                to=to or None,
                reason=description,
                wait=False,
                watch_interval=watch_interval,
                enforce_public_remote_targets=True,
                args=args,
            )
        except Exception as exc:
            return f"Error adding resource: {exc}"
        root_uri = result.get("root_uri", "")
        if watch_interval > 0:
            watch_suffix = f" (watch enabled, refresh every {watch_interval:g} minute(s))"
        else:
            watch_suffix = ""
        message = (
            f"Resource added: {root_uri}{watch_suffix}"
            if root_uri
            else f"Resource added (processing in background){watch_suffix}."
        )
        # Detect-and-suggest: if this single page belongs to a site that exposes a
        # sitemap/RSS feed, hint at whole-site ingestion. Never auto-crawls; the
        # add above is already done, so a slow/failed probe has no functional impact.
        hint = await _maybe_sitemap_hint(path)
        if hint:
            message += "\n" + hint
        return message

    # Branch 4: local path — mint token, return upload instruction
    server_config = get_server_config()
    ttl_seconds = (
        server_config.upload_signed_ttl_seconds
        if server_config is not None
        else _DEFAULT_UPLOAD_TTL_SECONDS
    )

    token, expires_at = upload_token_store.issue(
        ctx.user.account_id,
        ctx.user.user_id,
        ttl_seconds=ttl_seconds,
        to=to,
        reason=description,
        actor_peer_id=ctx.actor_peer_id or "",
    )
    base_url, url_source = _resolve_public_base_url()
    upload_url = f"{base_url}/api/v1/resources/temp_upload?token={quote(token, safe='')}"
    expires_iso = datetime.fromtimestamp(expires_at, tz=timezone.utc).isoformat(timespec="seconds")
    minutes = max(1, ttl_seconds // 60)

    prose = (
        "Local file detected — upload it to ingest this resource.\n"
        "\n"
        'HTTP POST the file bytes (multipart/form-data, field name "file") to:\n'
        "\n"
        f"  {upload_url}\n"
        "\n"
        "The URL's token authorizes the upload (no API key needed); the server ingests "
        "the file automatically once received — you do NOT need to call add_resource again.\n"
        "\n"
        f"This upload URL expires in ~{minutes} minutes ({expires_iso})."
    )

    if url_source not in ("env", "config"):
        prose += (
            "\n\n"
            "Note for the user: this upload URL was auto-detected from the incoming "
            "request because OPENVIKING_PUBLIC_BASE_URL is not set on the server. "
            "If the upload fails (connection refused, wrong host, TLS error), ask the "
            "server operator to set OPENVIKING_PUBLIC_BASE_URL to the agent-facing "
            "URL of the OpenViking server (e.g. via docker-compose `environment:` "
            "or systemd unit) and retry."
        )

    return prose


# -- watch management ------------------------------------------------------
# MCP exposes the minimum closure: list + cancel. Pause/resume/trigger and
# the unified `update` verb are intentionally NOT exposed — they're either
# low-value for agents or invite unwanted autonomous decisions. Power users
# should reach for the REST API or the `ov task watch *` CLI (`pause`,
# `resume`, `trigger`, `update --interval`, etc.) for those operations.


@mcp.tool()
async def list_watches() -> str:
    """List watch tasks (auto-refresh subscriptions) visible to the current user."""
    service = get_service()
    ctx = _get_ctx()
    scheduler = getattr(service, "watch_scheduler", None)
    if scheduler is None or not scheduler.is_running:
        return "Error: Watch scheduler not running"
    wm = scheduler.watch_manager
    if wm is None:
        return "Error: Watch scheduler not running"
    # get_all_tasks does not raise PermissionDeniedError — it silently filters
    # tasks the caller cannot see (watch_manager.py:596-624), so we just
    # accept the filtered list.
    tasks = await wm.get_all_tasks(
        ctx.account_id,
        ctx.user.user_id,
        str(ctx.role),
        active_only=False,
    )
    if not tasks:
        return "No watch tasks."
    lines = []
    for t in tasks:
        status = "active" if t.is_active else "paused"
        nxt = t.next_execution_time.isoformat() if t.next_execution_time else "n/a"
        lines.append(
            f"- {t.to_uri or '(no uri)'}  interval={t.watch_interval:g}m  {status}  next={nxt}"
        )
    return "\n".join(lines)


@mcp.tool()
async def cancel_watch(to_uri: str) -> str:
    """Cancel a watch task by its target URI (e.g. "viking://resources/volcengine/OpenViking")."""
    from openviking.resource import watch_manager as _wm_mod

    service = get_service()
    ctx = _get_ctx()
    scheduler = getattr(service, "watch_scheduler", None)
    if scheduler is None or not scheduler.is_running:
        return "Error: Watch scheduler not running"
    wm = scheduler.watch_manager
    if wm is None:
        return "Error: Watch scheduler not running"
    task = await wm.get_task_by_uri(
        to_uri,
        ctx.account_id,
        ctx.user.user_id,
        str(ctx.role),
    )
    if task is None:
        return f"No watch task found for {to_uri}"
    try:
        # Return value (bool) is intentionally ignored: delete_task returns
        # False only when the task was removed between our lookup and the
        # delete call (a concurrent cancel from another caller). In that case
        # the post-condition the caller wanted ("no watch on this URI") still
        # holds, so we report the same success message either way. Permission
        # errors still surface via the explicit except below.
        _ = await wm.delete_task(
            task.task_id,
            ctx.account_id,
            ctx.user.user_id,
            str(ctx.role),
        )
    except _wm_mod.PermissionDeniedError:
        return f"Permission denied for {to_uri}"
    return f"Watch cancelled: {to_uri}"


# -- grep ------------------------------------------------------------------


@mcp.tool()
async def grep(
    uri: str, pattern: str | list[str], case_insensitive: bool = False, node_limit: int = 10
) -> str:
    """Search content in viking:// files using regex patterns (like grep). Supports multiple patterns searched concurrently. Use this for exact text matching; use the search tool for semantic retrieval."""
    import asyncio

    service = get_service()
    ctx = _get_ctx()
    patterns = [pattern] if isinstance(pattern, str) else pattern
    semaphore = asyncio.Semaphore(10)

    async def _grep_one(p: str) -> tuple[str, list[dict]]:
        async with semaphore:
            try:
                result = await service.fs.grep(
                    uri,
                    p,
                    ctx=ctx,
                    case_insensitive=case_insensitive,
                    node_limit=node_limit,
                )
                return (p, result.get("matches", []))
            except Exception:
                return (p, [])

    results = await asyncio.gather(*[_grep_one(p) for p in patterns])

    merged: dict[str, list[tuple]] = {}
    total = 0
    for p, matches in results:
        total += len(matches)
        for m in matches:
            m_uri = m.get("uri", "?")
            merged.setdefault(m_uri, []).append((m.get("line", "?"), m.get("content", ""), p))

    if not merged:
        return f"No matches found for pattern(s): {', '.join(patterns)}"

    lines = [f"Found {total} match(es) across {len(patterns)} pattern(s):"]
    for m_uri, hits in merged.items():
        hits.sort(key=lambda x: int(x[0]) if str(x[0]).isdigit() else 0)
        lines.append(f"\n{m_uri}")
        for line_no, content, p in hits:
            lines.append(f"  L{line_no} [{p}]: {content}")
    return "\n".join(lines)


# -- glob ------------------------------------------------------------------


@mcp.tool()
async def glob(pattern: str, uri: str = "viking://", node_limit: int = 100) -> str:
    """Find viking:// files matching a glob pattern (e.g. **/*.md, *.py). Use this for filename matching; use the search tool for content-based retrieval."""
    service = get_service()
    ctx = _get_ctx()

    try:
        result = await service.fs.glob(pattern, ctx=ctx, uri=uri, node_limit=node_limit)
    except Exception as e:
        return f"Error: {e}"

    matches = result.get("matches", [])
    if not matches:
        return f"No files found matching: {pattern}"

    lines = [f"Found {len(matches)} file(s):"]
    for m in matches:
        m_uri = m.get("uri", str(m)) if isinstance(m, dict) else str(m)
        lines.append(f"  {m_uri}")
    return "\n".join(lines)


# -- forget ----------------------------------------------------------------


@mcp.tool()
async def forget(uri: str, recursive: bool = False) -> str:
    """Permanently delete a viking:// URI from OpenViking. Irreversible — confirm with user before calling."""
    service = get_service()
    ctx = _get_ctx()
    await service.fs.rm(uri, ctx=ctx, recursive=recursive)
    return f"Deleted: {uri}"


# -- code navigation -------------------------------------------------------


def _require_viking_uri(uri: str) -> Optional[str]:
    """Return error message if uri is not a viking:// URI, else None."""
    if not isinstance(uri, str) or not uri.startswith("viking://"):
        return (
            "Error: only viking:// URIs are supported; "
            "use add_resource to ingest local code as a viking:// resource first."
        )
    return None


@mcp.tool()
async def code_outline(uri: str) -> str:
    """Show symbol structure (classes, functions, methods, line ranges) of a viking:// source file without reading full content."""
    err = _require_viking_uri(uri)
    if err:
        return err
    service = get_service()
    ctx = _get_ctx()
    try:
        content = await service.fs.read(uri, ctx=ctx)
    except Exception as exc:
        return f"Error: failed to read {uri}: {exc}"
    if not isinstance(content, str):
        return f"Error: {uri} is not text"
    return outline_file(content, uri)


@mcp.tool()
async def code_search(query: str, uri: str) -> str:
    """Search symbol names (class/function/method) by substring across a viking:// code repository. Returns symbol type, class context, file URI, and line range. Scans up to 200 source files — narrow uri for deeper coverage."""
    err = _require_viking_uri(uri)
    if err:
        return err
    if not query:
        return "Error: empty query"

    service = get_service()
    ctx = _get_ctx()
    try:
        entries = await service.fs.ls(uri, ctx=ctx, recursive=True, output="original")
    except Exception as exc:
        return f"Error: failed to list {uri}: {exc}"

    code_uris, capped = filter_code_uris(entries or [])
    if not code_uris:
        return f"No supported source files found under {uri}"

    semaphore = asyncio.Semaphore(CODE_SEARCH_CONCURRENCY)

    async def _read(u: str) -> Optional[tuple[str, str]]:
        async with semaphore:
            try:
                body = await service.fs.read(u, ctx=ctx)
            except Exception as exc:
                logger.warning("code_search: read failed for %s: %s", u, exc)
                return None
            if isinstance(body, str):
                return body, u
            return None

    fetched = await asyncio.gather(*[_read(u) for u in code_uris])
    files = [pair for pair in fetched if pair is not None]
    result = search_symbols(query, files)
    if capped:
        result += "\n\n(scanning stopped at 200-file cap; narrow uri to search more)"
    return result


@mcp.tool()
async def code_expand(uri: str, symbol: str) -> str:
    """Return the full source of a single named symbol from a viking:// source file. symbol accepts 'bar' (top-level) or 'Foo.bar' (method)."""
    err = _require_viking_uri(uri)
    if err:
        return err
    if not symbol:
        return "Error: empty symbol"

    service = get_service()
    ctx = _get_ctx()
    try:
        content = await service.fs.read(uri, ctx=ctx)
    except Exception as exc:
        return f"Error: failed to read {uri}: {exc}"
    if not isinstance(content, str):
        return f"Error: {uri} is not text"
    return expand_symbol(content, uri, symbol)


# -- health ----------------------------------------------------------------


@mcp.tool()
async def health() -> str:
    """Check whether the OpenViking server is healthy."""
    try:
        service = get_service()
        return f"OpenViking is healthy (service initialized, storage: {type(service.viking_fs).__name__})"
    except Exception as e:
        return f"OpenViking is unhealthy: {e}"


# ---------------------------------------------------------------------------
# App factory + lifespan
# ---------------------------------------------------------------------------


@asynccontextmanager
async def mcp_lifespan():
    """Run the MCP session manager. Call this inside the FastAPI lifespan."""
    async with mcp.session_manager.run():
        logger.info(
            "MCP endpoint ready (13 tools: find, search, read, list, remember, add_resource, grep, glob, code_outline, code_search, code_expand, forget, health)"
        )
        yield


def create_mcp_app() -> ASGIApp:
    """Create the MCP ASGI app with identity middleware.

    IMPORTANT: call `mcp_lifespan()` inside the FastAPI lifespan BEFORE
    serving requests. The session manager task group must be initialized.
    """
    starlette_app = mcp.streamable_http_app()
    handler = starlette_app.routes[0].app
    return _IdentityASGIMiddleware(handler)
