from __future__ import annotations

import os

import smolagents.models
from smolagents import ToolCallingAgent

# ---------------- Agent Prompts ----------------
GENERATOR_INSTRUCTIONS = """
You are the Generator. Your goal is to produce a concise draft that strictly satisfies the caller's constraints.
Use the managed agent named "critic_agent" to iteratively improve your draft.
When the draft is complete and ready, return it to the user with the final answer tool.
"""

CRITIC_INSTRUCTIONS = """
You are the Critic. Your job is to provide constructive, actionable feedback on drafts produced by the Generator.
You should not iterate or make multiple tool calls.
Instead, simply call the final answer tool with your evaluation and feedback.
"""


# ---------------- Factory ----------------
def generate_generator_with_managed_critic(
    *,
    gen_max_steps: int = 4,
    crt_max_steps: int = 1,
) -> 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"),
    )

    critic_agent = ToolCallingAgent(
        tools=[],  # critic needs no tools; it just returns JSON text
        model=model,
        name="critic_agent",
        description="Evaluates drafts against constraints and returns a compact set of recommendations report.",
        instructions=CRITIC_INSTRUCTIONS,
        add_base_tools=False,
        max_steps=crt_max_steps,
    )

    generator = ToolCallingAgent(
        tools=[],  # keep toolbox minimal
        model=model,
        name="generator_with_managed_critic",
        description="Generates drafts and iteratively improves them using a managed critic agent.",
        managed_agents=[critic_agent],  # <-- critic attached here
        instructions=GENERATOR_INSTRUCTIONS,
        add_base_tools=False,
        max_steps=gen_max_steps,
    )
    return generator