"""This module is used to submit jobs to a queue on a cluster. It can submit a single job, or
if used in "rapid-fire" mode, can submit multiple jobs within a directory structure.
The details of job submission and queue communication are handled using Queueadapter,
which specifies a QueueAdapter as well as desired properties of the submit script.
"""

import glob
import os
import time
from datetime import datetime

from monty.os import cd

from fireworks.core.fworker import FWorker
from fireworks.fw_config import (
    ALWAYS_CREATE_NEW_BLOCK,
    QSTAT_FREQUENCY,
    QUEUE_JOBNAME_MAXLEN,
    QUEUE_RETRY_ATTEMPTS,
    QUEUE_UPDATE_INTERVAL,
    RAPIDFIRE_SLEEP_SECS,
    STREAM_LOGLEVEL,
    SUBMIT_SCRIPT_NAME,
)
from fireworks.utilities.fw_serializers import load_object
from fireworks.utilities.fw_utilities import create_datestamp_dir, get_fw_logger, get_slug, log_exception

__author__ = "Anubhav Jain, Michael Kocher"
__copyright__ = "Copyright 2012, The Materials Project"
__maintainer__ = "Anubhav Jain"
__email__ = "ajain@lbl.gov"
__date__ = "Dec 12, 2012"


def launch_rocket_to_queue(
    launchpad,
    fworker,
    qadapter,
    launcher_dir=".",
    reserve=False,
    strm_lvl=STREAM_LOGLEVEL,
    create_launcher_dir=False,
    fill_mode=False,
    fw_id=None,
):
    """Submit a single job to the queue.

    Args:
        launchpad (LaunchPad)
        fworker (FWorker)
        qadapter (QueueAdapterBase)
        launcher_dir (str): The directory where to submit the job
        reserve (bool): Whether to queue in reservation mode
        strm_lvl (str): level at which to stream log messages
        create_launcher_dir (bool): Whether to create a subfolder launcher+timestamp, if needed
        fill_mode (bool): whether to submit jobs even when there is nothing to run
            (only in non-reservation mode)
        fw_id (int): specific fw_id to reserve (reservation mode only)
    """
    fworker = fworker or FWorker()
    launcher_dir = os.path.abspath(launcher_dir)
    l_logger = get_fw_logger("queue.launcher", l_dir=launchpad.logdir, stream_level=strm_lvl)

    l_logger.debug("getting queue adapter")
    qadapter = load_object(qadapter.to_dict())  # make a defensive copy, mainly for reservation mode

    fw, launch_id = None, None  # only needed in reservation mode

    if not os.path.exists(launcher_dir):
        raise ValueError(f"Desired launch directory {launcher_dir} does not exist!")

    if "--offline" in qadapter["rocket_launch"] and not reserve:
        raise ValueError("Must use reservation mode (-r option) of qlaunch when using offline option of rlaunch!!")

    if reserve and "singleshot" not in qadapter.get("rocket_launch", ""):
        raise ValueError("Reservation mode of queue launcher only works for singleshot Rocket Launcher!")

    if fill_mode and reserve:
        raise ValueError("Fill_mode cannot be used in conjunction with reserve mode!")

    if fw_id and not reserve:
        raise ValueError("qlaunch for specific fireworks may only be used in reservation mode.")

    if fill_mode or launchpad.run_exists(fworker):
        launch_id = None
        try:
            if reserve:
                if fw_id:
                    l_logger.debug("finding a FW to reserve...")
                fw, launch_id = launchpad.reserve_fw(fworker, launcher_dir, fw_id=fw_id)
                if not fw:
                    l_logger.info("No jobs exist in the LaunchPad for submission to queue!")
                    return False
                l_logger.info(f"reserved FW with fw_id: {fw.fw_id}")

                # update qadapter job_name based on FW name
                job_name = get_slug(fw.name)[0:QUEUE_JOBNAME_MAXLEN]
                qadapter.update(job_name=job_name)

                if "_queueadapter" in fw.spec:
                    l_logger.debug("updating queue params using Firework spec..")
                    qadapter.update(fw.spec["_queueadapter"])

                # reservation mode includes --fw_id in rocket launch
                qadapter["rocket_launch"] += f" --fw_id {fw.fw_id}"

                # update launcher_dir if _launch_dir is selected in reserved fw
                if "_launch_dir" in fw.spec:
                    fw_launch_dir = os.path.expandvars(fw.spec["_launch_dir"])

                    if not os.path.isabs(fw_launch_dir):
                        fw_launch_dir = os.path.join(launcher_dir, fw_launch_dir)

                    launcher_dir = fw_launch_dir

                    os.makedirs(launcher_dir, exist_ok=True)

                    launchpad.change_launch_dir(launch_id, launcher_dir)
                elif create_launcher_dir:
                    # create launcher_dir
                    launcher_dir = create_datestamp_dir(launcher_dir, l_logger, prefix="launcher_")
                    launchpad.change_launch_dir(launch_id, launcher_dir)

            elif create_launcher_dir:
                # create launcher_dir
                launcher_dir = create_datestamp_dir(launcher_dir, l_logger, prefix="launcher_")

            # move to the launch directory
            l_logger.info(f"moving to launch_dir {launcher_dir}")

            with cd(launcher_dir):
                if "--offline" in qadapter["rocket_launch"]:
                    setup_offline_job(launchpad, fw, launch_id)

                l_logger.debug("writing queue script")
                with open(SUBMIT_SCRIPT_NAME, "w") as f:
                    queue_script = qadapter.get_script_str(launcher_dir)
                    f.write(queue_script)

                l_logger.info("submitting queue script")
                reservation_id = qadapter.submit_to_queue(SUBMIT_SCRIPT_NAME)
                if not reservation_id:
                    raise RuntimeError(
                        "queue script could not be submitted, check queue script/queue adapter/queue server status!"
                    )
                if reserve:
                    launchpad.set_reservation_id(launch_id, reservation_id)
            return reservation_id

        except Exception:
            log_exception(l_logger, "Error writing/submitting queue script!")
            if reserve and launch_id is not None:
                try:
                    l_logger.info(f"Un-reserving FW with fw_id, launch_id: {fw.fw_id}, {launch_id}")
                    launchpad.cancel_reservation(launch_id)
                    launchpad.forget_offline(launch_id)
                except Exception:
                    log_exception(l_logger, f"Error unreserving FW with fw_id {fw.fw_id}")

            return False

    else:
        l_logger.info("No jobs exist in the LaunchPad for submission to queue!")
        return None  # note: this is a hack (rather than False) to indicate a soft failure to rapidfire()


def rapidfire(
    launchpad,
    fworker,
    qadapter,
    launch_dir=".",
    block_dir=None,
    nlaunches=0,
    njobs_queue=0,
    njobs_block=500,
    sleep_time=None,
    reserve=False,
    strm_lvl=STREAM_LOGLEVEL,
    timeout=None,
    fill_mode=False,
) -> None:
    """Submit many jobs to the queue.

    Args:
        launchpad (LaunchPad)
        fworker (FWorker)
        qadapter (QueueAdapterBase)
        launch_dir (str): directory where we want to write the blocks
        block_dir (str): directory to use as block dir. Can be a new or existing block. Dirname must
            start with 'block_'.
        nlaunches (int): total number of launches desired; "infinite" for loop, 0 for one round
        njobs_queue (int): stops submitting jobs when njobs_queue jobs are in the queue, 0 for no limit.
            If 0 skips the check on the number of jobs in the queue.
        njobs_block (int): automatically write a new block when njobs_block jobs are in a single block
        sleep_time (int): secs to sleep between rapidfire loop iterations
        reserve (bool): Whether to queue in reservation mode
        strm_lvl (str): level at which to stream log messages
        timeout (int): # of seconds after which to stop the rapidfire process
        fill_mode (bool): whether to submit jobs even when there is nothing to run (only in
            non-reservation mode)
    """
    sleep_time = sleep_time or RAPIDFIRE_SLEEP_SECS
    launch_dir = os.path.abspath(launch_dir)
    nlaunches = -1 if nlaunches == "infinite" else int(nlaunches)
    l_logger = get_fw_logger("queue.launcher", l_dir=launchpad.logdir, stream_level=strm_lvl)

    # make sure launch_dir exists:
    if not os.path.exists(launch_dir):
        raise ValueError(f"Desired launch directory {launch_dir} does not exist!")

    num_launched = 0
    start_time = datetime.now()

    try:
        l_logger.info("getting queue adapter")

        prev_blocks = sorted(glob.glob(os.path.join(launch_dir, "block_*")), reverse=True)
        if block_dir is not None:
            if not block_dir.startswith("block_"):
                raise ValueError(f"Invalid name {block_dir}, block dirs must start with 'block_")
            block_dir = os.path.abspath(os.path.join(launch_dir, block_dir))
            os.mkdir(block_dir, exist_ok=True)
        elif prev_blocks and not ALWAYS_CREATE_NEW_BLOCK:
            block_dir = os.path.abspath(os.path.join(launch_dir, prev_blocks[0]))
            l_logger.info(f"Found previous block, using {block_dir}")
        else:
            block_dir = create_datestamp_dir(launch_dir, l_logger)

        while True:
            # get number of jobs in queue if a maximum has been set.
            jobs_in_queue = 0
            if njobs_queue:
                jobs_in_queue = _get_number_of_jobs_in_queue(qadapter, njobs_queue, l_logger)
            job_counter = 0  # this is for QSTAT_FREQUENCY option

            while launchpad.run_exists(fworker) or (fill_mode and not reserve):
                if timeout and (datetime.now() - start_time).total_seconds() >= timeout:
                    l_logger.info("Timeout reached.")
                    break

                if njobs_queue and jobs_in_queue >= njobs_queue:
                    l_logger.info(f"Jobs in queue ({jobs_in_queue}) meets/exceeds maximum allowed ({njobs_queue})")
                    break

                l_logger.info("Launching a rocket!")

                # switch to new block dir if it got too big
                if _njobs_in_dir(block_dir) >= njobs_block:
                    l_logger.info(f"Block got bigger than {njobs_block} jobs.")
                    block_dir = create_datestamp_dir(launch_dir, l_logger)

                # launch a single job
                return_code = launch_rocket_to_queue(
                    launchpad, fworker, qadapter, block_dir, reserve, strm_lvl, True, fill_mode
                )
                if return_code is None:
                    l_logger.info("No READY jobs detected...")
                    break
                if not return_code:
                    raise RuntimeError("Launch unsuccessful!")
                num_launched += 1
                if nlaunches > 0 and num_launched == nlaunches:
                    l_logger.info(f"Launched allowed number of jobs: {num_launched}")
                    break
                # wait for the queue system to update
                l_logger.info(f"Sleeping for {QUEUE_UPDATE_INTERVAL} seconds...zzz...")
                time.sleep(QUEUE_UPDATE_INTERVAL)
                jobs_in_queue += 1
                job_counter += 1
                if job_counter % QSTAT_FREQUENCY == 0 and njobs_queue:
                    job_counter = 0
                    jobs_in_queue = _get_number_of_jobs_in_queue(qadapter, njobs_queue, l_logger)

            if (
                (nlaunches > 0 and num_launched == nlaunches)
                or (timeout and (datetime.now() - start_time).total_seconds() >= timeout)
                or (nlaunches == 0 and not launchpad.future_run_exists(fworker))
            ):
                break

            l_logger.info(f"Finished a round of launches, sleeping for {sleep_time} secs")
            time.sleep(sleep_time)
            l_logger.info("Checking for Rockets to run...")

    except Exception:
        log_exception(l_logger, "Error with queue launcher rapid fire!")


def _njobs_in_dir(block_dir):
    """Internal method to count the number of jobs inside a block.

    Args:
        block_dir: (str) the block directory we want to count the jobs in

    Return:
        (int): job count
    """
    return len(glob.glob(f"{os.path.abspath(block_dir)}/launcher_*"))


def _get_number_of_jobs_in_queue(qadapter, njobs_queue, l_logger):
    """Internal method to get the number of jobs in the queue using the given job params.
    In case of failure, automatically retries at certain intervals...

    Args:
        qadapter (QueueAdapter)
        njobs_queue (int): The desired maximum number of jobs in the queue
        l_logger (logger): A logger to put errors/info/warnings/etc.

    Return:
        (int)
    """
    RETRY_INTERVAL = 30  # initial retry in 30 sec upon failure

    for _i in range(QUEUE_RETRY_ATTEMPTS):
        try:
            jobs_in_queue = qadapter.get_njobs_in_queue()
            if jobs_in_queue is not None:
                l_logger.info(f"{jobs_in_queue} jobs in queue. Maximum allowed by user: {njobs_queue}")
                return jobs_in_queue
        except Exception:
            log_exception(l_logger, f"Could not get number of jobs in queue! Sleeping {RETRY_INTERVAL} secs...zzz...")
        time.sleep(RETRY_INTERVAL)
        RETRY_INTERVAL *= 2

    raise RuntimeError("Unable to determine number of jobs in queue, check queue adapter and queue server status!")


def setup_offline_job(launchpad, fw, launch_id) -> None:
    # separate this function out for reuse in unit testing
    fw.to_file("FW.json")
    with open("FW_offline.json", "w") as f:
        f.write(f'{{"launch_id":{launch_id}}}')
    launchpad.add_offline_run(launch_id, fw.fw_id, fw.name)
