#! /usr/bin/env python3

# ICON
#
# ---------------------------------------------------------------
# Copyright (C) 2004-2026, DWD, MPI-M, DKRZ, KIT, ETH, MeteoSwiss
# Contact information: icon-model.org
#
# See AUTHORS.TXT for a list of authors
# See LICENSES/ for license information
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------

import json
import os
import signal
import sys
import time
from pathlib import Path

# ===============================================================================
# this should prevent a blas-related runtime error when using numpy, which is
# loaded as dependency of the buildbot scripting (pandas)
# error: https://buildbot.dkrz.de/#/builders/73/builds/4294/steps/4/logs/stdio
os.environ["OPENBLAS_NUM_THREADS"] = "1"
# ===============================================================================
# import local modules
import icon_env
from buildbot_config import BuildbotConfig
from cmd_line_job import CmdLineJob
from icon_paths import base_path, buildbot_list_path, run_path
from pbs_job import PBSJob
from slurm_job import SerialSlurmJob, SlurmJob

# global constants
EX_FAILED = 1
STATUS_FILE = Path(base_path) / "LOOP_STATUS_EXP_FILE"


# method to write a buildbot status overview file. This needs to be called if
# - all submitted tests are successfull or failed
# - submitted tests are canceled
# - the script itself is canceled, for example with ctrl-c
def writeStatusFile(filename, experiments):
    experimentsFailed = []
    with open(str(filename), "w") as file:
        for eobj in experiments:
            exp_jobid = eobj.batch_job.jobid
            exp_file = eobj.get_run_name(relative=False)
            if eobj.batch_job.returncode is not None:
                if eobj.failed():
                    experimentsFailed.append(True)
                    exit_msg = f"FAILED (jobID:{exp_jobid})"
                else:
                    experimentsFailed.append(False)
                    exit_msg = f"OK (jobID:{exp_jobid})"
            else:
                experimentsFailed.append(True)
                print(
                    f"internal error, did not get returncode for {eobj.name} (jobID:{exp_jobid})"
                )
                exit_msg = "FAILED"

            file.write("{0:50}: {1}\n".format(exp_file.name, exit_msg))

    return experimentsFailed


# submit all tests, wait for return values/status or wait for cancelling
def runexp():
    # Get build settings
    setup_dict = icon_env.info()

    builder = "{}_{}".format(
        setup_dict["use_target"].upper(), setup_dict["use_compiler"]
    )

    with open(str(run_path / "buildbot_config.json"), "r") as file:
        config = json.load(file)

    builder = config["builder"]
    list_name = config["list_name"]

    print(f"initializing experiments for builder {builder} in list {list_name}")

    full_list_name = buildbot_list_path / list_name

    if Path(full_list_name).exists():
        thisList = BuildbotConfig.from_pickle(full_list_name)
    else:
        print(f"did not find experiment list {full_list_name}")
        sys.exit(EX_FAILED)

    bobj = thisList.builder_meta[builder]
    bobj.submit = setup_dict["use_submit"]

    # Set up mkexp run-time environment (currently breaks on DWD and Balfrin)
    if bobj.machine not in ["dwd_nec", "balfrin"]:
        icon_env.load()

    serial_sbatch = setup_dict.setdefault("use_serial_submit", "no") == "yes"

    if bobj.submit.startswith("sbatch"):
        if serial_sbatch:
            Job = SerialSlurmJob
        else:
            Job = SlurmJob
    elif bobj.submit.startswith("qsub"):
        Job = PBSJob
    elif bobj.submit == "":
        Job = CmdLineJob
    else:
        print(
            f"no batch job helper class implemented for machine {bobj.machine}"
        )
        sys.exit(EX_FAILED)

    # set up the experiments for job submission
    # builder info and batch system information need to be added
    eobjs = thisList.get_experiments_by_builder(builder).flatten()

    for eobj in eobjs:
        # BatchJob keeps track of dependencies and is responsible for submitting to the queue
        eobj.batch_job = Job(bobj.submit, run_path)

    for eobj in eobjs:
        # finally submit the jobs
        eobj.submit(bobj)
        exp_jobid = eobj.batch_job.jobid
        if eobj.batch_job.jobid is not None:
            print(
                f"Submitted {eobj.get_run_name(relative=True)} (jobID:{eobj.batch_job.jobid})"
            )

    # Signal handling for stopping this script with ctrl-c on the command line:
    # stop all submitted jobs
    def cleanupOnInterupt(signum, frame):
        print("runexp: caught signal", signum, frame)
        print("runexp: Try to cancel all jobs .........................")
        for eobj in eobjs:
            eobj.batch_job.cancel()
        experimentsFailed = writeStatusFile(STATUS_FILE, eobjs)
        sys.exit(EX_FAILED)

    signal.signal(signal.SIGINT, cleanupOnInterupt)
    signal.signal(signal.SIGTERM, cleanupOnInterupt)

    # wait for completion of all gathered processes and do some output while waiting
    #
    # special handling for job cancelling:
    #   currently used slurm versions at dkrz, kit and cscs do NOT return
    #   non-zero exit codes when pending jobs are cancelled with scancel
    #
    #   that's why the jobs status has to be checked with sacct in addition to
    #   the return code for slurm
    #
    #   PBS works reasonable - job status can be ignored, since the exit code
    #   does what it is supposed to
    incomplete_jobs = set(eobjs)
    while len(incomplete_jobs) > 0:
        completed_jobs = set()
        for eobj in incomplete_jobs:
            if eobj.batch_job.poll(timeout=60):
                if eobj.failed():
                    _status = "FAILED"
                else:
                    _status = "OK"

                print(
                    f"job finished: {eobj.get_run_name(relative=True)} (jobID:{eobj.batch_job.jobid}): {_status}"
                )
                completed_jobs.add(eobj)
            else:
                N = len(incomplete_jobs) - len(completed_jobs)
                if N > 1:
                    print(f"Waiting for {N} jobs:")
                    for e in set(incomplete_jobs) - set(completed_jobs):
                        print(
                            f"\t{e.get_run_name(relative=True)} ({str(e.batch_job.jobid)})"
                        )
                elif N == 1:
                    e = [*incomplete_jobs][0]
                    print(
                        f"Waiting for the last job: {e.get_run_name(relative=True)} ({e.batch_job.jobid})"
                    )
        incomplete_jobs.difference_update(completed_jobs)

    # wait a bit longer, because queing or file system might create logfiles
    # with a certain delay
    time.sleep(20)

    # 1) parse returncodes of all experiments and fail if at least one of them did fail
    # 2) write STATUS file which is supposed to be shown by buildbot for better overview
    experimentsFailed = writeStatusFile(STATUS_FILE, eobjs)

    exit_code = EX_FAILED if True in experimentsFailed else os.EX_OK
    sys.exit(exit_code)


if __name__ == "__main__":
    runexp()
