"""Vector/semantic search tools."""

from pydantic_ai import RunContext
from pydantic_ai.toolsets import FunctionToolset

from pydanticai_multiagent.dependencies import SearchDeps

vector_search_toolset: FunctionToolset[SearchDeps] = FunctionToolset()


@vector_search_toolset.tool
async def search_knowledge_base(
    ctx: RunContext[SearchDeps],
    query: str,
    top_k: int = 5,
    category: str | None = None,
) -> str:
    """Search the internal knowledge base using semantic search.

    Args:
        ctx: The run context with dependencies.
        query: Natural language search query.
        top_k: Number of results to return (1-20).
        category: Optional category filter.
    """
    filters = {"category": category} if category else None

    results = await ctx.deps.vector_store.search(
        query=query,
        top_k=min(top_k, 20),
        filters=filters,
    )

    if not results:
        return f"No documents found matching: {query}"

    formatted = []
    for i, doc in enumerate(results, 1):
        title = doc.get("title", "Untitled")
        content = doc.get("content", "")[:500]
        score = doc.get("score", 0)
        formatted.append(f"{i}. {title} (relevance: {score:.2f})")
        formatted.append(f"   {content}")

    return "\n".join(formatted)


@vector_search_toolset.tool
async def get_document(
    ctx: RunContext[SearchDeps],
    doc_id: str,
) -> str:
    """Retrieve a specific document by its ID.

    Args:
        ctx: The run context with dependencies.
        doc_id: The unique document identifier.
    """
    doc = await ctx.deps.vector_store.get(doc_id)

    if not doc:
        return f"Document not found: {doc_id}"

    title = doc.get("title", "Untitled")
    content = doc.get("content", "")
    metadata = doc.get("metadata", {})

    meta_str = ", ".join(f"{k}: {v}" for k, v in metadata.items())
    return f"# {title}\n\nMetadata: {meta_str}\n\n{content}"


@vector_search_toolset.tool
async def find_similar_documents(
    ctx: RunContext[SearchDeps],
    doc_id: str,
    top_k: int = 5,
) -> str:
    """Find documents similar to a given document.

    Args:
        ctx: The run context with dependencies.
        doc_id: The document ID to find similar documents for.
        top_k: Number of similar documents to return.
    """
    source_doc = await ctx.deps.vector_store.get(doc_id)
    if not source_doc:
        return f"Source document not found: {doc_id}"

    # Use the source document's content for similarity search
    content = source_doc.get("content", "")
    results = await ctx.deps.vector_store.search(query=content, top_k=top_k + 1)

    # Filter out the source document itself
    similar = [r for r in results if r.get("id") != doc_id][:top_k]

    if not similar:
        return "No similar documents found."

    formatted = []
    for i, doc in enumerate(similar, 1):
        title = doc.get("title", "Untitled")
        score = doc.get("score", 0)
        formatted.append(f"{i}. {title} (similarity: {score:.2f})")

    return "\n".join(formatted)
