"""The copilot as an MCP server (MCP Python SDK v2, lessons 13–14).

    pip install "mcp[cli]"                       # not installed by default: optional dependency
    python -m copilot.mcp_server                 # stdio, for Claude Desktop / Claude Code / IDEs
    python -m copilot.mcp_server --http          # Streamable HTTP on http://127.0.0.1:8001/mcp
    mcp dev copilot/mcp_server.py                # MCP Inspector (run from the project folder)

Tools:     search_docs(query, course?, k?) · ask(question)
Resource:  lesson://{course}/{name}          e.g. lesson://sql/06-order-limit.html
Prompt:    explain(topic)
"""

from __future__ import annotations

import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated

if __package__ in (None, ""):                           # `mcp dev copilot/mcp_server.py` runs this as a file
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from mcp.server import MCPServer
from mcp.server.mcpserver import Context
from mcp.server.mcpserver.exceptions import ResourceNotFoundError, ToolError
from pydantic import BaseModel, Field

from copilot import Copilot, build


@asynccontextmanager
async def lifespan(server: MCPServer) -> AsyncIterator[Copilot]:
    copilot = build()                                   # index built once, shared by every request
    try:
        yield copilot
    finally:
        copilot.tel.shutdown()


mcp = MCPServer(
    "docs-copilot", lifespan=lifespan,
    instructions="Search and read the learn-with-project lessons (PySpark, SQL, MySQL, Python, DSA, "
                 "Databricks, OpenAI SDK, AI Engineering and more). Search first, then cite lesson URLs.")


class SearchHit(BaseModel):
    url: str = Field(description="Lesson section; the page part is usable with the lesson:// resource")
    title: str
    section: str
    snippet: str


class CopilotAnswer(BaseModel):
    answer: str
    grounded: bool = Field(description="True only if every citation refers to a retrieved source")
    sources: list[str]
    trace_id: str = Field(description="OpenTelemetry trace id, for debugging and feedback")


def _copilot(ctx: Context) -> Copilot:
    return ctx.request_context.lifespan_context


@mcp.tool()
def search_docs(query: Annotated[str, Field(description="What to look for, in plain words")],
                ctx: Context,
                course: Annotated[str | None, Field(description="Course folder, e.g. 'sql'")] = None,
                k: Annotated[int, Field(ge=1, le=10)] = 5) -> list[SearchHit]:
    """Search the course lessons (keyword + vector hybrid). Returns sections with URLs and snippets."""
    copilot = _copilot(ctx)
    courses = {c.course for c in copilot.index.chunks}
    if course is not None and course not in courses:
        raise ToolError(f"Unknown course {course!r}. Valid values: {sorted(courses)}")
    return [SearchHit(url=h.chunk.url, title=h.chunk.title, section=h.chunk.section, snippet=h.chunk.text[:200])
            for h in copilot.index.search(query, k=k, course=course)]


@mcp.tool()
def ask(question: Annotated[str, Field(min_length=3, max_length=2000)], ctx: Context) -> CopilotAnswer:
    """Answer a question from the lessons with verified citations, or say I don't know."""
    a = _copilot(ctx).ask(question, user_id="mcp")
    return CopilotAnswer(answer=a.text, grounded=a.grounded, sources=[c.url for c in a.sources], trace_id=a.trace_id)


@mcp.resource("lesson://{course}/{name}")
def lesson(course: str, name: str, ctx: Context) -> str:
    """Full text of one lesson, e.g. lesson://sql/06-order-limit.html"""
    parts = [c.text for c in _copilot(ctx).index.chunks if c.course == course and c.lesson == name]
    if not parts:
        raise ResourceNotFoundError(f"No lesson {course}/{name}")
    return "\n\n".join(parts)


@mcp.prompt()
def explain(topic: str) -> str:
    """Explain a topic from the courses for a beginner, with citations."""
    return (f"Use search_docs to find lessons about {topic!r}, read the best one with the lesson resource, "
            f"then explain it for a beginner in under 150 words and cite the lesson URLs.")


if __name__ == "__main__":
    if "--http" in sys.argv:
        mcp.run(transport="streamable-http", port=8001)
    else:
        mcp.run()