import os
import numpy as np
from datetime import datetime, timedelta
from glob import glob
from re import sub
from traceback import extract_tb
from step0 import prepare_file
from step1 import gather_data
# from step2 import cal_zenith
from step3 import pair_ro_ref
from step4 import cal_exph
from sys import argv

if __name__ == '__main__':
    if len(argv) == 3:
        working_leo = argv[2]
        nProcess = int(argv[1])
    elif len(argv) == 2:
        working_leo = ''
        nProcess = int(argv[1])
    else:
        working_leo = ''
        nProcess = 0
    working_dir = '/data3/tcliu/gps_pod/python/planetiq/'
    data_dir = '/data3/xinjiaz/data_GPSRO/planetiq_scdr/GNOMES_PYXIS_L1A/'
    os.chdir(working_dir)
    root_out = working_dir + f'L1_extract_{working_leo}'

    logfile = f'progress_rt_{working_leo}.txt'
    def log_message(msg):
        with open(logfile, 'a') as fid:
            fid.write(f"{datetime.now().strftime('%H:%M:%S')} {msg}\n")

    log_message('- checking progress...')
    progfile = f'progress_rt_{working_leo}.npz'
    if os.path.exists(progfile):
        with np.load(progfile, allow_pickle=True) as fid:
            tarfiles = fid['tarfiles']
            progress = fid['progress']
            timetag = fid['timetag']
    else:
        tarfiles = np.array([],dtype='<U119')
        progress = np.array([],dtype='int8')
        timetag = np.array([],dtype=datetime)

    def update_progress(pnum,state):
        progress[pnum] = state
        timetag[pnum] = datetime.now()

    idx = np.where(progress%2==0)[0]
    for ii in idx:
        if progress[ii] == -2 or datetime.now() - timetag[ii] > timedelta(minutes=30):
            update_progress(ii,-1)
            # os.system('rm -r ' + root_out + '/*')
            # os.system('rm -r t01_RO/*.npz')
            # os.system('rm -r t01_REF/*.npz')
            # os.system('rm -r t03_pair/*.npz')
            log_message("-- Old process timeout and was terminated.")

    def save_progress():
        np.savez(f'progress_rt_{working_leo}.npz',
            tarfiles = tarfiles,
            progress = progress,
            timetag = timetag,
        )

    if np.any(progress==0):
        log_message("-- Still progressing.")
    else:
        manual_date = False # set True to manually set dates to process
        nloop = 0
        while not manual_date:
            new_proc = False
            nloop += 1
            if manual_date: # manually assign dates to process
                # process_dates = [datetime(2025, 9, 30) + timedelta(days=x) for x in range(33)]
                process_dates = [datetime(2025, 10, 31) + timedelta(days=x) for x in range(1)]
            else:
                dt_today = datetime.today()
                process_dates = [dt_today + timedelta(days=x) for x in range(-2,2)] # 3 recent days
            root_files = np.hstack([glob(data_dir + f'{x:%Y-%m-%d}/*{working_leo[:4]}.tar') for x in process_dates])

            idx = np.array([tarfile in root_files for tarfile in tarfiles], dtype=bool)
            tarfiles = tarfiles[idx]
            progress = progress[idx]
            timetag = timetag[idx]
            if np.any(idx==False):
                log_message(f'- {np.sum(idx==False):d} files are removed from the queue.')

            idx = np.array([root_file not in tarfiles for root_file in root_files], dtype=bool)
            tarfiles = np.concatenate((tarfiles, root_files[idx]))
            progress = np.concatenate((progress, np.full((np.sum(idx),),-1,dtype='int8')))
            timetag = np.concatenate((timetag, np.full((np.sum(idx),),datetime.now(),dtype=datetime)))
            if np.any(idx):
                log_message(f'- {np.sum(idx == True):d} files are added to the queue.')

            idx = np.argsort(tarfiles)
            tarfiles = tarfiles[idx]
            progress = progress[idx]
            timetag = timetag[idx]
            save_progress()

            try:
                idx = np.where(progress==-1)[0]
                for ii in range(len(idx)):
                    update_progress(idx[ii],0)
                    save_progress()
                    log_message(f"-- Processing {sub(r'.*/',r'',tarfiles[idx[ii]])} ({ii+1:d}/{len(idx):d})")
                    dt = datetime.strptime(sub(r'.*/(\d{4}-\d\d-\d\d)/.*',r'\1',tarfiles[idx[ii]]),'%Y-%m-%d')
                    os.makedirs(root_out, exist_ok=True)
                    os.makedirs(f't01_RO_{working_leo}', exist_ok=True)
                    os.makedirs(f't01_REF_{working_leo}', exist_ok=True)
                    os.makedirs(f't03_pair_{working_leo}', exist_ok=True)
                    os.system('rm -r ' + root_out + '/*')
                    os.system(f'rm t01_RO_{working_leo}/*.npz')
                    os.system(f'rm t01_REF_{working_leo}/*.npz')
                    os.system(f'rm t03_pair_{working_leo}/*.npz')
                    prepare_file(tarfiles[idx[ii]], working_leo)
                    gather_data(working_leo)
                    # cal_zenith(dt)
                    pair_ro_ref(dt, working_leo)
                    cal_exph(working_leo, nProcess, sub(r'.*/', r'', tarfiles[idx[ii]]))
                    update_progress(idx[ii], 1)
                if len(idx)>0:
                    save_progress()
                    new_proc = True
                    log_message(f'- Progress completed.')
            except Exception as e:
                log_message('- Unexpected error occurs.')
                log_message(f"-- {type(e).__name__} >> {str(e)}")
                for tb in extract_tb(e.__traceback__):
                    log_message(f"--  line {tb.lineno} of {tb.filename}")
                update_progress(progress == 0, -2)
                save_progress()
                new_proc = True
            if not new_proc:
                if nloop > 1:
                    log_message(f'- No new item, exit program.')
                break