"""Actual upstream optimizer entry point; called only by launch.py --execute."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--plan", type=Path, required=True)
    args = parser.parse_args()
    plan = json.loads(args.plan.read_text())
    from examples.calc_x.train_calc_agent import build_config
    from agentlightning.verl.entrypoint import run_ppo
    from omegaconf import OmegaConf
    import ray

    config = build_config(model=plan["model"], agl_base_url=plan["gateway"],
                          agl_key=os.environ["HESSIAN_AGL_KEY"], run_name="hessian-calculator",
                          config_overrides=plan["overrides"] + [
                              "trainer.default_local_dir=" + json.dumps(str(Path(plan["output"]) / "checkpoints")),
                              "data.seed=" + str(plan["seed"]),
                          ])
    redacted = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
    redacted.agentlightning.agl_key = "<local-ephemeral-key>"
    OmegaConf.save(redacted, Path(plan["output"]) / "resolved-config.yaml")
    tasks = [{"question": f"What is {a} multiplied by {b}?", "result": str(a * b)}
             for a, b in [(6, 7), (8, 9), (4, 11), (5, 13)]]
    validation = [{"question": "What is 7 multiplied by 9?", "result": "63"}]
    (Path(plan["output"]) / "tasks.json").write_text(json.dumps({"train": tasks, "validation": validation}, indent=2))
    try:
        run_ppo(config, train_dataset=tasks, val_dataset=validation)
    finally:
        ray.shutdown()


if __name__ == "__main__":
    main()