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

working_leo = ''
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=60):
        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 = True # set True to manually set dates to process
    if manual_date: # manually assign dates to process
        process_dates = [datetime(2025, 11, 19) + 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,1)] # 3 recent days
    root_files = np.hstack([glob(data_dir + f'{x:%Y-%m-%d}/*{working_leo}.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)
            update_progress(idx[ii], 1)
        if len(idx)>0:
            save_progress()
            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()