from __future__ import annotations

import os

import smolagents.models
from smolagents import ToolCallingAgent
import agents.code_writing_agents
import agents.agent_with_custom_beam_design_tools

INSTRUCTIONS = (
    "You are a manager agent with access to two specialized subagents."
    "Your primary task is to delegate tasks to the appropriate subagent based on their expertise."
)
AGENTS = [
    agents.code_writing_agents.generate_code_writing_agent_with_search(),
    agents.agent_with_custom_beam_design_tools.generate_beam_agent(),
]


# ---------------- Factory ----------------
def generate_manager_with_heterogeneous_agents() -> ToolCallingAgent:
    """
    Returns a CodeAgent (generator) that manages a critic sub-agent.
    The critic is exposed to the generator as a managed agent (callable like a tool).
    """

    model = smolagents.models.OpenAIServerModel(
        model_id=os.getenv("AGENT_MODEL", ""),
        api_base=os.getenv("UPSTREAM_OPENAI_BASE", "").rstrip("/"),
        api_key=os.getenv("OPENAI_API_KEY"),
    )

    generator = ToolCallingAgent(
        tools=[],  # keep toolbox minimal
        model=model,
        name="manager_with_heterogeneous_agents",
        managed_agents=AGENTS,
        instructions=INSTRUCTIONS,
        add_base_tools=False,
        max_steps=2,
    )
    return generator