"""Router agent for delegating to specialized agents.

The router uses dynamic toolset selection via @agent.toolset to include
domain-specific tools passed through deps.domain_toolset. This allows
users to extend the router with custom tools at runtime.

Example:
    ```python
    from pydanticai_multiagent import router_agent, create_mock_search_deps
    from my_domain import my_toolset

    deps = create_mock_search_deps()
    deps.domain_toolset = my_toolset  # Add domain-specific tools

    result = await router_agent.run("Use my domain tools", deps=deps)
    ```
"""

from pydantic_ai import Agent, RunContext
from pydantic_ai.toolsets import AbstractToolset, FunctionToolset

from pydanticai_multiagent.dependencies import SearchDeps

from .analyst import analyst_agent
from .code import code_agent
from .research import research_agent
from .support import support_agent
from .writer import writer_agent

router_agent: Agent[SearchDeps, str] = Agent(
    "openai:gpt-4o-mini",
    deps_type=SearchDeps,
    instructions="""You are a routing agent that directs user requests to specialized agents.

Available specialists:
1. Research Agent - For research questions, fact-finding, and information synthesis
2. Analyst Agent - For data analysis, statistics, trends, and business insights
3. Code Agent - For programming questions, code review, and debugging
4. Writer Agent - For content creation, emails, documentation, and copywriting
5. Support Agent - For customer service inquiries and account-related questions

Your job:
1. Analyze the user's request to understand their intent
2. Delegate to the most appropriate specialist agent OR use domain-specific tools if available
3. Return the specialist's response to the user

Guidelines:
- If domain-specific tools are available, use them for domain-related queries
- If the request spans multiple domains, break it into parts and delegate separately
- If unclear, ask the user to clarify before delegating
- You can delegate multiple times for complex requests
- Summarize results if needed for clarity
""",
)


@router_agent.toolset
def dynamic_domain_toolset(ctx: RunContext[SearchDeps]) -> AbstractToolset:
    """Dynamically include domain-specific tools based on deps.

    This allows users to extend the router with custom tools by setting
    `deps.domain_toolset` before running the agent.
    """
    if ctx.deps.domain_toolset is not None:
        return ctx.deps.domain_toolset
    # Return empty toolset if no domain tools configured
    return FunctionToolset()


@router_agent.tool
async def delegate_to_research(
    ctx: RunContext[SearchDeps],
    query: str,
) -> str:
    """Delegate a research question to the research agent.

    Args:
        ctx: The run context with dependencies.
        query: The research question to investigate.
    """
    result = await research_agent.run(query, deps=ctx.deps, usage=ctx.usage)
    return result.output.format_with_sources()


@router_agent.tool
async def delegate_to_analyst(
    ctx: RunContext[SearchDeps],
    query: str,
) -> str:
    """Delegate data analysis to the analyst agent.

    Args:
        ctx: The run context with dependencies.
        query: The analysis request describing what to analyze.
    """
    result = await analyst_agent.run(query, deps=ctx.deps.to_base_deps(), usage=ctx.usage)
    return result.output.model_dump_json(indent=2)


@router_agent.tool
async def delegate_to_code(
    ctx: RunContext[SearchDeps],
    query: str,
) -> str:
    """Delegate programming questions to the code agent.

    Args:
        ctx: The run context with dependencies.
        query: The programming question or code to review.
    """
    result = await code_agent.run(query, deps=ctx.deps.to_base_deps(), usage=ctx.usage)
    output = result.output

    response = f"**Explanation:**\n{output.explanation}"
    if output.code:
        response += f"\n\n**Code ({output.language or 'unknown'}):**\n```{output.language or ''}\n{output.code}\n```"
    if output.suggestions:
        response += "\n\n**Suggestions:**\n" + "\n".join(f"- {s}" for s in output.suggestions)

    return response


@router_agent.tool
async def delegate_to_writer(
    ctx: RunContext[SearchDeps],
    task: str,
    context: str | None = None,
) -> str:
    """Delegate content creation to the writer agent.

    Args:
        ctx: The run context with dependencies.
        task: Description of what to write (e.g., 'blog post about AI').
        context: Additional context or requirements for the content.
    """
    prompt = task
    if context:
        prompt = f"{task}\n\nContext: {context}"

    result = await writer_agent.run(prompt, deps=ctx.deps.to_base_deps(), usage=ctx.usage)
    return result.output


@router_agent.tool
async def delegate_to_support(
    ctx: RunContext[SearchDeps],
    query: str,
) -> str:
    """Delegate customer support inquiries to the support agent.

    Args:
        ctx: The run context with dependencies.
        query: The customer's question or issue.
    """
    # In production, roles/permissions would come from auth context
    result = await support_agent.run(query, deps=ctx.deps.to_auth_deps(), usage=ctx.usage)
    output = result.output

    response = output.answer
    if output.escalate:
        response += f"\n\n⚠️ **Escalated to human support:** {output.escalation_reason}"
    if output.follow_up_actions:
        response += "\n\n**Follow-up actions:**\n" + "\n".join(
            f"- {a}" for a in output.follow_up_actions
        )

    return response
