"""Custom agent using haiku.rag composable toolsets.

Demonstrates how to compose search, QA, and document toolsets into a
pydantic-ai Agent using AgentDeps and prepare_context.

Requirements:
    - An Ollama instance running locally (default embedder)
    - An Anthropic API key (for the QA model) or adjust the model below
    - A haiku.rag database with documents already ingested

Usage:
    uv run python examples/custom_agent.py /path/to/db.lancedb
"""

import asyncio
import sys

from pydantic_ai import Agent

from haiku.rag.client import HaikuRAG
from haiku.rag.tools import (
    AgentDeps,
    ToolContext,
    create_document_toolset,
    create_qa_toolset,
    create_search_toolset,
    prepare_context,
)


async def main(db_path: str) -> None:
    async with HaikuRAG(db_path) as client:
        # Compose toolsets into an agent
        config = client.config
        search_toolset = create_search_toolset(config)
        qa_toolset = create_qa_toolset(config)
        document_toolset = create_document_toolset(config)

        agent = Agent(
            "anthropic:claude-haiku-4-5-20251001",
            deps_type=AgentDeps,
            output_type=str,
            instructions=(
                "You are a helpful assistant with access to a knowledge base. "
                "Use the available tools to answer questions."
            ),
            toolsets=[search_toolset, qa_toolset, document_toolset],
        )

        # Prepare a shared ToolContext
        context = ToolContext()
        prepare_context(context, features=["search", "documents", "qa"])

        deps = AgentDeps(client=client, tool_context=context)

        print("Custom agent ready. Ctrl+C to exit.\n")
        while True:
            try:
                user_input = input("You: ").strip()
            except (EOFError, KeyboardInterrupt):
                break

            if not user_input:
                continue

            result = await agent.run(user_input, deps=deps)
            print(f"\nAgent: {result.output}\n")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <db_path>")
        sys.exit(1)
    asyncio.run(main(sys.argv[1]))