from __future__ import annotations

import os

import smolagents.models
from smolagents import CodeAgent, ToolCallingAgent

# ---------------- Agent Prompts ----------------
GENERATOR_INSTRUCTIONS = """
You are the Generator/Refiner.

Goal
- Produce a concise draft that strictly satisfies the caller's constraints.
- Use the managed agent named "critic_agent" to validate your draft.
- Repeat: draft → call critic_agent → if ok, return draft; else revise and re-check.
- Output ONLY the final draft text (no JSON, no commentary).

Constraints to enforce in every revision:
- Respect the maximum word count.
- Use bullet points where appropriate (lines starting with '-' or '•').
- Include all required phrases verbatim.
- End with a line starting with 'Next steps:'.

Implementation guidance (you write/run the code):
- Use a small loop with a bounded number of rounds provided by the caller.
- Call critic_agent with a single string payload that contains:
  - the DRAFT
  - the list of required phrases
  - the word limit
  - any other constraints you need
- critic_agent returns a JSON object with keys: ok (bool), violations (list[str]), suggestions (str).
- If ok == true, immediately return the current draft (and stop).
- Otherwise, revise the draft to fix ALL violations, then re-check.

Important:
- Do NOT print anything except the final draft at the end.
- Avoid verbose interleaved logs; keep code minimal.
"""

CRITIC_INSTRUCTIONS = """
You are the Critic.

Input
- A single string payload that includes:
  - DRAFT text
  - explicit constraints (e.g., max_words, must_include list, bullets rule, ending 'Next steps:')

Task
- Evaluate the DRAFT against the constraints and general quality (clarity, correctness, structure, tone).
- Return ONLY compact JSON (no text outside JSON) with shape:
{
  "ok": true|false,
  "violations": ["short, concrete issues..."],
  "suggestions": "one short paragraph of actionable guidance"
}

Rules
- ok=true ONLY if all constraints are satisfied AND the draft reads clearly.
- Keep violations terse and actionable.
- Keep suggestions short and prescriptive.
"""


# ---------------- Factory ----------------
def generate_generator_with_managed_critic(
    *,
    gen_max_steps: int = 12,
    crt_max_steps: int = 2,
) -> CodeAgent:
    """
    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 JSON report.",
        instructions=CRITIC_INSTRUCTIONS,
        add_base_tools=False,
        max_steps=crt_max_steps,
    )

    generator = CodeAgent(
        tools=[],  # keep toolbox minimal
        model=model,
        name="generator_with_managed_critic",
        managed_agents=[critic_agent],  # <-- critic attached here
        instructions=GENERATOR_INSTRUCTIONS,
        add_base_tools=False,
        max_steps=gen_max_steps,
    )
    return generator