#!/usr/bin/env python3
"""Run one manifest task selected by Slurm array index."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from dr_sfc.stage1_bc import run_stage1_bc_task
from dr_sfc.stage1_climatology import build_climatology
from dr_sfc.stage2_operational import run_stage2_operational_lead


def read_task(path: Path, array_index: int) -> dict[str, object]:
    if array_index < 1:
        raise ValueError("Task index is 1-based; Slurm array IDs should start at 1")
    with path.open("r", encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            if i == array_index:
                return json.loads(line)
    raise IndexError(f"No task {array_index} in {path}")


def run_task(task: dict[str, object]) -> None:
    stage = task.get("stage")
    if stage == "stage1_climatology":
        build_climatology(Path(str(task["input_dir"])), Path(str(task["output_dir"])))
        return
    if stage == "stage1_bc":
        output_dir = run_stage1_bc_task(task)
        print(f"Completed Stage 1 BC task: {output_dir}")
        return
    if stage == "stage2_operational_lead":
        outputs = run_stage2_operational_lead(task)
        print("Completed Stage 2 operational lead task:")
        for output in outputs:
            print(output)
        return
    raise ValueError(f"Unsupported task stage: {stage}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--task-index", type=int, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    task = read_task(args.manifest, args.task_index)
    run_task(task)


if __name__ == "__main__":
    main()