# Process a dataset, given
# - project_id = "none"
# - obs_type = "*"
from csst_dag.dag import CsstDAG
import argparse

# python -m csst_dag.trigger.csst-msc-l1-mbi -h
# python -m csst_dag.trigger.csst-msc-l1-mbi --dataset csst-msc-c9-25sqdeg-v3 --batch-id
parser = argparse.ArgumentParser(
    description="Scheduler for MSC L1 pipeline.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)

parser.add_argument("--dataset", type=str, help="Dataset name")
# parser.add_argument("--instrument", type=str, help="Instrument name", default="MSC")
parser.add_argument("--project-id", type=str, help="Project ID", default="None")
parser.add_argument("--obs-type", type=str, help="Observation type", default="WIDE")
parser.add_argument("--batch-id", type=str, help="Batch ID", default="default")
args = parser.parse_args()
print(args)

DAG_LOOP_MAP = {
    "WIDE": ["csst-msc-l1-mbi", "csst-msc-l1-sls"],
    "DEEP": ["csst-msc-l1-mbi", "csst-msc-l1-sls"],
    "BIAS": ["csst-msc-l1-qc0"],
    "DARK": ["csst-msc-l1-qc0"],
    "FLAT": ["csst-msc-l1-qc0"],
}

for obs_type, dag_ids in DAG_LOOP_MAP.items():
    print(f"* Processing {obs_type}")
    for dag_id in dag_ids:
        print(f"   - Scheduling `{dag_id}` -> ", end="")
        dag = CsstDAG.get_dag(dag_id=dag_id)
        msgs = dag.schedule(
            dataset=args.dataset,
            obs_type=args.obs_type,
            project_id=args.project_id,
            batch_id=args.batch_id,
            initial_prc_status=-2,
            final_prc_status=-2,
            demo=True,
        )
        print(f"{len(msgs)} tasks.")