"""Tool implementations for PKB. Imported by server.py."""

from pydantic import BaseModel, Field
from mcp.shared.exceptions import McpError
from mcp.types import ErrorData, INVALID_PARAMS

# Per-user note store. NOTES (legacy name) views into the default user's bucket
# so prior lesson code keeps working in the lab.
USER_NOTES: dict[str, dict[str, str]] = {
    "local": {
        "hello": "Hello note body",
        "todo": "Pick up groceries",
        "ideas": "Build an MCP server tutorial",
    }
}

# Legacy single-tenant alias used by L04-L08 examples. Points at the "local" bucket.
NOTES = USER_NOTES["local"]


def export_notes(user: str) -> dict[str, str]:
    """Return a copy of all notes for USER. Used by the export tool."""
    return dict(USER_NOTES.get(user, {}))


class SearchArgs(BaseModel):
    query: str = Field(..., description="case-insensitive substring match")
    limit: int = Field(10, ge=1, le=100, description="max number of hits")


def search(args: SearchArgs) -> list[str]:
    """Search note titles and bodies by case-insensitive substring."""
    q = args.query.lower()
    hits = [k for k, v in NOTES.items() if q in v.lower() or q in k.lower()]
    return hits[: args.limit]


from mcp.server.fastmcp import Context
from mcp.types import SamplingMessage, TextContent


async def summarize_via_sampling(note_id: str, ctx: Context) -> str:
    """Summarize a note by asking the client's LLM via sampling.

    Raises McpError(INVALID_PARAMS) if note_id is unknown.
    """
    if note_id not in NOTES:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message=f"unknown note id: {note_id}",
        ))
    body = NOTES[note_id]
    result = await ctx.session.create_message(
        messages=[SamplingMessage(
            role="user",
            content=TextContent(
                type="text",
                text=f"Summarize the following note in 2 sentences:\n\n{body}",
            ),
        )],
        max_tokens=200,
    )
    # result.content is a TextContent or ImageContent; we requested text.
    if isinstance(result.content, TextContent):
        return result.content.text
    return "(sampling returned non-text content)"


class WriteArgs(BaseModel):
    title: str = Field(..., min_length=1, max_length=120,
                       description="note title (must be unique unless overwrite=true)")
    body: str = Field(..., max_length=10_000, description="note body")
    overwrite: bool = Field(False, description="allow replacing an existing note")


def fetch_note(note_id: str) -> str:
    """Return the body of one note by id. Tool-friendly counterpart to the resource."""
    if note_id not in NOTES:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message=f"unknown note id: {note_id}",
        ))
    return NOTES[note_id]


def write_note(args: WriteArgs) -> str:
    """Save a note. Refuses to overwrite unless args.overwrite is True."""
    if args.title in NOTES and not args.overwrite:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message=f"note exists: {args.title}; pass overwrite=true to replace",
        ))
    NOTES[args.title] = args.body
    return f"saved: {args.title}"


async def summarize_and_save(note_id: str, ctx: Context) -> str:
    """Compose summarize + write: sample, then save as `{id}-summary`."""
    summary = await summarize_via_sampling(note_id, ctx)
    new_title = f"{note_id}-summary"
    NOTES[new_title] = summary
    return f"saved: {new_title}"