from dataclasses import dataclass
from typing import Any

from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware, ToolCallLimitMiddleware

from kbqa.agents.context import ChatRunContext, ResearchContext
from kbqa.agents.llm import create_chat_model
from kbqa.agents.middleware import RequireResearchMiddleware
from kbqa.agents.prompts import MAIN_PROMPT, RESEARCH_PROMPT
from kbqa.agents.tools import create_research_tool, search_knowledge_base
from kbqa.config import Settings


@dataclass(frozen=True, slots=True)
class AgentBundle:
    main_agent: Any
    research_agent: Any


def create_agent_bundle(settings: Settings) -> AgentBundle:
    model = create_chat_model(settings)
    research_agent = create_agent(
        model=model,
        tools=[search_knowledge_base],
        system_prompt=RESEARCH_PROMPT,
        context_schema=ResearchContext,
        middleware=[
            RequireResearchMiddleware(tool_name="search_knowledge_base"),
            ModelCallLimitMiddleware(run_limit=4, exit_behavior="error"),
            ToolCallLimitMiddleware(
                tool_name="search_knowledge_base",
                run_limit=settings.max_research_searches,
                exit_behavior="continue",
            ),
        ],
        name="research-agent",
    )
    research_tool = create_research_tool(research_agent)
    main_agent = create_agent(
        model=model,
        tools=[research_tool],
        system_prompt=MAIN_PROMPT,
        context_schema=ChatRunContext,
        middleware=[
            RequireResearchMiddleware(),
            ModelCallLimitMiddleware(run_limit=3, exit_behavior="error"),
            ToolCallLimitMiddleware(tool_name="research", run_limit=1, exit_behavior="error"),
        ],
        name="main-agent",
    )
    return AgentBundle(main_agent=main_agent, research_agent=research_agent)
