#!/usr/bin/env python3
"""Submit the fail-closed build, smoke, and resumable production chain."""
from __future__ import annotations

from datetime import datetime
import json
from pathlib import Path
import subprocess


HERE = Path(__file__).resolve().parent
STATUS = HERE / "caption_pipeline_jobs.json"


def submit(args, dependency=None):
    command = ["sbatch", "--parsable"]
    if dependency:
        command += [f"--dependency=afterok:{dependency}"]
    command += [str(x) for x in args]
    result = subprocess.run(command, cwd=HERE, check=True, text=True, capture_output=True)
    return result.stdout.strip().split(";")[0]


def main():
    if STATUS.exists():
        previous = json.loads(STATUS.read_text())
        previous_asset = previous.get("jobs", {}).get("assets", "unknown")
        archive = HERE / f"caption_pipeline_jobs.failed-{previous_asset}.json"
        assert not archive.exists(), f"Refusing to overwrite pipeline archive: {archive}"
        STATUS.replace(archive)
    jobs = {}
    jobs["assets"] = submit(["build_caption_assets.sbatch"])
    jobs["stages"] = submit(["build_caption_stages.sbatch"], jobs["assets"])
    jobs["freeze"] = submit(["freeze_caption_plans.sbatch"], jobs["stages"])
    jobs["smoke_4gpu"] = submit([
        "--nodes=1", "--time=01:00:00", "--job-name=m2-cap-smoke4",
        "train.sbatch", "plan_caption_smoke_4gpu.json"], jobs["freeze"])
    jobs["smoke_32nodes"] = submit([
        "--nodes=32", "--time=01:00:00", "--job-name=m2-cap-smoke32",
        "train.sbatch", "plan_caption_smoke_32nodes.json"], jobs["smoke_4gpu"])
    dependency = jobs["smoke_32nodes"]
    production = []
    for index in range(1, 6):
        job = submit(["--nodes=32", "--time=12:00:00", f"--job-name=m2-cap-prod{index}",
                      "train.sbatch", "plan_caption_production_32nodes.json"], dependency)
        production.append(job); dependency = job
    jobs["production_chain"] = production
    payload = {"status": "SUBMITTED", "created_at": datetime.now().astimezone().isoformat(),
               "dependency_policy": "afterok", "jobs": jobs,
               "production_authorized": True,
               "throughput_gate_seconds_per_update": 5.0}
    STATUS.write_text(json.dumps(payload, indent=2) + "\n")
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()