#!/usr/bin/env python3
"""Set a job's status, e.g. from the Argo exit handler when a workflow fails.

A succeeded job is never overwritten: the loader's success is the source of truth.
"""
import argparse
import os
import sys

from load_db import engine_for
from sqlalchemy import Engine, text


def set_status(engine: Engine, job_id: str, status: str, log: str | None) -> None:
    with engine.begin() as conn:
        conn.execute(
            text("""
                UPDATE jobs
                SET status = CAST(:status AS jobstatus), log = :log,
                    finished_at = CASE WHEN :status IN ('succeeded', 'failed') THEN now() END
                WHERE id = :id AND status <> 'succeeded'
            """),
            {"id": job_id, "status": status, "log": log},
        )


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--job-id", required=True)
    p.add_argument("--status", required=True, choices=["running", "failed"])
    p.add_argument("--log")
    a = p.parse_args()
    url = os.environ.get("DATABASE_URL")
    if not url:
        sys.exit("DATABASE_URL is not set")
    set_status(engine_for(url), a.job_id, a.status, a.log)


if __name__ == "__main__":
    main()