from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext

from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.sandbox import Sandbox
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import (
    build_document_filter,
    combine_filters,
)


class AnalysisResult(BaseModel):
    """Result from the analysis toolset (RLM execution)."""

    answer: str = Field(description="The answer produced by analysis")
    code_executed: bool = Field(
        default=True,
        description="Whether code was executed to produce this answer",
    )


def create_analysis_toolset(
    config: AppConfig,
    base_filter: str | None = None,
    tool_name: str = "analyze",
) -> FunctionToolset[RAGDeps]:
    """Create a toolset with code analysis capabilities via RLM agent.

    Args:
        config: Application configuration.
        base_filter: Optional base SQL WHERE clause applied to searches.
        tool_name: Name for the analyze tool. Defaults to "analyze".

    Returns:
        FunctionToolset with an analyze tool.
    """

    async def analyze(  # pragma: no cover
        ctx: RunContext[RAGDeps],
        task: str,
        document_name: str | None = None,
    ) -> AnalysisResult:
        """Execute a computational task via code execution.

        Uses the RLM (Recursive Language Model) agent to write and execute
        Python code to answer the task.

        Args:
            task: A specific, actionable instruction describing what to compute.
            document_name: Optional document name/title to focus on.

        Returns:
            AnalysisResult with answer and execution metadata.
        """
        client = ctx.deps.client

        doc_filter = build_document_filter(document_name) if document_name else None
        effective_filter = combine_filters(base_filter, doc_filter)

        rlm_context = RLMContext(filter=effective_filter)

        sandbox = Sandbox(
            client=client,
            config=config,
            context=rlm_context,
        )
        deps = RLMDeps(
            sandbox=sandbox,
            context=rlm_context,
        )

        rlm_agent = create_rlm_agent(config)
        result = await rlm_agent.run(task, deps=deps)

        program = result.output.program

        return AnalysisResult(
            answer=result.output.answer,
            code_executed=bool(program),
        )

    toolset: FunctionToolset[RAGDeps] = FunctionToolset()
    toolset.add_function(analyze, name=tool_name)
    return toolset