#!/usr/bin/env python3
import argparse
import os

# Create the parser
parser = argparse.ArgumentParser(description='User parameter YAML file.')

# Add the arguments
parser.add_argument('user_parameter_file',
                    metavar='file_path',
                    type=str,
                    help='the path to the user parameter file')

parser.add_argument('--submit',
                    '-s',
                    metavar='job_submission_system',
                    type=str,
                    choices=['condor', 'slurm'],
                    help='the submit option, the available options are condor and slurm')

# working dir
parser.add_argument('--work-dir',
                    '-d',
                    metavar='work_dir',
                    type=str,
                    default='.',
                    help='the working directory')

# force overwrite
parser.add_argument('--overwrite',
                    '-o',
                    action='store_true',
                    default=False,
                    help='overwrite the existing results')

# conda env, default is current env
parser.add_argument('--conda-env',
                    '-e',
                    metavar='conda_env',
                    type=str,
                    default=os.environ.get('CONDA_DEFAULT_ENV'),
                    help='the conda environment')

# accounting_group
parser.add_argument('--accounting-group',
                    '-g',
                    metavar='accounting_group',
                    type=str,
                    default='ligo.sim.o4.burst.allsky.cwboffline',
                    help='the condor accounting group')

# threads
parser.add_argument('--threads',
                    '-n',
                    metavar='threads',
                    type=int,
                    default=0,
                    help='the number of threads, if it set to 0, it will use the value from the user parameter file')

# Parse the arguments
args = parser.parse_args()

from pycwb.search import search
from pycwb.modules.condor.condor import generate_job_script, generate_condor_sub, submit

if args.submit:
    print(f"Submitting the search on {args.submit}.")
    # create a dictionary with the submit option
    if args.submit == 'condor':
        generate_job_script(args.user_parameter_file, args.conda_env, args.work_dir, args.threads)
        generate_condor_sub(args.work_dir, args.accounting_group, args.threads)
        submit(args.work_dir)
    elif args.submit == 'slurm':
        print("Not implemented yet.")

else:
    print("Running the search locally.")
    # Run the search function with the specified user parameter file
    search(args.user_parameter_file, working_dir=args.work_dir, overwrite=args.overwrite, nproc=args.threads)
