"""Deposit-flight playbook, step 6: train the with-token model and its flag-off twin.

Both runs use the quickstart's CPU CI profile (forty epochs) and the
last-event readout; the with-token run also consumes the elapsed-silence
token. They are otherwise identical, so the ablation compares exactly one
thing.
"""

import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # playbooks/playbook_api.py
from playbook_api import Api, State

api, state = Api(), State("deposit-flight")

# The M16 switch this playbook proves: "on" (the default) builds, trains and
# serves with observation time; "off" runs the same chain without it, and the
# quiet-book step then fails, because nothing rises under silence.
FLAG_ON = os.environ.get("TESSERA_PLAYBOOK_OBSERVATION_TIME", "on") != "off"

profile_path = Path(__file__).resolve().parents[2] / "quickstart" / "ci-train-profile.json"
profile = json.loads(profile_path.read_text())
profile.pop("_comment", None)
models = {}
for side, dataset_id in sorted(state["datasets"].items()):
    body = json.loads(json.dumps(profile))
    body["dataset_id"] = dataset_id
    # Forty epochs instead of the quickstart's ten: on this small book the twins
    # need them to settle, so the ablation compares the token, not training noise.
    body["trainer"]["epochs"] = 40
    body["overrides"] = {**body["overrides"], "readout": "last_event"}
    if side == "on":
        body["overrides"]["observation_time"] = True
    run = api.post("/runs", body, key=f"deposit-flight-train-{side}")
    run = api.wait(
        f"/runs/{run['run_id']}",
        lambda item: item.get("status") in ("succeeded", "failed", "cancelled"),
        what=f"the {side} training run",
    )
    print(f"training {side}: {run['status']}")
    if run["status"] != "succeeded":
        sys.exit(json.dumps(run)[:3000])
    model = api.wait(
        f"/models/{run['model_id']}",
        lambda item: item.get("status") in ("evaluated", "promoted", "archived", "failed"),
        what=f"the {side} evaluation",
    )
    report = api.get(f"/models/{run['model_id']}/eval-report")
    failing = sorted(g for g, block in report["gates"].items() if block["status"] != "pass")
    print(f"model {side}: {model['status']}, failing gates: {failing or 'none'}")
    models[side] = run["model_id"]
print("trainings: succeeded")
state.update(models=models)