"""
Built-in tool: kb_search

Searches the tenant's knowledge base using pgvector cosine similarity.

The query is embedded using the same all-MiniLM-L6-v2 model as conversation
embeddings (vector(384)), then matched against kb_chunks via HNSW ANN search.

The tool accepts an optional tenant_id and agent_id via kwargs — these are
injected by the executor using context that the LLM doesn't provide directly.
"""

from __future__ import annotations

import logging

logger = logging.getLogger(__name__)

_TOP_K = 3


async def kb_search(
    query: str,
    tenant_id: str | None = None,
    agent_id: str | None = None,
    **kwargs: object,
) -> str:
    """
    Search the knowledge base for content relevant to the query.

    Args:
        query:     The search query string.
        tenant_id: Injected by executor — the current tenant UUID string.
        agent_id:  Injected by executor — the current agent UUID string.

    Returns:
        Formatted string with top matching KB chunks, or a message if the KB
        is empty or search fails.
    """
    if not tenant_id:
        return "Knowledge base search unavailable: tenant context not set."

    try:
        import uuid as _uuid

        from shared.db import async_session_factory, engine
        from shared.rls import configure_rls_hook, current_tenant_id
        from orchestrator.memory.embedder import embed_text

        from sqlalchemy import text

        tenant_uuid = _uuid.UUID(tenant_id)
        query_embedding = embed_text(query)

        # Format embedding as PostgreSQL vector literal
        embedding_str = "[" + ",".join(str(x) for x in query_embedding) + "]"

        configure_rls_hook(engine)
        token = current_tenant_id.set(tenant_uuid)
        try:
            async with async_session_factory() as session:
                result = await session.execute(
                    text("""
                        SELECT content, chunk_index,
                               embedding <=> CAST(:embedding AS vector) AS distance
                        FROM kb_chunks
                        WHERE tenant_id = :tenant_id
                        ORDER BY embedding <=> CAST(:embedding AS vector)
                        LIMIT :top_k
                    """),
                    {
                        "embedding": embedding_str,
                        "tenant_id": str(tenant_uuid),
                        "top_k": _TOP_K,
                    },
                )
                rows = result.fetchall()
        finally:
            current_tenant_id.reset(token)

    except Exception:
        logger.exception("KB search failed for tenant=%s", tenant_id)
        return "Knowledge base search encountered an error. Please try again."

    if not rows:
        return f"No relevant knowledge base content found for: {query}"

    lines = [f"Knowledge base results for: {query}\n"]
    for i, row in enumerate(rows, start=1):
        lines.append(f"{i}. {row.content}\n")

    return "\n".join(lines)
