#!python
# coding: utf-8

import argparse

import numpy as np
import os
import sys
import time
from astropy.table import Table
import glob

## Import some helper functions, you can see their definitions by uncomenting the bash shell command
from desiutil.log import get_logger
from desispec.workflow.tableio import load_table, write_table
from desispec.workflow.proctable import get_processing_table_pathname
from desispec.workflow.processing import update_and_recursively_submit
from desispec.workflow.queue import get_resubmission_states

def parse_args():  # options=None):
    """
    Creates an arguments parser for the desi run production
    """
    parser = argparse.ArgumentParser(description="Submit a one past night of data for processing with the DESI data pipeline.")

    parser.add_argument("-n","--night", type=str, default=None,
                        required=False, help="The night you want processed.")
    parser.add_argument("-e","--expids", type=str, default=None,
                        required=False, help="The exposure ids to resubmit (along "
                                             + "with the jobs they depend on)."
                                             + " Should be a comma separated list.")
    parser.add_argument("-t","--tileids", type=str, default=None,
                        required=False, help="The tile ids to resubmit (along "
                                             + "with the jobs they depend on)."
                                             + " Should be a comma separated list.")
    parser.add_argument("--proc-table-pathname", type=str, required=False, default=None,
                        help="Directory name where the output processing table should be saved.")
    parser.add_argument("--tab-filetype", type=str, required=False, default='csv',
                        help="File format and extension for the exp and proc tables.")
    parser.add_argument("-r", "--reservation", type=str, required=False, default=None,
                        help="The reservation to submit jobs to. If None, it is not submitted to a reservation.")
    parser.add_argument("--dry-run-level", type=int, default=0,
                        help="What level of dry_run to perform, if any. Default is 0. "
                              + "0 which runs the code normally. "
                              + "1 writes all files but doesn't submit any jobs to Slurm. "
                              + "2 writes tables but doesn't write scripts or submit anything. "
                              + "3 Doesn't write or submit anything but queries Slurm normally for job status. "
                              + "4 Doesn't write, submit jobs, or query Slurm."
                              + "5 Doesn't write, submit jobs, or query Slurm; instead it makes up the status of the jobs.")
    parser.add_argument("--resub-states", type=str, default=None, required=False,
                        help="The SLURM queue states that should be resubmitted. " +
                             "E.g. UNSUBMITTED, BOOT_FAIL, DEADLINE, NODE_FAIL, " +
                             "OUT_OF_MEMORY, PREEMPTED, TIMEOUT, CANCELLED, FAILED.")
    parser.add_argument("--no-resub-failed", action="store_true", required=False,
                        help="Give this flag if you do NOT want to resubmit " +
                             "jobs with Slurm status 'FAILED' by default. " +
                             "This should not be used if defining " +
                             "--resub-states explicitly.")

    args = parser.parse_args()

    if args.expids is not None:
        args.expids = np.array([eid.strip() for eid in args.expids.split(',')]).astype(int)
    if args.tileids is not None:
        args.tileids = np.array([tid.strip() for tid in args.tileids.split(',')]).astype(int)

    if args.resub_states is not None:
        ## User should never provide custom list of states and request to remove FAILED
        if args.no_resub_failed:
            log = get_logger()
            msg = f"Provided user-defined resubmision states {args.resub_states} but " \
                  + f"also requested --no-resub-failed. Please choose one or the other."
            log.critical(msg)
            raise ValueError(msg)
        ## Clean up the input string into a list of strings
        args.resub_states = [state.strip().upper() for state in args.resub_states.split(',')]

    return args


if __name__ == '__main__':
    args = parse_args()
    log = get_logger()
    ptable_pathname = args.proc_table_pathname
    if ptable_pathname is None:
        if args.night is None:
            ValueError("Either night or proc-table-path must be specified")
        ## Determine where the processing table will be written
        ptable_pathname = get_processing_table_pathname(prodmod=args.night,
                                             extension=args.tab_filetype)

    if not os.path.exists(ptable_pathname):
        ValueError(f"Processing table: {ptable_pathname} doesn't exist.")

    if args.dry_run_level > 0:
        log.info(f"{args.dry_run_level=}, so will be simulating some features."
                 + f" See parser for what each level limits.")

    ## Load in the files defined above
    ptable = load_table(tablename=ptable_pathname, tabletype='proctable')
    log.info(f"Identified ptable with {len(ptable)} entries.")
    ptable, nsubmits = update_and_recursively_submit(ptable, submits=0,
                                                     resubmission_states=args.resub_states,
                                                     no_resub_failed=args.no_resub_failed,
                                                     ptab_name=ptable_pathname,
                                                     dry_run_level=args.dry_run_level,
                                                     reservation=args.reservation,
                                                     expids=args.expids, tileids=args.tileids)

    if args.dry_run_level < 3:
        write_table(ptable, tablename=ptable_pathname)

    log.info("Completed all necessary queue resubmissions from processing "
             + f"table: {ptable_pathname}")
