#!/bin/env python 
import argparse
import os,sys
import subprocess
from aiub_gnss_data import AIUBData, wget_aiub_gnss
from nasa_gnss_data import NASAData, wget_nasa_gnss
from ucar_gnssro_data import UCARData, wget_ucar_gnssro, local_ucar_filelist
from ucar2umd_att_rt import ucar2umd_att_rt
from ucar2umd_rnx_rt import ucar2umd_rnx_rt

from s05_input2bernese_rt import bernese_datapool_subdir
from s05_input2bernese_rt import bernese_gnss_inputs
from s05_input2bernese_rt import bernese_leoatt_inputs, bernese_leornx_inputs
from satellite_info import sat_ant_info
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import List, Dict, Set
import gpsdate
import yaml
from rt_file_monitor import update_rt_satant_dict,update_rt_filelist
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import Pool


__authors__ = ["Ming Chen"]
__license__ = "Public Domain"
__maintainer__ = "Ming Chen"
__email__ = "ming.chen@noaa.gov"
__date__  = "May 03, 2026"

""" 
#####################################################################
# 
# rt_bern_proc.py
#
# Purpose:
# This Python script acts as a driver for the BERNESE workflow:
# 1) preparing required input files,
# 2) running the core BERNESE PCF, and
# 3) handling post-processing of the generated outputs.
#
# History:  
#       Created by: Ming Chen 2026/05/03
#                   ming.chen@noaa.gov
# 
# Usage: python  rt_bern_proc.py  --config-file rt_cosmic2.yaml
#
#
#####################################################################

"""

# ----------------------------------------------------------------------
# YAML loading
# ----------------------------------------------------------------------
def load_config(config_file: Path) -> Dict:
    if not config_file.exists():
        raise FileNotFoundError(f"Config file not found: {config_file}")
    with open(config_file, "r") as f:
        config = yaml.safe_load(f)
    return config

def rt_leo_preproc(config, satants=None):
    #config = load_config(Path(args.config_file))
    leo = config['leo']
    mission = leo['mission']
    window_minutes = config['bern_session_minutes']
    total_minutes = window_minutes 
    umd_att_repo = leo['att_repo_root']
    umd_rnx_repo = leo['rnx_repo_root']
    att_processed = []
    if config['real_time'].lower() != "on":  satants =  config['satant']
    print(satants)
    for k in satants:

        if isinstance(satants, dict):
           if satants[k]['status'] == 0: continue
           s, a =  k.split("_")
           satidx =int(s) ; antidx = int(a)
           tend = satants[k]["end_time"]
           ddt = datetime.strptime(tend, "%Y%m%dT%H:%M:%S")
        else:
           sa, tend =  k.split("=")
           s, a =  sa.split("_")
           satidx =int(s) ; antidx = int(a)
          
           ddt = datetime.strptime(tend, "%Y%m%dT%H:%M:%S")

        if total_minutes < 0:
           t_start = ddt + timedelta(minutes=total_minutes)
           t_end   = ddt
        else:
           t_start = ddt ; t_end   = ddt + timedelta(minutes=total_minutes)

        tstart = t_start.strftime("%Y%m%dT%H%M%S")
        tend = t_end.strftime("%Y%m%dT%H%M%S")

        mydate = t_end
        datestr = mydate.strftime('%Y%m%d') 
        hh = mydate.strftime('%H') 

        #satidx = 6;  antidx = 1

        print(f'ffff{mission} {satidx}, {t_end} {umd_att_repo}')
        ucar2umd_att_rt(mission, t_start,t_end, satidx=satidx, umd_repo =umd_att_repo)
        #if s not in att_processed:
        #   ucar2umd_att_rt(mission, t_start,t_end, satidx=satidx, umd_repo =umd_att_repo)
        #   att_processed.append(s)


        ucar2umd_rnx_rt(mission, t_start,t_end, satidx=satidx, antidx=antidx, umd_repo =umd_rnx_repo)
        print( t_start,t_end, satidx,antidx)


def rt_run_bern(config, satants=None, dry_run = True):
    gnss = config['gnss']
    leo = config['leo']
    bern = config['bern']
    gnss_prod = gnss['product'] #'COD0OPSULT'
    gnss_source = gnss['source'] #'nasa aiub'
    mission = leo['mission']
    window_minutes = config['bern_session_minutes']
    total_minutes = window_minutes 

    valid_runs = [] ; att_status = {}
    start_end = []

    if config['real_time'].lower() != "on": satants =  config['satant']
    for k in satants:
        
        if isinstance(satants, dict):
           if satants[k]['status'] == 0: continue
           s, a =  k.split("_")
           satidx =int(s) ; antidx = int(a)
           tend = satants[k]["end_time"]
           ddt = datetime.strptime(tend, "%Y%m%dT%H:%M:%S")
        else:
           sa, tend =  k.split("=")
           s, a =  sa.split("_")
           satidx =int(s) ; antidx = int(a)

           ddt = datetime.strptime(tend, "%Y%m%dT%H:%M:%S")
        if total_minutes < 0:
           t_start = ddt + timedelta(minutes=total_minutes)
           t_end   = ddt
        else:
           t_start = ddt ; t_end   = ddt + timedelta(minutes=total_minutes)

        tstart = t_start.strftime("%Y%m%dT%H%M%S")
        tend = t_end.strftime("%Y%m%dT%H%M%S")

        mydate = t_end
        myleo = {
           'mission': mission,
           'satidx':  satidx, 'antidx': antidx,
           'session': '0',
           'tstart': t_start,
           'tend': t_end,
        }


        datestr = mydate.strftime('%Y%m%d') 
        timestr = mydate.strftime('%Y%m%dT%H%M') 
        gnss_latency = gnss['latency'] 
        datapool_subdir =   bernese_datapool_subdir(myleo['mission'], myleo['satidx'], codcom='RTM')
        gnss_status = bernese_gnss_inputs(gnss_prod, t_end, data_source=gnss_source, 
                gnss_latency = gnss_latency, datapool_subdir = datapool_subdir )
        #print(gnss_status)
        if not gnss_status:
           print(f"GNSS input rror: {gnss_prod} {timestr} {gnss_latency} {gnss_source}")
           return valid_runs

        #if s not in att_status:
        #   leoatt_status = bernese_leoatt_inputs(myleo,datapool_subdir = datapool_subdir)
        #   att_status[s] = leoatt_status
        leoatt_status = bernese_leoatt_inputs(myleo,datapool_subdir = datapool_subdir)
        att_status[s] = leoatt_status
            
        leornx_status = bernese_leornx_inputs(myleo, datapool_subdir = datapool_subdir)
        leo_status = att_status[s] and leornx_status
        if leo_status: 
           valid_runs.append([satidx, antidx])
           start_end.append([tstart, tend, timestr])
        else: 
           print(f"LEO input rror: {myleo['mission']}  {myleo['satidx']} {myleo['antidx']}")
    
    if len(valid_runs) < 1: valid_runs
    commands = [] ; processes = []
    podgen_home  = os.environ['podgen_home'] 
    for k, t in zip(valid_runs,  start_end):
       tstart, tend, timestr = t[0],t[1],t[2]
       print(f"tstart-tend {tstart}, {tend}  {timestr}")
       cmd = f'{podgen_home}/rt_cosmic2_bern.bash'
       satidx, antidx = k[0], k[1]
       cmd = f'{cmd} {satidx} {antidx} "{timestr}" "{podgen_home}"'
       print(cmd)
       commands.append(cmd)
       
    # Explicitly define necessary environment variables
    my_env = os.environ.copy()
    my_env["PATH"] = "/usr/local/bin:/usr/bin:/bin:" + my_env.get("PATH", "")


    # Start all processes
    if dry_run: return

    for cmd in commands:
        p = subprocess.Popen(cmd.split(), env=my_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        processes.append(p)
    # Wait for all to complete and collect output
    for p in processes:
        output, error = p.communicate()
        print(output.decode())
 


# ----------------------------------------------------------------------
def parse_args():

    # Set up main parser
    parser = argparse.ArgumentParser(description="batch driver of the leopod generation")
    parser.add_argument("-v","--verbose", action='store_true',  help="print info for debug")
    parser.add_argument("--config-file", type=str, default='rt_cosmic2_reproc.yaml', help="config yaml file")
    parser.add_argument("--dry-run", action='store_true', help="dry run for testing")

    #subparsers = parser.add_subparsers(dest="command", required=True)
    #leo_preproc = subparsers.add_parser("leo_preproc", help="preprocessing leo mission data")
    #leo_preproc.set_defaults(func=rt_leo_preproc)
    #run_bern = subparsers.add_parser("bern_run", help="preprocessing leo mission data")
    #run_bern.set_defaults(func=rt_run_bern)
 
 
    return parser.parse_args()
 
def main():
    args = parse_args()
   
    mission ='cosmic2'
    config = load_config(Path(args.config_file))
    satant_file = config['rt_satant_file']
    satants = None
    
    umd_repo = config['leo']['rnx_repo_root']
    if config['real_time']:
       #ctime = datetime.strptime("2026.121.05.45", "%Y.%j.%H.%M")
       now_utc = datetime.now(timezone.utc)
       ctime =  now_utc - timedelta(days=2) 
       #satants = update_rt_satant_dict(mission, satant_file, umd_repo =umd_repo)
       satants = update_rt_filelist(mission, satant_file, umd_repo = umd_repo)
       print(satants)
    dry_run = args.dry_run
    rt_leo_preproc(config, satants=satants)
    rt_run_bern(config, satants=satants, dry_run = dry_run)

def main_para():
    args = parse_args()
   
    # Call the function assigned to the parser

    myfunc = args.func 
    args = vars(args) 
    print(args)

    inputs = [[args, 5, 1], [args, 5, 2], [args, 2, 1]]

    with Pool(processes=len(inputs)) as pool:
        pool.starmap(myfunc, inputs)

if __name__ == "__main__":
    main()

