import os
import sys
import argparse

parser = argparse.ArgumentParser(description='wrapper')
parser.add_argument('--job_data_root', default="job_data_root", type=str)
args = parser.parse_args()

def run():
    print('execute predicting', flush=True)
    syscommand = fr"export JOB_DATA_ROOT={args.job_data_root} &&" + \
                fr"export TEST=1 &&" + \
                "bash ./predict_ddp_k8spod_webui.sh"

    print(fr'wrapper command {syscommand}', flush=True)
    retcode = os.WEXITSTATUS(os.system(syscommand))
    return retcode

if __name__ == "__main__":
    job_data_root = args.job_data_root
    tar_preprocess_dir = os.path.join(job_data_root, 'output', 'preprocess')
    if not os.path.exists(tar_preprocess_dir) or len(os.listdir(tar_preprocess_dir)) == 0:
        print('re-preprocessing...', flush=True)
        from run_preprocess import run as run_preprocess_func
        # run_preprocess
        ret_code = run_preprocess_func()
        if ret_code != 0:
            sys.exit(ret_code)
        print('done...', flush=True)
        
    print('running predict...', flush=True)
    # run infer
    ret_code = run()
    if ret_code != 0:
        sys.exit(ret_code)
    print('all done', flush=True)