"""Security-lake playbook, step 7: build the dataset, then train.

The build comes after the protocol (step 6), so the report can judge it. It
keeps each user's last 64 events: a derived label is supervised only at
anchors inside the kept window, so the window must reach back to the takeovers.
The run uses the quickstart's documented CPU CI profile; production runs drop its
`overrides` and `trainer`.
"""

import json
import sys
from datetime import datetime, timedelta
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("security-lake")

label_id = state["label_ids"]["repeated_logon_failure_1d"]
start = datetime.fromisoformat(state["window"]["start"])
end = start + timedelta(days=state["window"]["days"] + 1)
start_text, end_text = (moment.strftime("%Y-%m-%dT%H:%M:%SZ") for moment in (start, end))
build = api.post(
    "/datasets",
    {
        "entity_type": "user",
        "labels": [label_id],
        "window": {
            "max_events_per_entity": 64,
            "time_range": {"start": start_text, "end": end_text},
        },
        "split": {"method": "entity_hash", "train": 0.7, "val": 0.1, "test": 0.2, "seed": 7},
    },
    key="security-lake-dataset",
)
dataset = api.wait(
    f"/datasets/{build['dataset_id']}",
    lambda ds: ds.get("status") in ("succeeded", "failed", "aborted_leakage"),
    what="the dataset build",
)
print(f"dataset: {dataset['status']}")
if dataset["status"] != "succeeded":
    sys.exit(json.dumps(dataset)[:3000])

profile_path = Path(__file__).resolve().parents[2] / "quickstart" / "ci-train-profile.json"
profile = json.loads(profile_path.read_text())
profile.pop("_comment", None)
profile["dataset_id"] = dataset["dataset_id"]
run = api.post("/runs", profile, key="security-lake-train")
run = api.wait(
    f"/runs/{run['run_id']}",
    lambda item: item.get("status") in ("succeeded", "failed", "cancelled"),
    what="the training run",
)
print(f"training: {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="the model's evaluation",
)
print(f"model: {model['status']}")
state.update(dataset_id=dataset["dataset_id"], model_id=run["model_id"], cutoff=end_text)