# ============================================================================
# STEP 0 — Prerequisites
# ============================================================================
# export OPENAI_API_KEY="sk-..."
# `uv add openai-agents`
# export RELAI_API_KEY="relai-..."
# Save as `stock-assistant.py`

import asyncio

from agents import Agent, Runner

from relai import (
    AgentOutputs,
    AsyncRELAI,
    AsyncSimulator,
    SimulationTape,
    random_env_generator,
)
from relai.critico import Critico
from relai.critico.evaluate import RELAIFormatEvaluator
from relai.maestro import Maestro, params, register_param
from relai.mocker import Persona
from relai.simulator import simulated

# ============================================================================
# STEP 1.1 — Decorate inputs/tools that will be simulated
# ============================================================================


@simulated
async def get_user_query() -> str:
    """Get user's query about stock prices."""
    # In a real agent, this function might get input from a chat interface.
    return input("Enter you stock query: ")


# ============================================================================
# STEP 1.2 — Register parameters for optimization
# ============================================================================

register_param(
    "prompt",
    type="prompt",
    init_value="You are a helpful assistant for stock price questions.",
    desc="system prompt for the agent",
)

# ============================================================================
# STEP 2 — Your agent core
# ============================================================================


async def agent_fn(tape: SimulationTape) -> AgentOutputs:
    # It is good practice to catch exceptions in agent function
    # especially if the agent might raise errors with different configs
    try:
        question = await get_user_query()
        agent = Agent(
            name="Stock assistant",
            instructions=params.prompt,  # access registered parameter
            model="gpt-5-mini",
        )
        result = await Runner.run(agent, question)
        tape.extras["format_rubrics"] = {"Prices must include cents (eg: $XXX.XX)": 1.0}
        tape.agent_inputs["question"] = question  # trace inputs for later auditing
        return {"summary": result.final_output}
    except Exception as e:
        return {"summary": str(e)}
    


async def main() -> None:
    # Set up your simulation environment
    # Bind Personas/MockTools to fully-qualified function names
    env_generator = random_env_generator(
        config_set={
            "__main__.get_user_query": [Persona(user_persona="A polite and curious user.")],
        }
    )

    async with AsyncRELAI() as client:
        # ============================================================================
        # STEP 3 — Simulate
        # ============================================================================
        simulator = AsyncSimulator(agent_fn=agent_fn, env_generator=env_generator, client=client)
        agent_logs = await simulator.run(num_runs=1)

        # ============================================================================
        # STEP 4 — Evaluate with Critico
        # ============================================================================
        critico = Critico(client=client)
        format_evaluator = RELAIFormatEvaluator(client=client)
        critico.add_evaluators({format_evaluator: 1.0})
        critico_logs = await critico.evaluate(agent_logs)

        # Publish evaluation report to the RELAI platform
        await critico.report(critico_logs)

        maestro = Maestro(client=client, agent_fn=agent_fn, log_to_platform=True, name="Stock assistant")
        maestro.add_setup(simulator=simulator, critico=critico)

        # ============================================================================
        # STEP 5.1 — Optimize configs with Maestro (the parameters registered earlier in STEP 2)
        # ============================================================================

        # params.load("saved_config.json")  # load previous params if available
        await maestro.optimize_config(
            total_rollouts=20,  # Total number of rollouts to use for optimization.
            batch_size=2,  # Base batch size to use for individual optimization steps. Defaults to 4.
            explore_radius=1,  # A positive integer controlling the aggressiveness of exploration during optimization.
            explore_factor=0.5,  # A float between 0 to 1 controlling the exploration-exploitation trade-off.
            verbose=True,  # If True, additional information will be printed during the optimization step.
        )
        params.save("saved_config.json")  # save optimized params for future usage

        # ============================================================================
        # STEP 5.2 — Optimize agent structure with Maestro (changes that cannot be achieved by setting parameters alone)
        # ============================================================================

        await maestro.optimize_structure(
            total_rollouts=10,  # Total number of rollouts to use for optimization.
            code_paths=["stock-assistant.py"],  # A list of paths corresponding to code implementations of the agent.
            verbose=True,  # If True, additional information will be printed during the optimization step.
        )


if __name__ == "__main__":
    asyncio.run(main())