"""Search tool for agent."""

from langchain_core.tools import tool
from src.rag.retriever import retriever
from sqlalchemy.ext.asyncio import AsyncSession
from src.middlewares.logging import get_logger

logger = get_logger("search_tool")


@tool
async def search_documents(
    query: str,
    user_id: str,
    db: AsyncSession,
    num_results: int = 5
) -> str:
    """Search user's uploaded documents for relevant information.

    Args:
        query: The search query or question
        user_id: The user's ID
        db: Database session
        num_results: Number of results to return (default: 5)

    Returns:
        Relevant document excerpts with source and page information
    """
    try:
        results = await retriever.retrieve(query, user_id, db, num_results)

        if not results:
            return "No relevant information found in the documents."

        formatted_results = []
        for result in results:
            filename = result["metadata"].get("filename", "Unknown")
            page = result["metadata"].get("page_label")
            source_label = f"{filename}, p.{page}" if page else filename
            formatted_results.append(f"[Source: {source_label}]\n{result['content']}\n")

        return "\n".join(formatted_results)

    except Exception as e:
        logger.error("Search failed", error=str(e))
        return "Sorry, I encountered an error while searching the documents."